322 lines
11 KiB
Python
322 lines
11 KiB
Python
"""
|
||
Agent-like chat orchestration for grounded QA over documents using LlamaIndex.
|
||
|
||
This module separates:
|
||
1) retrieval of source nodes
|
||
2) answer synthesis with an explicit grounded prompt
|
||
3) response formatting with sources/metadata
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import re
|
||
from dataclasses import dataclass
|
||
from typing import Any, Iterable, List
|
||
|
||
from llama_index.core import PromptTemplate
|
||
from llama_index.core.agent import AgentWorkflow
|
||
from llama_index.core.retrievers import VectorIndexRetriever
|
||
from llama_index.core.schema import NodeWithScore
|
||
from llama_index.core.tools import FunctionTool
|
||
from loguru import logger
|
||
|
||
from config import get_llm_model, setup_global_models
|
||
from vector_storage import get_vector_store_and_index
|
||
|
||
|
||
GROUNDED_SYNTHESIS_PROMPT = PromptTemplate(
|
||
"""You are a grounded QA assistant for a document knowledge base.
|
||
|
||
Rules:
|
||
- Answer ONLY from the provided context snippets.
|
||
- If the context is insufficient, say directly that the information is not available in the retrieved sources.
|
||
- Prefer exact dates/years and cite filenames/pages when possible.
|
||
- Avoid generic claims that are not supported by the snippets.
|
||
- If multiple sources disagree, mention the conflict briefly.
|
||
|
||
User question:
|
||
{query}
|
||
|
||
Optional draft from tool-using agent (may be incomplete):
|
||
{agent_draft}
|
||
|
||
Context snippets (JSON):
|
||
{context_json}
|
||
|
||
Return a concise answer with source mentions in plain text.
|
||
"""
|
||
)
|
||
|
||
|
||
@dataclass
|
||
class RetrievalSnippet:
|
||
content: str
|
||
score: float | None
|
||
metadata: dict[str, Any]
|
||
|
||
def to_api_dict(self) -> dict[str, Any]:
|
||
metadata = self.metadata or {}
|
||
content_preview = self.content.strip().replace("\n", " ")
|
||
if len(content_preview) > 400:
|
||
content_preview = content_preview[:400] + "..."
|
||
return {
|
||
"content_snippet": content_preview,
|
||
"score": self.score,
|
||
"metadata": {
|
||
"filename": metadata.get("filename", "unknown"),
|
||
"file_path": metadata.get("file_path", "unknown"),
|
||
"page_label": metadata.get("page_label", metadata.get("page", "unknown")),
|
||
"chunk_number": metadata.get("chunk_number", "unknown"),
|
||
"total_chunks": metadata.get("total_chunks", "unknown"),
|
||
"file_type": metadata.get("file_type", "unknown"),
|
||
"processed_at": metadata.get("processed_at", "unknown"),
|
||
},
|
||
}
|
||
|
||
|
||
def _normalize_text(value: Any) -> str:
|
||
if value is None:
|
||
return ""
|
||
if isinstance(value, bytes):
|
||
return value.decode("utf-8", errors="replace")
|
||
return str(value)
|
||
|
||
|
||
def _extract_years(query: str) -> list[int]:
|
||
years = []
|
||
for match in re.findall(r"\b(19\d{2}|20\d{2}|21\d{2})\b", query):
|
||
try:
|
||
years.append(int(match))
|
||
except ValueError:
|
||
continue
|
||
return sorted(set(years))
|
||
|
||
|
||
def _extract_keywords(query: str) -> list[str]:
|
||
words = re.findall(r"[A-Za-zА-Яа-я0-9_-]{4,}", query.lower())
|
||
stop = {"what", "when", "where", "which", "that", "this", "with", "from", "about"}
|
||
keywords = [w for w in words if w not in stop]
|
||
return list(dict.fromkeys(keywords))[:6]
|
||
|
||
|
||
def _node_text(node_with_score: NodeWithScore) -> str:
|
||
node = getattr(node_with_score, "node", None)
|
||
if node is None:
|
||
return _normalize_text(getattr(node_with_score, "text", ""))
|
||
return _normalize_text(getattr(node, "text", getattr(node_with_score, "text", "")))
|
||
|
||
|
||
def _node_metadata(node_with_score: NodeWithScore) -> dict[str, Any]:
|
||
node = getattr(node_with_score, "node", None)
|
||
if node is None:
|
||
return dict(getattr(node_with_score, "metadata", {}) or {})
|
||
return dict(getattr(node, "metadata", {}) or {})
|
||
|
||
|
||
def _similarity_key(text: str) -> str:
|
||
text = re.sub(r"\s+", " ", text.strip().lower())
|
||
return text[:250]
|
||
|
||
|
||
def _is_low_information_chunk(text: str) -> bool:
|
||
compact = " ".join(text.split())
|
||
if len(compact) < 20:
|
||
return True
|
||
alpha_chars = sum(ch.isalpha() for ch in compact)
|
||
if alpha_chars < 8:
|
||
return True
|
||
# Repetitive headers/footers often contain too few unique tokens.
|
||
tokens = [t for t in re.split(r"\W+", compact.lower()) if t]
|
||
return len(tokens) >= 3 and len(set(tokens)) <= 2
|
||
|
||
|
||
def post_process_nodes(nodes: list[NodeWithScore]) -> list[NodeWithScore]:
|
||
"""
|
||
Post-process retrieved nodes:
|
||
- drop empty / near-empty chunks
|
||
- optionally drop low-information chunks
|
||
- deduplicate near-identical chunks
|
||
"""
|
||
filtered: list[NodeWithScore] = []
|
||
seen = set()
|
||
|
||
for nws in nodes:
|
||
text = _node_text(nws)
|
||
if not text or not text.strip():
|
||
continue
|
||
if _is_low_information_chunk(text):
|
||
continue
|
||
|
||
meta = _node_metadata(nws)
|
||
dedup_key = (
|
||
meta.get("file_path", ""),
|
||
meta.get("page_label", meta.get("page", "")),
|
||
meta.get("chunk_number", ""),
|
||
_similarity_key(text),
|
||
)
|
||
if dedup_key in seen:
|
||
continue
|
||
seen.add(dedup_key)
|
||
filtered.append(nws)
|
||
|
||
return filtered
|
||
|
||
|
||
def retrieve_source_nodes(query: str, top_k: int = 5, search_multiplier: int = 3) -> list[NodeWithScore]:
|
||
"""
|
||
Retrieve source nodes with light metadata-aware query expansion (years/keywords).
|
||
"""
|
||
setup_global_models()
|
||
_, index = get_vector_store_and_index()
|
||
retriever = VectorIndexRetriever(index=index, similarity_top_k=max(top_k * search_multiplier, top_k))
|
||
|
||
queries = [query]
|
||
years = _extract_years(query)
|
||
keywords = _extract_keywords(query)
|
||
queries.extend(str(year) for year in years)
|
||
queries.extend(keywords[:3])
|
||
|
||
collected: list[NodeWithScore] = []
|
||
for q in list(dict.fromkeys([q for q in queries if q and q.strip()])):
|
||
try:
|
||
logger.info(f"Retrieving nodes for query variant: {q}")
|
||
collected.extend(retriever.retrieve(q))
|
||
except Exception as e:
|
||
logger.warning(f"Retrieval variant failed for '{q}': {e}")
|
||
|
||
processed = post_process_nodes(collected)
|
||
processed.sort(key=lambda n: (getattr(n, "score", None) is None, -(getattr(n, "score", 0.0) or 0.0)))
|
||
return processed[:top_k]
|
||
|
||
|
||
def build_structured_snippets(nodes: list[NodeWithScore]) -> list[dict[str, Any]]:
|
||
"""Return structured snippets for tools/API responses."""
|
||
snippets: list[dict[str, Any]] = []
|
||
for nws in nodes:
|
||
snippet = RetrievalSnippet(
|
||
content=_node_text(nws),
|
||
score=getattr(nws, "score", None),
|
||
metadata=_node_metadata(nws),
|
||
)
|
||
snippets.append(snippet.to_api_dict())
|
||
return snippets
|
||
|
||
|
||
def retrieval_tool_search(query: str, top_k: int = 5) -> str:
|
||
"""
|
||
Tool wrapper for document retrieval returning structured JSON snippets.
|
||
"""
|
||
nodes = retrieve_source_nodes(query=query, top_k=top_k)
|
||
snippets = build_structured_snippets(nodes)
|
||
payload = {
|
||
"query": query,
|
||
"count": len(snippets),
|
||
"snippets": snippets,
|
||
}
|
||
return json.dumps(payload, ensure_ascii=False)
|
||
|
||
|
||
def synthesize_answer(query: str, sources: list[dict[str, Any]], agent_draft: str = "") -> str:
|
||
"""
|
||
Answer synthesis from retrieved sources using an explicit grounded prompt.
|
||
"""
|
||
llm = get_llm_model()
|
||
context_json = json.dumps(sources, ensure_ascii=False, indent=2)
|
||
prompt = GROUNDED_SYNTHESIS_PROMPT.format(
|
||
query=query,
|
||
agent_draft=agent_draft or "(none)",
|
||
context_json=context_json,
|
||
)
|
||
logger.info("Synthesizing grounded answer from retrieved sources")
|
||
response = llm.complete(prompt)
|
||
return _normalize_text(getattr(response, "text", response))
|
||
|
||
|
||
def format_chat_response(query: str, final_answer: str, sources: list[dict[str, Any]], mode: str) -> dict[str, Any]:
|
||
"""
|
||
Response formatting with answer + structured sources.
|
||
"""
|
||
return {
|
||
"query": query,
|
||
"answer": final_answer,
|
||
"sources": sources,
|
||
"mode": mode,
|
||
}
|
||
|
||
|
||
def _extract_agent_result_text(result: Any) -> str:
|
||
if result is None:
|
||
return ""
|
||
if hasattr(result, "response"):
|
||
return _normalize_text(getattr(result, "response"))
|
||
return _normalize_text(result)
|
||
|
||
|
||
async def _run_agent_workflow_async(query: str, top_k: int) -> str:
|
||
"""
|
||
Run LlamaIndex AgentWorkflow with a retrieval tool. Returns agent draft answer text.
|
||
"""
|
||
setup_global_models()
|
||
llm = get_llm_model()
|
||
tool = FunctionTool.from_defaults(
|
||
fn=retrieval_tool_search,
|
||
name="document_search",
|
||
description=(
|
||
"Search documents and return structured snippets as JSON with fields: "
|
||
"filename, file_path, page_label/page, chunk_number, content_snippet, score. "
|
||
"Use this before answering factual questions about documents."
|
||
),
|
||
)
|
||
|
||
system_prompt = (
|
||
"You are a QA agent over a document store. Use the document_search tool when factual "
|
||
"information may come from documents. If tool output is insufficient, say so."
|
||
)
|
||
agent = AgentWorkflow.from_tools_or_functions(
|
||
[tool],
|
||
llm=llm,
|
||
system_prompt=system_prompt,
|
||
verbose=False,
|
||
)
|
||
handler = agent.run(user_msg=query, max_iterations=4)
|
||
result = await handler
|
||
return _extract_agent_result_text(result)
|
||
|
||
|
||
def run_agent_workflow(query: str, top_k: int = 5) -> str:
|
||
"""
|
||
Synchronous wrapper around the async LlamaIndex agent workflow.
|
||
"""
|
||
try:
|
||
return asyncio.run(_run_agent_workflow_async(query=query, top_k=top_k))
|
||
except RuntimeError:
|
||
# Fallback if already in an event loop; skip agent workflow in that case.
|
||
logger.warning("Async event loop already running; skipping agent workflow and using direct retrieval+synthesis")
|
||
return ""
|
||
except Exception as e:
|
||
logger.warning(f"Agent workflow failed, will fallback to direct retrieval+synthesis: {e}")
|
||
return ""
|
||
|
||
|
||
def chat_with_documents(query: str, top_k: int = 5, use_agent: bool = True) -> dict[str, Any]:
|
||
"""
|
||
Full chat orchestration entrypoint:
|
||
- optionally run agent workflow (tool-calling)
|
||
- retrieve + post-process sources
|
||
- synthesize grounded answer
|
||
- format response
|
||
"""
|
||
logger.info(f"Starting chat orchestration for query: {query[:80]}")
|
||
agent_draft = ""
|
||
mode = "retrieval+synthesis"
|
||
if use_agent:
|
||
agent_draft = run_agent_workflow(query=query, top_k=top_k)
|
||
mode = "agent+retrieval+synthesis" if agent_draft else "retrieval+synthesis"
|
||
|
||
nodes = retrieve_source_nodes(query=query, top_k=top_k)
|
||
sources = build_structured_snippets(nodes)
|
||
final_answer = synthesize_answer(query=query, sources=sources, agent_draft=agent_draft)
|
||
return format_chat_response(query=query, final_answer=final_answer, sources=sources, mode=mode)
|