378 lines
13 KiB
Python
378 lines
13 KiB
Python
"""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")
|