418 lines
15 KiB
Python
418 lines
15 KiB
Python
"""Agentic research engine - coordinates agent skills and Ollama LLM interactions."""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import re
|
|
from typing import Any
|
|
import httpx
|
|
|
|
from app.agents.skills import SKILLS, Skill
|
|
from app.agents.tools import TOOLS
|
|
from app.config import get_settings
|
|
|
|
|
|
class Researcher:
|
|
"""Main research orchestrator that coordinates agents and tools."""
|
|
|
|
def __init__(self):
|
|
self.settings = get_settings()
|
|
self.active_skills: list[Skill] = []
|
|
self.findings: list[dict] = []
|
|
self.context: dict = {}
|
|
|
|
def select_skills(self, query: str, skill_names: list[str] | None = None) -> list[str]:
|
|
"""Select relevant agent skills for the query. Defaults to all if none specified."""
|
|
if skill_names:
|
|
selected = []
|
|
for name in skill_names:
|
|
s = SKILLS.get(name)
|
|
if s:
|
|
selected.append(name)
|
|
return selected
|
|
|
|
query_lower = query.lower()
|
|
selected = []
|
|
for name, skill in SKILLS.items():
|
|
for token in skill.verb_tokens:
|
|
if token in query_lower:
|
|
selected.append(name)
|
|
break
|
|
# If no token matched, default to researcher + qa
|
|
if not selected:
|
|
selected = ["researcher", "qa_agent"]
|
|
return selected
|
|
|
|
async def call_ollama(self, prompt: str, model: str | None = None, system: str | None = None) -> str:
|
|
"""Call Ollama LLM with streaming support."""
|
|
model = model or self.settings.gpt_oss_model
|
|
messages = []
|
|
if system:
|
|
messages.append({"role": "system", "content": system})
|
|
messages.append({"role": "user", "content": prompt})
|
|
|
|
async with httpx.AsyncClient(timeout=300) as client:
|
|
resp = await client.post(
|
|
f"{self.settings.ollama_url}/api/chat",
|
|
json={"model": model, "messages": messages, "stream": False},
|
|
)
|
|
resp.raise_for_status()
|
|
return resp.json().get("message", {}).get("content", "")
|
|
|
|
async def run_skill(
|
|
self, skill_name: str, query: str, doc_id: str | None = None
|
|
) -> str:
|
|
"""Execute a single agent skill."""
|
|
skill = SKILLS.get(skill_name)
|
|
if not skill:
|
|
return f"[Unknown skill: {skill_name}]"
|
|
|
|
system_prompt = f"""You are the {skill_name.replace('_', ' ').title()} agent.
|
|
{skill.instructions}"""
|
|
|
|
tool_prompts = []
|
|
for tool_name in skill.tools:
|
|
tool_fn = TOOLS.get(tool_name)
|
|
if not tool_fn:
|
|
continue
|
|
import inspect
|
|
sig = inspect.signature(tool_fn)
|
|
try:
|
|
if tool_name == "read_document" and doc_id:
|
|
result = await tool_fn(doc_id)
|
|
elif tool_name == "read_chunks" and doc_id:
|
|
result = await tool_fn(doc_id, limit=20)
|
|
elif tool_name == "get_page_text" and doc_id:
|
|
result = await tool_fn(doc_id, page_num=0)
|
|
elif tool_name == "extract_facts" and doc_id:
|
|
result = await tool_fn(doc_id, question=query)
|
|
elif tool_name == "list_documents":
|
|
result = await tool_fn()
|
|
elif tool_name in ("vector_search", "text_search", "memory_similarity_search"):
|
|
kwargs = {"query": query}
|
|
if doc_id and "doc_id" in sig.parameters:
|
|
kwargs["doc_id"] = doc_id
|
|
result = await tool_fn(**kwargs)
|
|
elif tool_name in ("get_findings", "get_memories", "get_pipeline_state"):
|
|
result = await tool_fn(session_id=doc_id or "")
|
|
elif tool_name == "save_finding":
|
|
result = await tool_fn(session_id=doc_id or "", question=query, answer="", finding_type="research", confidence=0.5)
|
|
else:
|
|
kwargs = {}
|
|
if "doc_id" in sig.parameters and doc_id:
|
|
kwargs["doc_id"] = doc_id
|
|
if "query" in sig.parameters:
|
|
kwargs["query"] = query
|
|
result = await tool_fn(**kwargs)
|
|
tool_prompts.append(f"\n--- {tool_name} output ---\n{result}")
|
|
except Exception as e:
|
|
tool_prompts.append(f"\n--- {tool_name} error ---\n{str(e)}")
|
|
|
|
context = "".join(tool_prompts)
|
|
|
|
prompt = f"""Research query: {query}
|
|
Document ID: {doc_id}
|
|
Context from tools:
|
|
{context}
|
|
|
|
Provide your analysis following the {skill_name} protocol."""
|
|
|
|
response = await self.call_ollama(prompt, system=system_prompt)
|
|
|
|
self.findings.append({
|
|
"skill": skill_name,
|
|
"query": query,
|
|
"response": response,
|
|
"doc_id": doc_id,
|
|
"timestamp": str(os.popen('date +%Y-%m-%dT%H:%M:%S').read()).strip(),
|
|
})
|
|
|
|
return response
|
|
|
|
async def run_research_session(
|
|
self, query: str, session_id: str,
|
|
doc_id: str | None = None,
|
|
skill_names: list[str] | None = None
|
|
) -> dict[str, str]:
|
|
"""Run a full research session with selected agents."""
|
|
selected_skills = self.select_skills(query, skill_names)
|
|
results = {}
|
|
|
|
for skill_name in selected_skills:
|
|
results[skill_name] = await self.run_skill(skill_name, query, doc_id)
|
|
|
|
return results
|
|
|
|
async def cross_reference(self, query: str, doc_ids: list[str]) -> str:
|
|
"""Cross-reference content across multiple documents."""
|
|
contexts = []
|
|
for doc_id in doc_ids:
|
|
chunks = await TOOLS["read_chunks"](doc_id)
|
|
contexts.append(f"--- {doc_id} ---\n{chunks[:2000]}")
|
|
|
|
combined = "\n\n".join(contexts)
|
|
prompt = f"""Cross-reference these documents for the query: {query}
|
|
|
|
{combined}
|
|
|
|
Direct comparison. No preamble."""
|
|
|
|
return await self.call_ollama(prompt, system="You are a cross-reference analyst. Output concise, comparative findings.")
|
|
|
|
|
|
# ── Pipeline skills ────────────────────────────────
|
|
|
|
PIPELINE_STAGES = ("document_triage", "evidence_extraction", "research_synthesis")
|
|
|
|
|
|
async def _parse_triage_output(response: str) -> dict:
|
|
"""Parse document-triage output into structured dict."""
|
|
obj = {}
|
|
for section in ["OBJECTIVE", "SUB-QUESTIONS", "CLASSIFICATION", "READING_ORDER", "EXTRACTION_CRITERIA", "RISKS_AND_GAPS"]:
|
|
marker = f"[{section}]"
|
|
# Find start of this section
|
|
start = response.find(marker)
|
|
end = response.find("[", start + len(marker)) if start != -1 else -1
|
|
if end != -1:
|
|
chunk = response[start + len(marker):end].strip()
|
|
elif start != -1:
|
|
chunk = response[start + len(marker):].strip()
|
|
else:
|
|
chunk = ""
|
|
obj[section] = chunk
|
|
return obj
|
|
|
|
|
|
async def _parse_evidence_output(response: str) -> list[dict]:
|
|
"""Parse evidence-extraction output into structured list."""
|
|
rows = []
|
|
# Find [EVIDENCE_ROWS] section
|
|
start_marker = "[EVIDENCE_ROWS]"
|
|
start = response.find(start_marker)
|
|
# Check if there's a cross-comparison section
|
|
cross_start = response.find("[CROSS_COMPARISON]")
|
|
end = cross_start if cross_start != -1 else len(response)
|
|
if start == -1:
|
|
start = 0
|
|
section = response[start:end].strip()
|
|
|
|
lines = section.split("\n")
|
|
for line in lines:
|
|
line = line.strip()
|
|
if not line or line.startswith("["):
|
|
continue
|
|
parts = [p.strip() for p in line.split("|")]
|
|
if len(parts) >= 7:
|
|
rows.append({
|
|
"topic": parts[1],
|
|
"evidence_type": parts[2],
|
|
"description": parts[3],
|
|
"trace_ref": parts[4],
|
|
"evidence": parts[5],
|
|
"analyst_note": parts[6],
|
|
"confidence": parts[7] if len(parts) > 7 else "Medium",
|
|
"review_needed": parts[8] if len(parts) > 8 else "No",
|
|
})
|
|
return rows
|
|
|
|
|
|
async def _parse_synthesis_mode(response: str) -> str:
|
|
"""Determine output mode from synthesis response content."""
|
|
for mode in ["brief", "report", "gap analysis", "matrix"]:
|
|
if mode in response.lower():
|
|
return mode
|
|
return "brief"
|
|
|
|
|
|
class ResearchPipeline:
|
|
"""Pipeline orchestrator: document_triage → evidence_extraction → research_synthesis."""
|
|
|
|
def __init__(self):
|
|
self.settings = get_settings()
|
|
self.plan: dict = {}
|
|
self.evidence: list[dict] = []
|
|
self.synthesis: str = ""
|
|
self.findings: list[dict] = []
|
|
|
|
async def run(
|
|
self,
|
|
query: str,
|
|
session_id: str,
|
|
db,
|
|
doc_ids: list[str] | None = None,
|
|
output_mode: str | None = None,
|
|
) -> dict:
|
|
"""Run all three pipeline stages sequentially, persisting after each."""
|
|
results = {}
|
|
|
|
# ── Stage 1: document_triage ─────────────────────
|
|
triage_skill = SKILLS["document_triage"]
|
|
system = f"You are document_triage. {triage_skill.instructions}"
|
|
|
|
# Gather context from available documents
|
|
if doc_ids:
|
|
doc_context = []
|
|
for did in doc_ids:
|
|
doc_info = await TOOLS["read_document"](did)
|
|
doc_context.append(str(doc_info))
|
|
context_input = "\n".join(doc_context)
|
|
else:
|
|
doc_list = await TOOLS["list_documents"]()
|
|
context_input = doc_list
|
|
|
|
triage_prompt = f"""Research query: {query}
|
|
Context from tools:
|
|
{context_input}
|
|
|
|
Apply the document-tireage protocol."""
|
|
|
|
triage_output = await self._call_ollama(triage_prompt, system=system)
|
|
self.plan = await _parse_triage_output(triage_output)
|
|
self.findings.append({
|
|
"stage": "triage",
|
|
"output": triage_output,
|
|
"state": self.plan,
|
|
})
|
|
await db.save_pipeline_stage(session_id, "triage", triage_output, self.plan)
|
|
results["triage"] = triage_output
|
|
results["triage_state"] = self.plan
|
|
|
|
# ── Stage 2: evidence_extraction ────────────────
|
|
evidence_skill = SKILLS["evidence_extraction"]
|
|
system = f"""You are evidence_extraction. {evidence_skill.instructions}
|
|
|
|
Triage plan (from previous stage):
|
|
{json.dumps(self.plan, indent=2, default=str)}
|
|
|
|
Focus on extracting evidence for the sub-questions and extraction criteria defined above."""
|
|
|
|
# Read chunks from all relevant documents
|
|
all_chunks = []
|
|
for did in (doc_ids or []):
|
|
chunks = await TOOLS["read_chunks"](did)
|
|
all_chunks.append(f"--- Document {did} ---\n{chunks}")
|
|
|
|
context_input = "\n\n".join(all_chunks) if all_chunks else "No documents loaded yet."
|
|
|
|
evidence_prompt = f"""Research query: {query}
|
|
Sub-questions to address:
|
|
{self.plan.get('SUB-QUESTIONS', 'N/A')}
|
|
|
|
Context from tools:
|
|
{context_input}
|
|
|
|
Apply the evidence_extraction protocol. Return structured evidence rows."""
|
|
|
|
evidence_output = await self._call_ollama(evidence_prompt, system=system)
|
|
self.evidence = await _parse_evidence_output(evidence_output)
|
|
|
|
# Persist to DB
|
|
if self.evidence:
|
|
await db.save_structured_evidence(session_id, self.evidence)
|
|
await db.save_pipeline_stage(session_id, "evidence", evidence_output, {"row_count": len(self.evidence)})
|
|
|
|
results["evidence"] = evidence_output
|
|
results["evidence_rows"] = self.evidence
|
|
|
|
# ── Stage 3: research_synthesis ─────────────────
|
|
synthesis_skill = SKILLS["research_synthesis"]
|
|
|
|
mode = output_mode or await _parse_synthesis_mode(evidence_output)
|
|
mode_prompts = {
|
|
"brief": "Use Research Brief output mode.",
|
|
"report": "Use Research Report output mode.",
|
|
"gap": "Use Gap Analysis output mode.",
|
|
"matrix": "Use Comparison Matrix output mode.",
|
|
}
|
|
mode_instruct = mode_prompts.get(mode, mode_prompts["brief"])
|
|
|
|
system = f"""You are research_synthesis. {synthesis_skill.instructions}
|
|
|
|
{mode_instruct}
|
|
|
|
Extracted evidence (from previous stage):
|
|
{json.dumps(self.evidence, indent=2, default=str)[:15000]}
|
|
|
|
Apply the research_synthesis protocol."""
|
|
|
|
synthesis_prompt = f"""Research query: {query}
|
|
|
|
Sub-questions:
|
|
{self.plan.get('SUB-QUESTIONS', 'N/A')}
|
|
|
|
Apply the research_synthesis protocol."""
|
|
|
|
synthesis_output = await self._call_ollama(synthesis_prompt, system=system)
|
|
|
|
self.synthesis = synthesis_output
|
|
synthesis_state = {
|
|
"mode": mode,
|
|
"sub_questions": self.plan.get("SUB-QUESTIONS", ""),
|
|
"evidence_count": len(self.evidence),
|
|
}
|
|
await db.save_pipeline_stage(session_id, "synthesis", synthesis_output, synthesis_state)
|
|
|
|
results["synthesis"] = synthesis_output
|
|
results["output_mode"] = mode
|
|
results["pipeline_complete"] = True
|
|
|
|
self.findings.append({
|
|
"stage": "synthesis",
|
|
"output": synthesis_output,
|
|
"state": synthesis_state,
|
|
})
|
|
|
|
return results
|
|
|
|
async def _call_ollama(self, prompt: str, system: str) -> str:
|
|
"""Call Ollama LLM."""
|
|
messages = [
|
|
{"role": "system", "content": system},
|
|
{"role": "user", "content": prompt},
|
|
]
|
|
async with httpx.AsyncClient(timeout=600) as client:
|
|
resp = await client.post(
|
|
f"{self.settings.ollama_url}/api/chat",
|
|
json={"model": self.settings.gpt_oss_model, "messages": messages, "stream": False},
|
|
)
|
|
resp.raise_for_status()
|
|
return resp.json().get("message", {}).get("content", "")
|
|
|
|
|
|
async def render_pipeline_results(self, results: dict) -> list[dict]:
|
|
"""Render pipeline results for frontend display."""
|
|
sections = []
|
|
|
|
# Triage stage
|
|
if "triage_state" in results:
|
|
plan = results["triage_state"]
|
|
sections.append({
|
|
"stage": "triage",
|
|
"title": "Stage 1: Source Triage",
|
|
"objective": plan.get("OBJECTIVE", ""),
|
|
"sub_questions": plan.get("SUB-QUESTIONS", ""),
|
|
"classification": plan.get("CLASSIFICATION", ""),
|
|
"reading_order": plan.get("READING_ORDER", ""),
|
|
"extraction_criteria": plan.get("EXTRACTION_CRITERIA", ""),
|
|
"risks_gaps": plan.get("RISKS_AND_GAPS", ""),
|
|
})
|
|
|
|
# Evidence stage
|
|
if "evidence_rows" in results:
|
|
rows = results["evidence_rows"]
|
|
sections.append({
|
|
"stage": "evidence",
|
|
"title": f"Stage 2: Extracted Evidence ({len(rows)} rows)",
|
|
"rows": rows,
|
|
})
|
|
|
|
# Synthesis stage
|
|
if "synthesis" in results:
|
|
sections.append({
|
|
"stage": "synthesis",
|
|
"title": f"Stage 3: Research Synthesis (mode: {results.get('output_mode', 'auto')})",
|
|
"synthesis": results["synthesis"],
|
|
})
|
|
|
|
return sections
|