from techdocs_graphrag.models import Answer, RetrievalMode
from techdocs_graphrag.ports import LanguageModel
from techdocs_graphrag.prompts import ANSWER_SYSTEM_PROMPT, render_answer_prompt
from techdocs_graphrag.retrieval import HybridRetriever

NO_CONTEXT_ANSWER = (
    "Je ne trouve aucune information pertinente dans la base indexée. "
    "Indexe d'abord un corpus ou reformule la question."
)


class GraphRAGAssistant:
    def __init__(self, *, retriever: HybridRetriever, llm: LanguageModel) -> None:
        self._retriever = retriever
        self._llm = llm

    def ask(self, question: str, mode: RetrievalMode = RetrievalMode.GRAPH) -> Answer:
        context = self._retriever.retrieve(question, mode)
        text = (
            NO_CONTEXT_ANSWER
            if context.is_empty
            else self._llm.complete(ANSWER_SYSTEM_PROMPT, render_answer_prompt(question, context))
        )
        return Answer(question=question, text=text, mode=mode, context=context)