agentic/main.py

"""FastAPI web routes for agentic research app."""
from __future__ import annotations

import os
import json
import uuid
import asyncio
import shutil
from pathlib import Path
from typing import Any

import httpx
from fastapi import FastAPI, File, UploadFile, HTTPException, Query
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel

from app.config import get_settings
from app.db.database import Database
from app.agents.engine import Researcher, ResearchPipeline
from app.agents.skills import SKILLS
from app.agents.tools import TOOLS
from app.core.processor import MarkerProcessor
from app.core.embedding_engine import get_embedding_sync

settings = get_settings()
app = FastAPI(title="Agentic Research", version="0.1.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

db_conn: Database | None = None
researcher = Researcher()
processor = MarkerProcessor()


# ── Models ──────────────────────────────────────────────

class ResearchRequest(BaseModel):
    query: str
    doc_id: str | None = None
    skills: list[str] | None = None

class SessionList(BaseModel):
    sessions: list


# ── Lifecycle ───────────────────────────────────────────

@app.on_event("startup")
async def on_startup():
    global db_conn
    pool = await Database.create_pool()
    db_conn = Database(pool)
    from app.db import database as _db_mod
    _db_mod.db = db_conn
    try:
        await db_conn.init_schema()
    except Exception as e:
        print(f"Schema init (non-fatal): {e}")
    
    # Create directories
    for d in [settings.workspace_dir, settings.documents_dir]:
        os.makedirs(d, exist_ok=True)


# ── Frontend ───────────────────────────────────────────

@app.get("/", response_class=HTMLResponse)
async def index():
    return FileResponse("frontend/index.html")

app.mount("/static", StaticFiles(directory="frontend"), name="static")

@app.get("/health")
async def health():
    return {
        "status": "ok",
        "ollama": settings.ollama_url,
        "marker_api": settings.marker_api_url,
        "db": "connected" if db_conn else "disconnected",
    }


# ── Documents ──────────────────────────────────────────

@app.post("/api/documents/upload")
async def upload_document(
    file: UploadFile = File(None),
    pdf_url: str | None = None,
):
    if not file and not pdf_url:
        raise HTTPException(400, "Provide file or pdf_url")
    
    if db_conn is None:
        raise HTTPException(500, "Database not initialized")
    
    doc_id = str(uuid.uuid4())
    filename = file.filename if file else pdf_url.split("/")[-1] if pdf_url else "unknown"
    
    # Read content
    text_content = ""
    file_data = b""
    
    if file:
        file_data = await file.read()
        ext = Path(filename).suffix.lower()
        if ext == ".pdf":
            pass  # PDF processing below
        else:
            # Store plain text directly
            text_content = file_data.decode("utf-8", errors="replace")
            file_path = os.path.join(settings.documents_dir, filename)
            os.makedirs(settings.documents_dir, exist_ok=True)
            with open(file_path, "wb") as f:
                f.write(file_data)
            
            # Chunk the text
            import re
            paragraphs = re.split(r'\n\s*\n', text_content)
            chunks = []
            for i, para in enumerate(paragraphs):
                if len(para.strip()) > 20:
                    chunks.append({
                        "content": para.strip(), "page_num": 0,
                        "block_index": i, "polygon": None, "chunk_type": "text",
                    })
            if chunks:
                await db_conn.batch_chunk(doc_id, chunks)
            
            await db_conn.upsert_document(
                filename, None, file.content_type or "text/plain",
                file_path, "indexed", 1, text_content[:50000], {"uploaded": True}
            )
            return {"doc_id": doc_id, "filename": filename, "chunks": len(chunks), "strategy": "direct_text"}
    
    # PDF processing
    try:
        file_path = os.path.join(settings.documents_dir, filename)
        os.makedirs(settings.documents_dir, exist_ok=True)
        with open(file_path, "wb") as f:
            f.write(file_data)
        
        # Create document record first (so FK constraint is satisfied)
        doc_uuid = str(uuid.uuid4())
        await db_conn.upsert_document(
            filename, doc_uuid, "application/pdf",
            file_path, "processing", 0,
            "", {"status": "processing"}
        )
        
        # Process through marker, using the real PK as doc_id
        result = await processor.process_document_file(file_data, filename, doc_uuid)
        
        # Update document with actual page count and result
        await db_conn.upsert_document(
            filename, doc_uuid, "application/pdf",
            file_path, "indexed", result.get("page_count", 0),
            json.dumps({"marker_result": result})[:10000], {"uploaded": True}
        )
        return {**result, "doc_id": doc_uuid, "filename": filename, "strategy": "marker_ocr"}
    except Exception as e:
        raise HTTPException(500, f"OCR failed: {str(e)}")


@app.get("/api/documents")
async def list_documents():
    if db_conn is None:
        return {"documents": []}
    docs = await db_conn.list_documents()
    return {"documents": docs}


@app.get("/api/documents/{doc_id}/chunks")
async def get_document_chunks(doc_id: str, page: int | None = None):
    if db_conn is None:
        return {"chunks": []}
    chunks = await db_conn.get_doc_chunks(doc_id)
    if page is not None:
        chunks = [c for c in chunks if c.get("page_num") == page]
    return {"chunks": chunks}


@app.get("/api/documents/{doc_id}")
async def get_document(doc_id: str):
    if db_conn is None:
        return {"document": None}
    doc = await db_conn.get_document(doc_id)
    return {"document": doc}


# ── Polygons / View ───────────────────────────────────

@app.get("/api/documents/{doc_id}/polygon-view")
async def get_polygon_view(doc_id: str, page: int = 0):
    view = await TOOLS["get_polygon_view"](doc_id, page)
    return view


@app.get("/api/documents/{doc_id}/page-text")
async def get_page_text(doc_id: str, page: int = 0):
    text = await TOOLS["get_page_text"](doc_id, page)
    return {"page": page, "text": text}


# ── Research ───────────────────────────────────────────

@app.post("/api/research/session")
async def create_session(body: dict):
    if db_conn is None:
        raise HTTPException(500, "DB not ready")
    query = body.get("query", "research query") if isinstance(body, dict) else str(body)
    session_id = await db_conn.create_session(query)
    return {"session_id": session_id, "query": query}


@app.post("/api/research/run")
async def run_research(req: ResearchRequest):
    if db_conn is None:
        raise HTTPException(500, "DB not ready")
    
    researcher = Researcher()
    results = await researcher.run_research_session(
        req.query, req.doc_id or "temp", skill_names=req.skills
    )
    
    # Save findings
    overall_answer = ""
    for skill, response in results.items():
        await db_conn.store_finding(
            req.doc_id or "temp", req.query, response, "Research complete",
            "researcher", 0.8
        )
        if response:
            overall_answer += f"### {skill}:\n{response}\n\n"
    
    return {"results": results, "doc_id": req.doc_id, "session_id": req.doc_id or "temp"}


@app.post("/api/research/semantic")
async def semantic_search(query: str, doc_id: str | None = None, limit: int = 20):
    vec = get_embedding_sync(query, settings.ollama_url)
    results = await db_conn.vector_search(vec, doc_id=doc_id, limit=limit)
    return {"results": results, "query": query}


@app.post("/api/research/text-search")
async def text_search(query: str, doc_id: str | None = None, limit: int = 20):
    results = await db_conn.search_chunks_text(query, doc_id=doc_id, limit=limit)
    return {"results": results, "query": query}


@app.get("/api/research/findings")
async def get_findings(session_id: str):
    if db_conn:
        findings = await db_conn.get_findings(session_id)
        return {"findings": findings}
    return {"findings": []}


@app.get("/api/research/memories")
async def get_memories(session_id: str):
    if db_conn:
        mems = await db_conn.get_memories(session_id)
        return {"memories": mems}
    return {"memories": []}


# ── Memories ───────────────────────────────────────────

@app.post("/api/memories/save")
async def save_memory(
    session_id: str,
    content: str,
    memory_type: str = "fact",
    importance: int = 3,
    source_doc_id: str | None = None,
):
    if not db_conn:
        raise HTTPException(500, "DB not ready")
    mem_id = await db_conn.store_memory(session_id, content, memory_type, importance, source_doc_id)
    return {"memory_id": mem_id}


@app.post("/api/memories/search")
async def search_memories(query: str, limit: int = 10):
    results = await db_conn.memory_similarity_search(query, limit)
    return {"memories": results}


# ── Ollama / Models ──────────────────────────────────

@app.get("/api/ollama/models")
async def check_ollama():
    try:
        async with httpx.AsyncClient(timeout=10) as client:
            resp = await client.get(f"{settings.ollama_url}/api/tags")
            return {"models": resp.json().get("models", [])}
    except Exception as e:
        return {"error": str(e), "ollama_url": settings.ollama_url}


@app.post("/api/ollama/chat")
async def ollama_chat(body: dict):
    model = body.get("model", settings.gpt_oss_model)
    messages = body.get("messages", [])
    
    async with httpx.AsyncClient(timeout=600) as client:
        resp = await client.post(
            f"{settings.ollama_url}/api/chat",
            json={"model": model, "messages": messages, "stream": False},
        )
        resp.raise_for_status()
        return resp.json()


# ── Agents ────────────────────────────────────────────

@app.get("/api/agents/skills")
async def get_skills():
    return {"skills": [{"name": n, "description": s.description} for n, s in SKILLS.items()]}


# ── Pipeline ────────────────────────────────────

class PipelineRequest(BaseModel):
    query: str
    doc_ids: list[str] | None = None
    output_mode: str | None = None


@app.post("/api/research/pipeline")
async def run_pipeline(req: PipelineRequest):
    """Run the full document-triage → evidence-extraction → research-synthesis pipeline."""
    if not db_conn:
        raise HTTPException(500, "DB not ready")
    
    pipeline = ResearchPipeline()
    results = await pipeline.run(
        query=req.query,
        session_id=req.doc_ids[0] if req.doc_ids else "temp",
        db=db_conn,
        doc_ids=req.doc_ids,
        output_mode=req.output_mode,
    )
    return results


@app.get("/api/research/pipeline/{session_id}/stages")
async def get_pipeline_stages(session_id: str):
    if db_conn:
        stages = await db_conn.get_pipeline_stages(session_id)
        return {"stages": stages}
    return {"stages": []}


@app.post("/api/research/pipeline/{session_id}/render")
async def render_pipeline_results(session_id: str):
    """Return rendered pipeline sections for display."""
    if not db_conn:
        raise HTTPException(500, "DB not ready")
    stages = await db_conn.get_pipeline_stages(session_id)
    pipeline = ResearchPipeline()
    sections = await pipeline.render_pipeline_results({s["stage"]: s["output"] for s in stages})
    return {"sections": sections, "session_id": session_id}


# ── Docs ──────────────────────────────────────────────

@app.get("/docs", include_in_schema=False)
async def docs_redirect():
    return FileResponse("/api/docs")