chat.py

from __future__ import annotations

import argparse

from rag_assistant.assistant import RagAssistant
from rag_assistant.errors import RagAssistantError


def main() -> None:
    parser = argparse.ArgumentParser(description="Ask questions about the indexed codebase.")
    parser.add_argument("--top-k", type=int, default=8, help="Number of chunks to retrieve per question.")
    args = parser.parse_args()

    assistant = RagAssistant()

    print("RAG Codebase Assistant")
    print("Ask a question, or type 'exit' to quit.")

    while True:
        question = input("\n> ").strip()
        if question.lower() in {"exit", "quit", "q"}:
            break
        if not question:
            continue

        try:
            print()
            for token in assistant.stream(question, top_k=args.top_k):
                print(token, end="", flush=True)
            print()
        except RagAssistantError as exc:
            print(f"\nError: {exc}")


if __name__ == "__main__":
    main()