rag.py

"""
RAG Pipeline complet
Combină Retriever + Generator într-un singur query flow.
"""

import time

from retriever import Retriever
from generator import generate_answer, generate_answer_stream
from config import TOP_K


class RAGPipeline:
    """Pipeline complet: întrebare → retrieval → generation → răspuns."""

    def __init__(self):
        print("Inițializăm RAG Pipeline ...")
        self.retriever = Retriever()
        print("RAG Pipeline gata!\n")

    def query(self, question: str, top_k: int = TOP_K) -> dict:
        """
        Procesează o întrebare end-to-end.

        Returns:
            dict cu: answer, sources, timing
        """
        start = time.time()

        # Pas 1: Retrieval
        context, sources = self.retriever.build_context(question, top_k)
        retrieval_time = time.time() - start

        # Pas 2: Generation
        gen_start = time.time()
        answer = generate_answer(question, context)
        generation_time = time.time() - gen_start

        return {
            "question": question,
            "answer": answer,
            "sources": [
                {
                    "text": s["text"][:200],
                    "source": s["metadata"].get("source", "?"),
                    "score": round(s["score"], 3),
                }
                for s in sources
            ],
            "timing": {
                "retrieval_ms": round(retrieval_time * 1000),
                "generation_ms": round(generation_time * 1000),
                "total_ms": round((time.time() - start) * 1000),
            },
        }

    def query_stream(self, question: str, top_k: int = TOP_K):
        """
        Procesează o întrebare cu streaming pentru răspuns.

        Yields:
            dict-uri cu type="source" sau type="token"
        """
        context, sources = self.retriever.build_context(question, top_k)

        # Trimite sursele mai întâi
        yield {
            "type": "sources",
            "sources": [
                {
                    "text": s["text"][:200],
                    "source": s["metadata"].get("source", "?"),
                    "score": round(s["score"], 3),
                }
                for s in sources
            ],
        }

        # Apoi stream-ul de tokens
        for token in generate_answer_stream(question, context):
            yield {"type": "token", "content": token}


if __name__ == "__main__":
    rag = RAGPipeline()

    questions = [
        "How do I create a basic FastAPI application?",
        "What is dependency injection in FastAPI?",
        "How do I handle errors?",
    ]

    for q in questions:
        print(f"\n{'='*60}")
        print(f"Q: {q}")
        print("-" * 60)

        result = rag.query(q)
        print(f"A: {result['answer']}")
        print(f"\nSurse:")
        for s in result["sources"]:
            print(f"  - {s['source']} (scor: {s['score']})")
        print(f"\nTimp: {result['timing']['total_ms']}ms "
              f"(retrieval: {result['timing']['retrieval_ms']}ms, "
              f"generation: {result['timing']['generation_ms']}ms)")