studyhub/core/topic.py

from __future__ import annotations

import numpy as np

from .answer import AnswerResult, generate_from, generate_in_thread
from .prompt import detect_language
from .retrieve import ScoredChunk
from .study import summarize_chunks


def retrieve_in_topic(query, index, chunk_ids, embedder, k: int = 6, reranker=None):
    """Like retrieve(), but only over a topic's own chunks (across all its sources)."""
    qv = embedder.encode([query]).astype("float32")[0]
    scored: list[ScoredChunk] = []
    for cid in chunk_ids:
        row = index.row_by_id.get(cid)
        if row is None:
            continue
        vec = np.asarray(index.faiss_index.reconstruct(int(row)), dtype="float32")
        scored.append(ScoredChunk(chunk=index.by_id[cid], score=float(np.dot(qv, vec))))
    scored.sort(key=lambda s: s.score, reverse=True)
    if reranker is not None and scored:
        pool = scored[: max(k, 16)]
        for sc, rs in zip(pool, reranker.scores(query, [s.chunk.text for s in pool])):
            sc.score = float(rs)
        pool.sort(key=lambda s: s.score, reverse=True)
        return pool[:k]
    return scored[:k]


def answer_topic(query, index, chunk_ids, embedder, llm, k: int = 6, token_budget: int = 4000, reranker=None):
    scored = retrieve_in_topic(query, index, chunk_ids, embedder, k=k, reranker=reranker)
    return generate_from(query, scored, llm, token_budget)


def answer_topic_in_thread(query, prior, index, chunk_ids, embedder, llm,
                           k: int = 6, token_budget: int = 4000, reranker=None):
    """answer_topic() with conversation context (same treatment as answer_in_thread)."""
    scored = retrieve_in_topic(query, index, chunk_ids, embedder, k=k, reranker=reranker)
    return generate_in_thread(query, prior, scored, llm, token_budget)


def topic_summary(index, chunk_ids, llm) -> AnswerResult:
    chunks = [index.by_id[c] for c in chunk_ids if c in index.by_id]
    if not chunks:
        return AnswerResult("No material linked to this topic yet.", [], "")
    lang = detect_language(" ".join(c.text for c in chunks[:3]))
    return summarize_chunks(chunks, llm, lang=lang)