studyhub/core/answer.py

from __future__ import annotations

from dataclasses import dataclass, field

from ..ingest.embedder import Embedder
from .citations import Citation, extract_citations, render_sources
from .llm import LLMClient
from .loader import RetrievalIndex
from .prompt import build_context, detect_language, system_prompt
from .retrieve import retrieve


@dataclass
class AnswerResult:
    answer: str
    sources: list[Citation]
    sources_md: str
    used_ids: list[str] = field(default_factory=list)


THREAD_RULE_EN = ("This is a continuing conversation. Cite ONLY [S#] ids that appear in the "
                  "CURRENT sources below; ids mentioned in earlier answers are not valid.")
THREAD_RULE_IT = ("Questa è una conversazione in corso. Cita SOLO gli id [S#] presenti nelle "
                  "fonti ATTUALI qui sotto; gli id citati nelle risposte precedenti non sono validi.")

_MAX_TURNS = 3        # turns of context shown to the model
_CLIP_CHARS = 600     # per-answer clip so old answers can't crowd out sources


def _conversation_block(prior: list[dict]) -> str:
    lines = []
    for t in prior[-_MAX_TURNS:]:
        a = str(t.get("a", ""))
        if len(a) > _CLIP_CHARS:
            a = a[:_CLIP_CHARS] + " …"
        lines.append(f"USER: {t.get('q', '')}\nASSISTANT: {a}")
    return "CONVERSATION SO FAR:\n" + "\n\n".join(lines)


def generate_from(query: str, scored, llm: LLMClient, token_budget: int = 4000) -> AnswerResult:
    """Build the grounded prompt from already-retrieved chunks, generate, attach citations."""
    context, id_map = build_context(scored, token_budget)
    if not id_map:
        return AnswerResult("I couldn't find anything about that in the materials.", [], "")
    text = llm.complete(system_prompt(detect_language(query)), f"SOURCES:\n{context}\n\nQUESTION: {query}")
    cites = extract_citations(text, id_map)
    return AnswerResult(answer=text, sources=cites, sources_md=render_sources(cites),
                        used_ids=[c.sid for c in cites])


def answer(query: str, index: RetrievalIndex, embedder: Embedder, llm: LLMClient,
           k: int = 6, token_budget: int = 4000, reranker=None) -> AnswerResult:
    return generate_from(query, retrieve(query, index, embedder, k=k, reranker=reranker), llm, token_budget)


def generate_in_thread(query: str, prior: list[dict], scored, llm: LLMClient,
                       token_budget: int = 4000) -> AnswerResult:
    """generate_from for an ongoing thread: same grounding, plus a clipped transcript of the
    last turns so follow-ups can resolve references; citations come only from current sources."""
    context, id_map = build_context(scored, token_budget)
    if not id_map:
        return AnswerResult("I couldn't find anything about that in the materials.", [], "")
    lang = detect_language(query)
    system = system_prompt(lang) + "\n" + (THREAD_RULE_IT if lang == "it" else THREAD_RULE_EN)
    convo = _conversation_block(prior) + "\n\n" if prior else ""
    text = llm.complete(system, f"{convo}SOURCES:\n{context}\n\nQUESTION: {query}")
    cites = extract_citations(text, id_map)
    return AnswerResult(answer=text, sources=cites, sources_md=render_sources(cites),
                        used_ids=[c.sid for c in cites])


def answer_in_thread(query: str, prior: list[dict], index: RetrievalIndex, embedder: Embedder,
                     llm: LLMClient, k: int = 6, token_budget: int = 4000, reranker=None) -> AnswerResult:
    """answer() with conversation context: retrieval still runs on the current query alone."""
    return generate_in_thread(query, prior, retrieve(query, index, embedder, k=k, reranker=reranker),
                              llm, token_budget)