346 lines
13 KiB
Python
346 lines
13 KiB
Python
import math
|
|
|
|
from pydantic import BaseModel, Field
|
|
from pydantic_ai import FunctionToolset, ToolReturn
|
|
|
|
from haiku.rag.agents.chat.context import (
|
|
cache_question_embedding,
|
|
get_cached_embedding,
|
|
trigger_background_summarization,
|
|
)
|
|
from haiku.rag.agents.chat.state import SessionContext
|
|
from haiku.rag.agents.research.dependencies import ResearchContext
|
|
from haiku.rag.agents.research.graph import build_research_graph
|
|
from haiku.rag.agents.research.models import Citation, SearchAnswer
|
|
from haiku.rag.agents.research.state import ResearchDeps, ResearchState
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.config.models import AppConfig
|
|
from haiku.rag.embeddings import get_embedder
|
|
from haiku.rag.tools.context import ToolContext
|
|
from haiku.rag.tools.filters import build_document_filter, combine_filters
|
|
from haiku.rag.tools.models import QAResult
|
|
from haiku.rag.tools.session import (
|
|
SESSION_NAMESPACE,
|
|
SessionState,
|
|
compute_combined_state_delta,
|
|
)
|
|
|
|
QA_NAMESPACE = "haiku.rag.qa"
|
|
|
|
PRIOR_ANSWER_RELEVANCE_THRESHOLD = 0.7
|
|
|
|
|
|
def _cosine_similarity(vec1: list[float], vec2: list[float]) -> float:
|
|
"""Compute cosine similarity between two vectors."""
|
|
dot_product = sum(a * b for a, b in zip(vec1, vec2))
|
|
norm1 = math.sqrt(sum(a * a for a in vec1))
|
|
norm2 = math.sqrt(sum(b * b for b in vec2))
|
|
if norm1 == 0 or norm2 == 0:
|
|
return 0.0
|
|
return dot_product / (norm1 * norm2)
|
|
|
|
|
|
class QAHistoryEntry(BaseModel):
|
|
"""A Q&A pair with optional cached embedding for similarity matching."""
|
|
|
|
question: str
|
|
answer: str
|
|
confidence: float = 0.9
|
|
citations: list[Citation] = []
|
|
question_embedding: list[float] | None = Field(default=None, exclude=True)
|
|
|
|
def to_search_answer(self) -> SearchAnswer:
|
|
"""Convert to SearchAnswer for research graph context."""
|
|
return SearchAnswer(
|
|
query=self.question,
|
|
answer=self.answer,
|
|
confidence=self.confidence,
|
|
cited_chunks=[c.chunk_id for c in self.citations],
|
|
citations=self.citations,
|
|
)
|
|
|
|
|
|
class QAState(BaseModel):
|
|
"""State for QA toolset.
|
|
|
|
Tracks Q&A history across tool invocations.
|
|
"""
|
|
|
|
history: list[QAResult] = []
|
|
|
|
|
|
class QASessionState(BaseModel):
|
|
"""Extended session state for QA with embedding cache."""
|
|
|
|
qa_history: list[QAHistoryEntry] = []
|
|
session_context: str | None = None
|
|
incoming_session_context: SessionContext | None = Field(
|
|
default=None, exclude=True
|
|
) # Track what client sent
|
|
|
|
|
|
QA_SESSION_NAMESPACE = "haiku.rag.qa_session"
|
|
MAX_QA_HISTORY = 50
|
|
|
|
|
|
def create_qa_toolset(
|
|
client: HaikuRAG,
|
|
config: AppConfig,
|
|
context: ToolContext | None = None,
|
|
base_filter: str | None = None,
|
|
tool_name: str = "ask",
|
|
session_context: str | None = None,
|
|
prior_answers: list[SearchAnswer] | None = None,
|
|
) -> FunctionToolset:
|
|
"""Create a toolset with Q&A capabilities using research graph.
|
|
|
|
Args:
|
|
client: HaikuRAG client for search operations.
|
|
config: Application configuration.
|
|
context: Optional ToolContext for state accumulation.
|
|
If provided, Q&A results are accumulated in QAState.
|
|
If SessionState is registered, it will be used for dynamic
|
|
document filtering and citation indexing.
|
|
base_filter: Optional base SQL WHERE clause applied to searches.
|
|
tool_name: Name for the ask tool. Defaults to "ask".
|
|
session_context: Optional session context for the research graph.
|
|
Overridden by QASessionState.session_context if available.
|
|
prior_answers: Optional list of prior answers for context.
|
|
Overridden by similarity-matched answers from QASessionState if available.
|
|
|
|
Returns:
|
|
FunctionToolset with an ask tool.
|
|
"""
|
|
# Get or create QA state if context provided
|
|
state: QAState | None = None
|
|
if context is not None:
|
|
state = context.get_or_create(QA_NAMESPACE, QAState)
|
|
|
|
async def ask(
|
|
question: str,
|
|
document_name: str | None = None,
|
|
) -> ToolReturn | QAResult:
|
|
"""Answer a question using the knowledge base.
|
|
|
|
Uses a research graph for searching and synthesizing answers.
|
|
|
|
Args:
|
|
question: The question to answer.
|
|
document_name: Optional document name/title to search within.
|
|
|
|
Returns:
|
|
QAResult with answer, confidence, and citations.
|
|
"""
|
|
# Get session states
|
|
session_state: SessionState | None = None
|
|
qa_session_state: QASessionState | None = None
|
|
old_state_snapshot: dict | None = None
|
|
|
|
if context is not None:
|
|
session_state = context.get_typed(SESSION_NAMESPACE, SessionState)
|
|
qa_session_state = context.get_typed(QA_SESSION_NAMESPACE, QASessionState)
|
|
|
|
# Capture combined state snapshot before changes
|
|
# Use incoming values (what client sent) so delta shows server-side updates
|
|
if session_state is not None:
|
|
old_state_snapshot = {
|
|
"session_id": session_state.incoming_session_id,
|
|
"document_filter": session_state.document_filter.copy(),
|
|
"citation_registry": session_state.citation_registry.copy(),
|
|
"citations": [c.model_dump() for c in session_state.citations],
|
|
}
|
|
if qa_session_state is not None:
|
|
old_state_snapshot["qa_history"] = [
|
|
qa.model_dump() for qa in qa_session_state.qa_history
|
|
]
|
|
# Use incoming_session_context so delta shows what client sent
|
|
if qa_session_state.incoming_session_context is not None:
|
|
old_state_snapshot["session_context"] = (
|
|
qa_session_state.incoming_session_context.model_dump(
|
|
mode="json"
|
|
)
|
|
)
|
|
else:
|
|
old_state_snapshot["session_context"] = None
|
|
|
|
# Build filter from session state, base_filter, and document_name
|
|
session_filter = None
|
|
if session_state is not None and session_state.document_filter:
|
|
from haiku.rag.tools.filters import build_multi_document_filter
|
|
|
|
session_filter = build_multi_document_filter(session_state.document_filter)
|
|
|
|
doc_filter = build_document_filter(document_name) if document_name else None
|
|
effective_filter = combine_filters(
|
|
combine_filters(base_filter, session_filter), doc_filter
|
|
)
|
|
|
|
# Determine session context
|
|
effective_session_context = session_context
|
|
if qa_session_state is not None and qa_session_state.session_context:
|
|
effective_session_context = qa_session_state.session_context
|
|
|
|
# Find relevant prior answers via similarity matching
|
|
effective_prior_answers = prior_answers or []
|
|
session_id = session_state.session_id if session_state is not None else ""
|
|
if qa_session_state is not None and qa_session_state.qa_history:
|
|
embedder = get_embedder(config)
|
|
question_embedding = await embedder.embed_query(question)
|
|
|
|
# Collect questions that need embedding
|
|
to_embed = []
|
|
to_embed_indices = []
|
|
for i, qa in enumerate(qa_session_state.qa_history):
|
|
if qa.question_embedding is None:
|
|
# Check per-session cache first
|
|
if session_id:
|
|
cached = get_cached_embedding(session_id, qa.question)
|
|
if cached:
|
|
qa.question_embedding = cached
|
|
continue
|
|
to_embed.append(qa.question)
|
|
to_embed_indices.append(i)
|
|
|
|
# Batch embed uncached questions
|
|
if to_embed:
|
|
new_embeddings = await embedder.embed_documents(to_embed)
|
|
for i, idx in enumerate(to_embed_indices):
|
|
embedding = new_embeddings[i]
|
|
qa_session_state.qa_history[idx].question_embedding = embedding
|
|
# Cache per-session for next request
|
|
if session_id:
|
|
cache_question_embedding(
|
|
session_id,
|
|
qa_session_state.qa_history[idx].question,
|
|
embedding,
|
|
)
|
|
|
|
# Find similar prior answers
|
|
matched_answers = []
|
|
for qa in qa_session_state.qa_history:
|
|
if qa.question_embedding is not None:
|
|
similarity = _cosine_similarity(
|
|
question_embedding, qa.question_embedding
|
|
)
|
|
if similarity >= PRIOR_ANSWER_RELEVANCE_THRESHOLD:
|
|
matched_answers.append(qa.to_search_answer())
|
|
|
|
if matched_answers:
|
|
effective_prior_answers = matched_answers
|
|
|
|
# Build and run the research graph
|
|
graph = build_research_graph(config=config, output_mode="conversational")
|
|
|
|
research_context = ResearchContext(
|
|
original_question=question,
|
|
session_context=effective_session_context,
|
|
qa_responses=effective_prior_answers,
|
|
)
|
|
research_state = ResearchState(
|
|
context=research_context,
|
|
max_iterations=1,
|
|
search_filter=effective_filter,
|
|
max_concurrency=config.research.max_concurrency,
|
|
)
|
|
deps = ResearchDeps(client=client)
|
|
|
|
result = await graph.run(state=research_state, deps=deps)
|
|
|
|
# Build citations with stable indices from session state
|
|
citations = []
|
|
for i, c in enumerate(result.citations):
|
|
if session_state is not None:
|
|
index = session_state.get_or_assign_index(c.chunk_id)
|
|
else:
|
|
index = i + 1
|
|
|
|
citations.append(
|
|
Citation(
|
|
index=index,
|
|
document_id=c.document_id,
|
|
chunk_id=c.chunk_id,
|
|
document_uri=c.document_uri,
|
|
document_title=c.document_title,
|
|
page_numbers=c.page_numbers,
|
|
headings=c.headings,
|
|
content=c.content,
|
|
)
|
|
)
|
|
|
|
qa_result = QAResult(
|
|
question=question,
|
|
answer=result.answer,
|
|
confidence=result.confidence,
|
|
citations=citations,
|
|
)
|
|
|
|
# Accumulate in QA state if context provided
|
|
if state is not None:
|
|
state.history.append(qa_result)
|
|
|
|
# Update session state with citations
|
|
if session_state is not None:
|
|
session_state.citations = citations
|
|
|
|
# Update QA session state with history entry
|
|
if qa_session_state is not None:
|
|
qa_session_state.qa_history.append(
|
|
QAHistoryEntry(
|
|
question=question,
|
|
answer=result.answer,
|
|
confidence=result.confidence,
|
|
citations=citations,
|
|
)
|
|
)
|
|
# Enforce FIFO limit
|
|
if len(qa_session_state.qa_history) > MAX_QA_HISTORY:
|
|
qa_session_state.qa_history = qa_session_state.qa_history[
|
|
-MAX_QA_HISTORY:
|
|
]
|
|
# Trigger background summarization
|
|
trigger_background_summarization(
|
|
qa_session_state=qa_session_state,
|
|
config=config,
|
|
session_id=session_id,
|
|
)
|
|
|
|
# Compute and return state delta if session state changed
|
|
if session_state is not None and old_state_snapshot is not None:
|
|
# Build new combined state snapshot
|
|
new_state_snapshot = {
|
|
"session_id": session_state.session_id,
|
|
"document_filter": session_state.document_filter,
|
|
"citation_registry": session_state.citation_registry,
|
|
"citations": [c.model_dump() for c in session_state.citations],
|
|
}
|
|
if qa_session_state is not None:
|
|
new_state_snapshot["qa_history"] = [
|
|
qa.model_dump() for qa in qa_session_state.qa_history
|
|
]
|
|
if qa_session_state.session_context:
|
|
new_state_snapshot["session_context"] = SessionContext(
|
|
summary=qa_session_state.session_context
|
|
).model_dump(mode="json")
|
|
else:
|
|
new_state_snapshot["session_context"] = None
|
|
|
|
state_event = compute_combined_state_delta(
|
|
old_state_snapshot,
|
|
new_state_snapshot,
|
|
state_key=session_state.state_key,
|
|
)
|
|
|
|
# Format answer with citation references
|
|
answer_text = result.answer
|
|
if citations:
|
|
citation_refs = " ".join(f"[{c.index}]" for c in citations)
|
|
answer_text = f"{answer_text}\n\nSources: {citation_refs}"
|
|
|
|
metadata = [state_event] if state_event is not None else None
|
|
return ToolReturn(return_value=answer_text, metadata=metadata)
|
|
|
|
return qa_result
|
|
|
|
toolset = FunctionToolset()
|
|
toolset.add_function(ask, name=tool_name)
|
|
return toolset
|