|
# retoor <retoor@molodetz.nl>
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import re
|
|
from dataclasses import dataclass, field
|
|
|
|
from .embeddings import embed_texts, local_embed
|
|
from .llm import complete_chat
|
|
from .store import Chunk, VectorStore
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
CITATION_MARKER = re.compile(r"\[(\d+)\]")
|
|
|
|
CHAT_TOP_K = 10
|
|
MAX_CONTEXT_CHARS = 16000
|
|
CHAT_MAX_TOKENS = 1400
|
|
|
|
SYSTEM_PROMPT = (
|
|
"You are the DeepSearch research assistant. Answer the user's question using ONLY "
|
|
"the numbered SOURCES below, which were gathered during a web research session. "
|
|
"Never use outside knowledge or guess. If the sources do not contain the answer, "
|
|
"say so plainly. Cite every claim inline with the bracket marker of the source it "
|
|
"comes from, like [1] or [2]. Keep the answer focused and well structured in "
|
|
"markdown."
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class ChatAnswer:
|
|
text: str
|
|
citations: list[dict] = field(default_factory=list)
|
|
|
|
|
|
def _build_context(chunks: list[Chunk]) -> tuple[str, list[dict]]:
|
|
blocks: list[str] = []
|
|
citations: list[dict] = []
|
|
used = 0
|
|
for index, chunk in enumerate(chunks, start=1):
|
|
snippet = chunk.text.strip()
|
|
if not snippet:
|
|
continue
|
|
header = f"[{index}] {chunk.title or chunk.url} ({chunk.url})"
|
|
block = f"{header}\n{snippet}"
|
|
if used + len(block) > MAX_CONTEXT_CHARS and blocks:
|
|
break
|
|
used += len(block)
|
|
blocks.append(block)
|
|
citations.append(
|
|
{
|
|
"index": index,
|
|
"url": chunk.url,
|
|
"title": chunk.title or chunk.url,
|
|
"score": round(chunk.score, 4),
|
|
}
|
|
)
|
|
return "\n\n".join(blocks), citations
|
|
|
|
|
|
class DeepsearchChat:
|
|
def __init__(self, collection_name: str, api_key: str) -> None:
|
|
self.store = VectorStore(collection_name)
|
|
self.api_key = api_key
|
|
|
|
async def _embed_query(self, question: str) -> list[float]:
|
|
result = await embed_texts([question], self.api_key)
|
|
if not result.vectors or not result.vectors[0]:
|
|
result = local_embed([question])
|
|
return result.vectors[0]
|
|
|
|
async def retrieve(self, question: str) -> list[Chunk]:
|
|
query_vector = await self._embed_query(question)
|
|
stored_dim = self.store.dims
|
|
if stored_dim is not None and len(query_vector) != stored_dim:
|
|
query_vector = local_embed([question]).vectors[0]
|
|
if len(query_vector) != stored_dim:
|
|
logger.warning(
|
|
"deepsearch query embedding dim %d != stored %d",
|
|
len(query_vector),
|
|
stored_dim,
|
|
)
|
|
return []
|
|
return await self.store.hybrid_search(question, query_vector, top_k=CHAT_TOP_K)
|
|
|
|
async def answer(self, question: str, history: list[dict] | None = None) -> ChatAnswer:
|
|
chunks = await self.retrieve(question)
|
|
if not chunks:
|
|
return ChatAnswer(
|
|
text=(
|
|
"The research session did not capture anything relevant to that "
|
|
"question. Try rephrasing or running a deeper search."
|
|
),
|
|
citations=[],
|
|
)
|
|
context, citations = _build_context(chunks)
|
|
messages = [{"role": "system", "content": SYSTEM_PROMPT}]
|
|
for turn in (history or [])[-6:]:
|
|
role = turn.get("role")
|
|
content = turn.get("content")
|
|
if role in ("user", "assistant") and content:
|
|
messages.append({"role": role, "content": content})
|
|
messages.append(
|
|
{
|
|
"role": "user",
|
|
"content": f"SOURCES:\n{context}\n\nQUESTION: {question}",
|
|
}
|
|
)
|
|
try:
|
|
text = await complete_chat(
|
|
messages, self.api_key, max_tokens=CHAT_MAX_TOKENS
|
|
)
|
|
except Exception as exc:
|
|
logger.warning("deepsearch chat synthesis failed: %s", exc)
|
|
text = (
|
|
"I could not reach the language model to synthesise an answer, but the "
|
|
"most relevant sources are listed below."
|
|
)
|
|
valid = {citation["index"] for citation in citations}
|
|
text = self._strip_unmatched_markers(text, valid)
|
|
return ChatAnswer(text=text, citations=citations)
|
|
|
|
def _strip_unmatched_markers(self, text: str, valid: set[int]) -> str:
|
|
def replace(match: re.Match) -> str:
|
|
index = int(match.group(1))
|
|
return match.group(0) if index in valid else ""
|
|
|
|
return CITATION_MARKER.sub(replace, text)
|