generator.py

"""
RAG Generator
Primește contextul din retriever + întrebarea utilizatorului → generează răspuns cu Phi-3 prin Ollama.
"""

import ollama

from config import OLLAMA_MODEL

SYSTEM_PROMPT = """You are a helpful technical documentation assistant. 
Answer the user's question based ONLY on the provided context. 
If the context doesn't contain enough information to answer, say so honestly.
Always cite which source document(s) you used in your answer.
Keep answers concise and accurate."""


def generate_answer(query: str, context: str, model: str = OLLAMA_MODEL) -> str:
    """
    Generează un răspuns bazat pe contextul recuperat din documente.

    Args:
        query: Întrebarea utilizatorului
        context: Contextul formatat din retriever (chunk-uri relevante)
        model: Modelul Ollama de folosit

    Returns:
        Răspunsul generat de LLM
    """
    user_message = f"""Context from documentation:
{context}

---

Question: {query}

Answer based on the context above. Cite sources using [1], [2], etc."""

    response = ollama.chat(
        model=model,
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_message},
        ],
    )

    return response["message"]["content"]


def generate_answer_stream(query: str, context: str, model: str = OLLAMA_MODEL):
    """
    Generează un răspuns cu streaming (text apare progresiv).

    Yields:
        Bucăți de text pe măsură ce sunt generate
    """
    user_message = f"""Context from documentation:
{context}

---

Question: {query}

Answer based on the context above. Cite sources using [1], [2], etc."""

    stream = ollama.chat(
        model=model,
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_message},
        ],
        stream=True,
    )

    for chunk in stream:
        yield chunk["message"]["content"]


if __name__ == "__main__":
    test_context = "[1] (Sursa: fastapi.md)\nFastAPI is a modern, fast web framework for building APIs with Python 3.7+."
    test_query = "What is FastAPI?"

    print(f"Query: {test_query}\n")
    print("Răspuns: ", end="", flush=True)
    for token in generate_answer_stream(test_query, test_context):
        print(token, end="", flush=True)
    print()