"""LangChain tools for the ReAct agent."""
import logging
import re
from dataclasses import dataclass, field
from langchain_core.runnables import Runnable
from langchain_core.tools import tool
from src.agent.prompts import render_prompt
from src.models import QueryResult
from src.retrieval.hybrid import HybridRetriever
from src.retrieval.reranker import Reranker
from src.retrieval.vector_store import VectorStore
logger = logging.getLogger(__name__)
_THINK_CLOSED_RE = re.compile(r".*?\s*", re.DOTALL)
_THINK_UNCLOSED_RE = re.compile(r".*", re.DOTALL)
def _strip_think(text: str) -> str:
"""Remove ```` blocks — both closed and unclosed."""
text = _THINK_CLOSED_RE.sub("", text)
text = _THINK_UNCLOSED_RE.sub("", text)
return text.strip()
def _extract_content(result: object) -> str:
"""Extract plain text from an LLM invoke result.
Handles AIMessage (content: str or list), plain strings, etc.
"""
if hasattr(result, "content"):
content = result.content
else:
content = result
if isinstance(content, list):
parts: list[str] = []
for block in content:
if isinstance(block, str):
parts.append(block)
elif isinstance(block, dict) and "text" in block:
parts.append(block["text"])
text = "\n".join(parts)
else:
text = str(content)
return _strip_think(text)
@dataclass
class ToolResultStore:
"""Captures structured retrieval results produced during tool invocations.
Attributes:
retrieved: Accumulated QueryResult list across all tool calls,
merged by chunk_id and sorted by descending score.
tool_calls: Log of (tool_name, query_or_arg) tuples in invocation order.
dense_results: Accumulated dense retrieval results across hybrid_search calls.
sparse_results: Accumulated sparse (BM25) retrieval results across hybrid_search calls.
fused_results: Accumulated RRF-fused results across hybrid_search calls.
"""
retrieved: list[QueryResult] = field(default_factory=list)
tool_calls: list[tuple[str, str]] = field(default_factory=list)
dense_results: list[QueryResult] = field(default_factory=list)
sparse_results: list[QueryResult] = field(default_factory=list)
fused_results: list[QueryResult] = field(default_factory=list)
def detect_document_languages(
vector_store: VectorStore,
llm: Runnable,
*,
max_documents: int = 5,
chunks_per_document: int = 2,
sample_chars: int = 2000,
) -> list[str]:
"""Detect all languages present in the document corpus via the LLM.
Samples chunks from up to ``max_documents`` distinct documents and asks the
LLM in a single call to identify every language present. Used by routers
so that intermediate retrieval queries can be phrased in the corpus
language(s) without hardcoding any specific language.
Args:
vector_store: VectorStore to sample chunks from.
llm: LLM runnable used for the single detection call.
max_documents: Maximum number of documents to sample from.
chunks_per_document: Chunks taken from each sampled document.
sample_chars: Cap on total sample text length sent to the LLM.
Returns:
List of detected language names in English (e.g. ``["Danish"]`` or
``["Danish", "English"]``), preserving the order returned by the LLM.
Returns an empty list when the corpus is empty or no readable text
could be sampled (e.g. when the vector store is mocked in tests).
"""
try:
ids = vector_store.list_document_ids()
except Exception:
return []
if not isinstance(ids, list) or not ids:
return []
samples: list[str] = []
for doc_id in ids[:max_documents]:
try:
chunks = vector_store.get_chunks_by_document_id(doc_id)
except Exception:
continue
if not isinstance(chunks, list):
continue
for c in chunks[:chunks_per_document]:
text = (getattr(c, "text", "") or "").strip()
if text:
samples.append(text)
sample_text = "\n---\n".join(samples)[:sample_chars].strip()
if not sample_text:
return []
prompt = render_prompt("detect_languages", sample_text=sample_text)
raw = _extract_content(llm.invoke(prompt))
seen: set[str] = set()
detected: list[str] = []
for line in raw.strip().splitlines():
name = line.strip().lstrip("-•*0123456789.) ").rstrip(".").strip()
if not name:
continue
name = name.capitalize()
if name.lower() not in seen:
seen.add(name.lower())
detected.append(name)
return detected
def _merge_results(existing: list[QueryResult], new: list[QueryResult]) -> list[QueryResult]:
"""Merge two QueryResult lists by chunk_id, keeping the highest score.
Args:
existing: Previously accumulated results.
new: New results to merge in.
Returns:
Merged list sorted by descending score.
"""
by_id = {r.chunk.chunk_id: r for r in existing}
for r in new:
cid = r.chunk.chunk_id
if cid not in by_id or r.score > by_id[cid].score:
by_id[cid] = r
return sorted(by_id.values(), key=lambda r: r.score, reverse=True)
def _format_results(results: list[QueryResult]) -> str:
"""Format a list of QueryResult into a readable string.
Args:
results: Ranked results to format.
Returns:
Formatted string with numbered entries, or a no-results message.
"""
if not results:
return "Ingen relevante dokumenter fundet. (No relevant documents found.)"
parts: list[str] = []
for i, r in enumerate(results, 1):
page_info = ""
page = r.chunk.metadata.get("page_number")
if page is not None:
page_info = f" side {page}"
parts.append(
f"[{i}] {r.chunk.document_id}{page_info} (relevance: {r.score:.3f})\n{r.chunk.text}"
)
return "\n\n---\n\n".join(parts)
def make_retrieval_tools(
hybrid_retriever: HybridRetriever,
reranker: Reranker,
vector_store: VectorStore,
store: ToolResultStore,
default_top_k: int = 5,
llm_chain: Runnable | None = None,
document_languages: list[str] | None = None,
) -> list:
"""Create retrieval tools bound to the given components and result store.
The returned tools write structured QueryResult objects into *store* on each
invocation so the calling router can surface them as sources without having
to re-parse the tool's text output.
Args:
hybrid_retriever: HybridRetriever instance.
reranker: Reranker instance.
vector_store: VectorStore instance for document-level access.
store: Shared ToolResultStore that captures structured results.
default_top_k: Default number of results to return per call.
llm_chain: Optional LLM chain for tools that need generation
(summarize_document, multi_query_search). When None, those
tools are excluded from the returned list.
document_languages: Detected languages of the document corpus
(e.g. ``["Danish"]`` or ``["Danish", "English"]``). Used by
multi_query_search to phrase sub-queries in the corpus
language(s) for best BM25 recall. When None or empty, the
sub-query language is left unconstrained.
Returns:
List of LangChain tool callables ready for bind_tools / ToolNode.
"""
if document_languages:
if len(document_languages) == 1:
_lang_clause = (
f"The queries should be in {document_languages[0]} "
f"(the document base is {document_languages[0]})."
)
else:
_lang_list = ", ".join(document_languages)
_lang_clause = (
f"The document base contains multiple languages: {_lang_list}. "
f"For each sub-query, write it in whichever of these languages "
f"best matches the topic; mix languages across sub-queries if "
f"the topic is likely covered by documents in different languages."
)
else:
_lang_clause = (
"Write each sub-query in the language most likely used by the "
"underlying documents."
)
# ------------------------------------------------------------------
# Core search tool
# ------------------------------------------------------------------
@tool
def hybrid_search(query: str, top_k: int = default_top_k) -> str:
"""Search the KU document knowledge base using hybrid retrieval.
Combines dense semantic search (Qdrant) and sparse keyword search (BM25),
then re-ranks results with a cross-encoder. Use this tool to find relevant
passages from ingested KU policy documents about rules, regulations, exam
procedures, employment conditions, and administrative guidelines.
Call this tool before answering any question that requires factual
information from KU documents. You may call it multiple times with
different queries if the first result is insufficient.
Args:
query: Search query. Danish gives the best recall against KU documents.
top_k: Number of top results to return (1–20). Default is 5.
Returns:
Formatted string of ranked document passages with source references
and relevance scores.
"""
logger.info("Tool hybrid_search: query=%r top_k=%d", query, top_k)
store.tool_calls.append(("hybrid_search", query))
hybrid_result = hybrid_retriever.search_detailed(query, top_k=top_k)
results = reranker.rerank(query, hybrid_result.fused_results, top_k=top_k)
store.dense_results = _merge_results(store.dense_results, hybrid_result.dense_results)
store.sparse_results = _merge_results(store.sparse_results, hybrid_result.sparse_results)
store.fused_results = _merge_results(store.fused_results, hybrid_result.fused_results)
store.retrieved = _merge_results(store.retrieved, results)
return _format_results(results)
# ------------------------------------------------------------------
# Document-level tools
# ------------------------------------------------------------------
@tool
def list_documents() -> str:
"""List all documents currently available in the KU knowledge base.
Use this tool when the user asks which documents are available, wants to
know what topics are covered, or before fetching a specific document by ID.
Returns:
Newline-separated list of document IDs, or a message if the
knowledge base is empty.
"""
logger.info("Tool list_documents called")
store.tool_calls.append(("list_documents", ""))
ids = vector_store.list_document_ids()
if not ids:
return "Ingen dokumenter i vidensbasen. (Knowledge base is empty.)"
lines = "\n".join(f"- {doc_id}" for doc_id in ids)
return f"Dokumenter i vidensbasen ({len(ids)} i alt):\n{lines}"
@tool
def fetch_document(document_id: str) -> str:
"""Fetch the full text of a specific document from the knowledge base.
Use this tool when the user asks for a summary or overview of a named
document, or when hybrid_search results reference a document that
warrants deeper reading. Prefer hybrid_search for targeted questions.
Args:
document_id: The exact document ID as returned by list_documents or
seen in hybrid_search results (e.g. 'ku_ai_policy.pdf').
Returns:
The concatenated text of all chunks belonging to the document, or
an error message if the document ID is not found.
"""
logger.info("Tool fetch_document: document_id=%r", document_id)
store.tool_calls.append(("fetch_document", document_id))
chunks = vector_store.get_chunks_by_document_id(document_id)
if not chunks:
return (
f"Dokumentet '{document_id}' blev ikke fundet i vidensbasen. "
f"(Document not found. Use list_documents to see available IDs.)"
)
chunks.sort(key=lambda c: c.metadata.get("chunk_index", 0))
existing = {r.chunk.chunk_id: r for r in store.retrieved}
for chunk in chunks:
if chunk.chunk_id not in existing:
existing[chunk.chunk_id] = QueryResult(chunk=chunk, score=1.0, source="fetch_document")
store.retrieved = sorted(existing.values(), key=lambda r: r.score, reverse=True)
full_text = "\n\n".join(c.text for c in chunks)
return (
f"Dokument: {document_id} ({len(chunks)} afsnit)\n\n"
f"{full_text}"
)
# ------------------------------------------------------------------
# Targeted within-document search
# ------------------------------------------------------------------
@tool
def search_within_document(document_id: str, query: str, top_k: int = 3) -> str:
"""Search for specific information within a single document.
Retrieves all chunks belonging to the document and uses the cross-encoder
reranker to find the most relevant passages for the query. Use this when
you already know which document to look in and need to pinpoint the exact
section (e.g. a specific clause, page, or paragraph).
Args:
document_id: The exact document ID to search within.
query: What to look for inside the document.
top_k: Number of top passages to return (1–10). Default is 3.
Returns:
The most relevant passages within the document, ranked by relevance.
"""
logger.info(
"Tool search_within_document: doc=%r query=%r top_k=%d",
document_id, query, top_k,
)
store.tool_calls.append(("search_within_document", f"{document_id}: {query}"))
chunks = vector_store.get_chunks_by_document_id(document_id)
if not chunks:
return (
f"Dokumentet '{document_id}' blev ikke fundet i vidensbasen. "
f"(Document not found. Use list_documents to see available IDs.)"
)
# Wrap chunks as QueryResult so the reranker can score them
candidates = [
QueryResult(chunk=c, score=0.0, source="search_within_document")
for c in chunks
]
results = reranker.rerank(query, candidates, top_k=top_k)
store.retrieved = _merge_results(store.retrieved, results)
return _format_results(results)
# ------------------------------------------------------------------
# LLM-powered tools (only available when llm_chain is provided)
# ------------------------------------------------------------------
tools: list = [hybrid_search, list_documents, fetch_document, search_within_document]
if llm_chain is not None:
@tool
def multi_query_search(question: str, top_k: int = default_top_k) -> str:
"""Decompose a complex question into sub-queries and search each independently.
Use this tool instead of hybrid_search when the question involves
multiple aspects, comparisons, or requires information from different
topics. For example: "How do exam rules differ between bachelor and
master programmes?" would be split into separate searches for each
programme's exam rules, then merged.
Args:
question: The complex user question to decompose and search.
top_k: Number of results to return per sub-query (1–10). Default is 5.
Returns:
Combined results from all sub-queries, deduplicated and ranked.
"""
logger.info("Tool multi_query_search: question=%r", question)
store.tool_calls.append(("multi_query_search", question))
# Step 1: Ask LLM to decompose the question
decompose_prompt = render_prompt(
"multi_query_decompose",
lang_clause=_lang_clause,
question=question,
)
raw = _extract_content(llm_chain.invoke(decompose_prompt))
sub_queries = [q.strip().lstrip("0123456789.-) ") for q in raw.splitlines() if q.strip()]
if not sub_queries:
sub_queries = [question]
logger.info("Decomposed into %d sub-queries: %s", len(sub_queries), sub_queries)
# Step 2: Search each sub-query independently
all_results: list[QueryResult] = []
for sq in sub_queries:
hybrid_result = hybrid_retriever.search_detailed(sq, top_k=top_k)
reranked = reranker.rerank(sq, hybrid_result.fused_results, top_k=top_k)
all_results = _merge_results(all_results, reranked)
store.dense_results = _merge_results(store.dense_results, hybrid_result.dense_results)
store.sparse_results = _merge_results(store.sparse_results, hybrid_result.sparse_results)
store.fused_results = _merge_results(store.fused_results, hybrid_result.fused_results)
# Step 3: Keep top results across all sub-queries
final = all_results[:top_k]
store.retrieved = _merge_results(store.retrieved, final)
header = f"Søgning opdelt i {len(sub_queries)} delforespørgsler:\n"
header += "\n".join(f" • {sq}" for sq in sub_queries)
header += "\n\n"
return header + _format_results(final)
@tool
def summarize_document(document_id: str) -> str:
"""Generate a structured summary of a document in the knowledge base.
Fetches the full document and uses the LLM to produce a concise summary
covering the main topics, key rules, and important details. Use this
when the user asks "what is this document about?" or wants an overview
before diving into specifics.
Args:
document_id: The exact document ID to summarize.
Returns:
A structured summary of the document, or an error if not found.
"""
logger.info("Tool summarize_document: document_id=%r", document_id)
store.tool_calls.append(("summarize_document", document_id))
chunks = vector_store.get_chunks_by_document_id(document_id)
if not chunks:
return (
f"Dokumentet '{document_id}' blev ikke fundet i vidensbasen. "
f"(Document not found. Use list_documents to see available IDs.)"
)
chunks.sort(key=lambda c: c.metadata.get("chunk_index", 0))
full_text = "\n\n".join(c.text for c in chunks)
# Register chunks as sources
existing = {r.chunk.chunk_id: r for r in store.retrieved}
for chunk in chunks:
if chunk.chunk_id not in existing:
existing[chunk.chunk_id] = QueryResult(
chunk=chunk, score=1.0, source="summarize_document",
)
store.retrieved = sorted(existing.values(), key=lambda r: r.score, reverse=True)
# Truncate to avoid exceeding context limits
max_chars = 8000
if len(full_text) > max_chars:
full_text = full_text[:max_chars] + "\n\n[... teksten er forkortet ... (text truncated)]"
summary_prompt = render_prompt(
"summarize_document",
document_id=document_id,
full_text=full_text,
)
summary = _extract_content(llm_chain.invoke(summary_prompt))
return f"Resumé af {document_id}:\n\n{summary}"
tools.extend([multi_query_search, summarize_document])
return tools