251 lines
9.1 KiB
Python
251 lines
9.1 KiB
Python
"""Agent tools - callable operations for the agentic research engine."""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Any
|
|
import httpx
|
|
|
|
from app.config import get_settings
|
|
from app.db.database import db
|
|
|
|
|
|
async def read_document(doc_id: str) -> dict[str, Any]:
|
|
"""Read complete document record."""
|
|
doc = await db.get_document(doc_id)
|
|
if not doc:
|
|
return {"error": "Document not found"}
|
|
return {"document": doc}
|
|
|
|
|
|
async def read_chunks(doc_id: str, page_num: int | None = None, limit: int = 50) -> str:
|
|
"""Read text chunks for a document."""
|
|
chunks = await db.get_doc_chunks(doc_id)
|
|
if page_num is not None:
|
|
chunks = [c for c in chunks if c.get("page_num") == page_num]
|
|
chunks = chunks[:limit]
|
|
parts = []
|
|
for c in chunks:
|
|
page_part = f"[p{c['page_num']}]" if c.get("page_num") else ""
|
|
parts.append(f"{page_part} {c['content'][:500]}")
|
|
return "\n\n".join(parts)
|
|
|
|
|
|
async def get_page_text(doc_id: str, page_num: int) -> str:
|
|
"""Extract text for a specific page, including polygon metadata."""
|
|
chunks = await db.get_doc_chunks(doc_id)
|
|
page_chunks = [c for c in chunks if c.get("page_num") == page_num]
|
|
if not page_chunks:
|
|
return f"No content available for page {page_num}"
|
|
|
|
# Build page view with polygon info
|
|
text_parts = []
|
|
for c in page_chunks:
|
|
polygon = c.get("polygon")
|
|
poly_info = ""
|
|
if polygon and isinstance(polygon, dict):
|
|
bbox = polygon.get("bbox", [])
|
|
if bbox:
|
|
poly_info = f" [bbox:{bbox[0]:.0f},{bbox[1]:.0f},{bbox[2]:.0f},{bbox[3]:.0f}]"
|
|
text_parts.append(f"[{c['block_index']}{poly_info}] {c['content'][:300]}")
|
|
|
|
return f"--- Page {page_num} ---\n" + "\n".join(text_parts)
|
|
|
|
|
|
async def vector_search(query: str, doc_id: str | None = None, limit: int = 20) -> str:
|
|
"""Search document chunks by semantic similarity."""
|
|
settings = get_settings()
|
|
async with httpx.AsyncClient(timeout=30) as client:
|
|
query_vec = []
|
|
# Try to get embedding via Ollama
|
|
try:
|
|
resp = await client.post(
|
|
f"{settings.ollama_url}/api/embeddings",
|
|
json={"model": "nomic-embed-text", "prompt": query},
|
|
)
|
|
query_vec = resp.json().get("embedding", [])
|
|
except Exception:
|
|
# Fall back to local embedding
|
|
from app.core import get_embedding
|
|
query_vec = get_embedding(query)
|
|
|
|
if not query_vec:
|
|
return "No embeddings available. Using text search fallback."
|
|
|
|
results = await db.vector_search(query_vec, doc_id=doc_id, limit=limit)
|
|
parts = []
|
|
for r in results:
|
|
score = r.get("similarity", 0)
|
|
page = r.get("page_num", "?")
|
|
parts.append(f"[{score:.2f}] (p{page}) {r['content'][:200]}")
|
|
return "\n".join(parts) if parts else "No similar content found."
|
|
|
|
|
|
async def text_search(query: str, doc_id: str | None = None, limit: int = 20) -> str:
|
|
"""Search chunks by keyword match."""
|
|
results = await db.search_chunks_text(query, doc_id=doc_id, limit=limit)
|
|
parts = []
|
|
for r in results:
|
|
rank = r.get("rank", 0)
|
|
page = r.get("page_num", "?")
|
|
parts.append(f"[{rank:.2f}] (p{page}) {r['content'][:200]}")
|
|
return "\n".join(parts) if parts else "No text matches found."
|
|
|
|
|
|
async def extract_facts(doc_id: str, question: str | None = None) -> str:
|
|
"""Extract key facts from document."""
|
|
chunks = await db.get_doc_chunks(doc_id)
|
|
# Extract from chunk metadata
|
|
facts = []
|
|
for c in chunks:
|
|
content = c.get("content", "")
|
|
if question and question.lower() not in content.lower():
|
|
continue
|
|
if "fact" in c.get("chunk_type", "").lower():
|
|
facts.append(content[:200])
|
|
elif len(content.strip()) > 50:
|
|
facts.append(content[:200])
|
|
|
|
seen = set()
|
|
unique = []
|
|
for f in facts:
|
|
key = f[:40]
|
|
if key not in seen:
|
|
seen.add(key)
|
|
unique.append(f)
|
|
|
|
return "\n".join(f"I: {f}" for f in unique[:20])
|
|
|
|
|
|
async def get_polygon_view(doc_id: str, page_num: int, block_index: int | None = None) -> dict:
|
|
"""Get page with polygon coordinates for visualization."""
|
|
chunks = await db.get_doc_chunks(doc_id)
|
|
page_chunks = [c for c in chunks if c.get("page_num") == page_num]
|
|
|
|
blocks = []
|
|
for c in page_chunks:
|
|
poly = c.get("polygon")
|
|
if poly and isinstance(poly, dict):
|
|
blocks.append({
|
|
"index": c.get("block_index", 0),
|
|
"page": page_num,
|
|
"polygon": poly,
|
|
"content": c.get("content", "")[:200],
|
|
"type": c.get("chunk_type", "text"),
|
|
})
|
|
elif block_index is None or c.get("block_index") == block_index:
|
|
blocks.append({
|
|
"index": c.get("block_index", 0),
|
|
"page": page_num,
|
|
"polygon": {"bbox": [0, 0, 1000, 1000]},
|
|
"content": c.get("content", "")[:200],
|
|
"type": c.get("chunk_type", "text"),
|
|
})
|
|
|
|
# Compute page bounding box from all polygons
|
|
all_polys = [b["polygon"] for b in blocks if b["polygon"] and isinstance(b["polygon"], dict)]
|
|
|
|
pages_data = {"page_num": page_num, "blocks": blocks}
|
|
|
|
if all_polys:
|
|
min_x = min(p.get("bbox", [0, 0, 0, 0])[:2] for p in all_polys)
|
|
max_x = max(p.get("bbox", [0, 0, 0, 0])[2:] for p in all_polys)
|
|
pages_data["page_bbox"] = {
|
|
"x_min": min(min_x),
|
|
"y_min": min(min_x, key=lambda x: x[1])[1],
|
|
"width": max(max_x) - min(min_x),
|
|
"height": max(max_x, key=lambda x: x[1])[1] - min(min_x, key=lambda x: x[1])[1],
|
|
}
|
|
|
|
return pages_data
|
|
|
|
|
|
async def get_memories(session_id: str) -> list[dict]:
|
|
"""Get research memories from pgvector."""
|
|
return await db.get_memories(session_id)
|
|
|
|
|
|
async def memory_similarity_search(query: str, limit: int = 10) -> str:
|
|
"""Search across all research memories."""
|
|
results = await db.memory_similarity_search(query, limit=limit)
|
|
parts = []
|
|
for r in results:
|
|
parts.append(f"[{r.get('importance', '?')}] ({r.get('memory_type', '?')}) {r['content'][:200]}")
|
|
return "\n".join(parts) if parts else "No memories found."
|
|
|
|
|
|
async def list_documents(status: str | None = None) -> str:
|
|
"""List all indexed documents."""
|
|
docs = await db.list_documents(status=status)
|
|
parts = []
|
|
for d in docs:
|
|
parts.append(f"- [{d['status']}] {d['filename']} (p{d.get('page_count', '?')}, {d['created_at']})")
|
|
return "\n".join(parts) if parts else "No documents found."
|
|
|
|
|
|
async def save_finding(session_id: str, question: str, answer: str,
|
|
summary: str, confidence: float = 0.8,
|
|
relevant_chunks: list | None = None) -> str:
|
|
"""Save a research finding to the database."""
|
|
finding_id = await db.store_finding(
|
|
session_id, question, answer, summary,
|
|
"researcher", confidence, relevant_chunks
|
|
)
|
|
await db.update_session(session_id, findings={"count": 1})
|
|
return finding_id
|
|
|
|
|
|
async def save_structured_evidence(
|
|
session_id: str, rows: list[dict]
|
|
) -> list[str]:
|
|
"""Save evidence-extraction rows with traceability fields."""
|
|
ids = await db.save_structured_evidence(session_id, rows)
|
|
return ids
|
|
|
|
|
|
async def get_pipeline_state(session_id: str) -> list[dict]:
|
|
"""Get saved pipeline intermediate stages."""
|
|
return await db.get_pipeline_stages(session_id)
|
|
|
|
|
|
async def merge_evidence(rows_a: list[dict], rows_b: list[dict]) -> list[dict]:
|
|
"""Merge two evidence sets, flagging agreement/conflict by topic."""
|
|
by_topic: dict[str, list] = {}
|
|
for r in rows_a + rows_b:
|
|
topic = r.get("topic", "unknown")
|
|
by_topic.setdefault(topic, []).append(r)
|
|
merged = []
|
|
for topic, rs in by_topic.items():
|
|
if len(rs) == 2:
|
|
merged.append({
|
|
"topic": topic,
|
|
"source_a": rs[0].get("evidence", ""),
|
|
"source_b": rs[1].get("evidence", ""),
|
|
"agreement": "Yes" if rs[0].get("evidence") == rs[1].get("evidence") else "No",
|
|
"conflict": "Yes" if rs[0].get("confidence") != rs[1].get("confidence") else "No",
|
|
"notes": f"Combined from {rs[0].get('source_doc','')} and {rs[1].get('source_doc','')}",
|
|
})
|
|
else:
|
|
merged.append({k: rs[0].get(k, "") for k in ("topic", "description", "trace_ref", "evidence", "confidence")})
|
|
return merged
|
|
|
|
|
|
# Tool registry
|
|
TOOLS = {
|
|
"read_document": read_document,
|
|
"read_chunks": read_chunks,
|
|
"get_page_text": get_page_text,
|
|
"vector_search": vector_search,
|
|
"text_search": text_search,
|
|
"extract_facts": extract_facts,
|
|
"get_polygon_view": get_polygon_view,
|
|
"get_memories": get_memories,
|
|
"memory_similarity_search": memory_similarity_search,
|
|
"list_documents": list_documents,
|
|
"save_finding": save_finding,
|
|
"save_structured_evidence": save_structured_evidence,
|
|
"get_pipeline_state": get_pipeline_state,
|
|
"merge_evidence": merge_evidence,
|
|
}
|