first commit

app/agents/engine.py

@@ -0,0 +1,417 @@
+"""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