forked from retoor/devplacepy
feat: add deepsearch research system with CLI prune/clear and database schema
Implement a multi-agent deep web research subsystem including CLI commands for pruning expired jobs and clearing all artifacts, database tables for sessions/messages/URL cache with indexes, config paths for chroma storage, and internal embed URL for vector operations.
This commit is contained in:
@@ -0,0 +1,430 @@
|
||||
# retoor <retoor@molodetz.nl>
|
||||
|
||||
import json
|
||||
import logging
|
||||
from typing import Annotated
|
||||
|
||||
from fastapi import APIRouter, Form, Request, WebSocket, WebSocketDisconnect
|
||||
from fastapi.responses import HTMLResponse, JSONResponse, Response
|
||||
|
||||
from devplacepy import database
|
||||
from devplacepy.config import DEEPSEARCH_DIR
|
||||
from devplacepy.models import DeepsearchChatForm, DeepsearchRunForm
|
||||
from devplacepy.responses import respond
|
||||
from devplacepy.schemas import DeepsearchJobOut, DeepsearchSessionOut
|
||||
from devplacepy.seo import base_seo_context, site_url, web_application_schema, website_schema
|
||||
from devplacepy.services.deepsearch.chat import DeepsearchChat
|
||||
from devplacepy.services.deepsearch.export import to_json, to_markdown, to_pdf
|
||||
from devplacepy.services.jobs import queue
|
||||
from devplacepy.services.jobs.deepsearch.progress import hub
|
||||
from devplacepy.services.jobs.deepsearch.service import DeepsearchService
|
||||
from devplacepy.services.manager import service_manager
|
||||
from devplacepy.templating import templates
|
||||
from devplacepy.utils import generate_uid, get_current_user, is_admin, not_found
|
||||
|
||||
from ._shared import owner_for
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
router = APIRouter()
|
||||
|
||||
ACTIVE_STATES = (queue.PENDING, queue.RUNNING)
|
||||
TERMINAL_STATES = (queue.DONE, queue.FAILED)
|
||||
CONTROL_STATES = ("running", "paused", "cancelled")
|
||||
TOUCH_EXTEND_SECONDS = 3600
|
||||
_service = DeepsearchService()
|
||||
|
||||
|
||||
def _collection_for(uid: str) -> str:
|
||||
return _service.collection_name(uid)
|
||||
|
||||
|
||||
def _control_path(uid: str):
|
||||
return DEEPSEARCH_DIR / uid / "control.json"
|
||||
|
||||
|
||||
def _set_control(uid: str, state: str) -> None:
|
||||
path = _control_path(uid)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(json.dumps({"state": state}), encoding="utf-8")
|
||||
|
||||
|
||||
def _owns(request: Request, session: dict) -> bool:
|
||||
owner_kind, owner_id = owner_for(request)
|
||||
return (
|
||||
session.get("owner_kind") == owner_kind
|
||||
and session.get("owner_id") == owner_id
|
||||
)
|
||||
|
||||
|
||||
def _job_payload(job: dict) -> dict:
|
||||
result = job.get("result", {})
|
||||
payload = job.get("payload", {})
|
||||
done = job.get("status") == queue.DONE
|
||||
uid = job.get("uid", "")
|
||||
return {
|
||||
"uid": uid,
|
||||
"kind": job.get("kind", ""),
|
||||
"status": job.get("status", ""),
|
||||
"query": payload.get("query"),
|
||||
"depth": int(payload.get("depth") or 0),
|
||||
"max_pages": int(payload.get("max_pages") or 0),
|
||||
"ws_url": f"/tools/deepsearch/{uid}/ws",
|
||||
"chat_ws_url": f"/tools/deepsearch/{uid}/chat" if done else None,
|
||||
"session_url": f"/tools/deepsearch/{uid}/session" if done else None,
|
||||
"score": result.get("score") if done else None,
|
||||
"confidence": result.get("confidence") if done else None,
|
||||
"source_diversity": result.get("source_diversity") if done else None,
|
||||
"page_count": int(result.get("page_count") or 0),
|
||||
"chunk_count": int(result.get("chunk_count") or 0),
|
||||
"error": job.get("error") or None,
|
||||
"created_at": job.get("created_at"),
|
||||
"completed_at": job.get("completed_at") or None,
|
||||
}
|
||||
|
||||
|
||||
@router.get("", response_class=HTMLResponse)
|
||||
async def deepsearch_page(request: Request):
|
||||
user = get_current_user(request)
|
||||
base = site_url(request)
|
||||
description = (
|
||||
"Run a multi-agent deep web research job: it plans queries, crawls and indexes "
|
||||
"sources, then synthesises a cited report with confidence scoring and gap "
|
||||
"analysis, plus a grounded chat over the results."
|
||||
)
|
||||
seo_ctx = base_seo_context(
|
||||
request,
|
||||
title="DeepSearch",
|
||||
description=description,
|
||||
breadcrumbs=[
|
||||
{"name": "Home", "url": "/feed"},
|
||||
{"name": "Tools", "url": "/tools"},
|
||||
{"name": "DeepSearch", "url": "/tools/deepsearch"},
|
||||
],
|
||||
schemas=[
|
||||
website_schema(base),
|
||||
web_application_schema("DeepSearch", description, "/tools/deepsearch", base),
|
||||
],
|
||||
)
|
||||
return templates.TemplateResponse(
|
||||
request,
|
||||
"tools/deepsearch.html",
|
||||
{**seo_ctx, "request": request, "user": user},
|
||||
)
|
||||
|
||||
|
||||
@router.post("/run")
|
||||
async def deepsearch_run(request: Request, data: Annotated[DeepsearchRunForm, Form()]):
|
||||
owner_kind, owner_id = owner_for(request)
|
||||
from devplacepy.services.audit import record as audit
|
||||
|
||||
active = [
|
||||
job
|
||||
for job in queue.list_jobs(kind="deepsearch", owner=(owner_kind, owner_id))
|
||||
if job.get("status") in ACTIVE_STATES
|
||||
]
|
||||
if active:
|
||||
audit.record(
|
||||
request,
|
||||
"deepsearch.run.request",
|
||||
result="denied",
|
||||
summary=f"DeepSearch denied for {data.query}: already running",
|
||||
metadata={"query": data.query, "reason": "active_job", "uid": active[0]["uid"]},
|
||||
links=[audit.job(active[0]["uid"])],
|
||||
)
|
||||
return JSONResponse(
|
||||
{
|
||||
"error": {
|
||||
"status": 429,
|
||||
"message": "You already have a DeepSearch running. Wait for it to finish.",
|
||||
"uid": active[0]["uid"],
|
||||
}
|
||||
},
|
||||
status_code=429,
|
||||
)
|
||||
user = get_current_user(request)
|
||||
api_key = (
|
||||
user.get("api_key") if user else ""
|
||||
) or database.internal_gateway_key()
|
||||
uid = generate_uid()
|
||||
payload = {
|
||||
"query": data.query,
|
||||
"depth": data.depth,
|
||||
"max_pages": data.max_pages,
|
||||
"api_key": api_key,
|
||||
"collection": _collection_for(uid),
|
||||
}
|
||||
database.create_deepsearch_session(
|
||||
uid, owner_kind, owner_id, data.query, data.depth, data.max_pages, _collection_for(uid)
|
||||
)
|
||||
_enqueue(uid, payload, owner_kind, owner_id, data.query)
|
||||
audit.record(
|
||||
request,
|
||||
"deepsearch.run.request",
|
||||
summary=f"requested DeepSearch for {data.query}",
|
||||
metadata={"query": data.query, "depth": data.depth, "max_pages": data.max_pages},
|
||||
links=[audit.job(uid)],
|
||||
)
|
||||
return JSONResponse(
|
||||
{
|
||||
"uid": uid,
|
||||
"status_url": f"/tools/deepsearch/{uid}",
|
||||
"ws_url": f"/tools/deepsearch/{uid}/ws",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _enqueue(uid: str, payload: dict, owner_kind: str, owner_id: str, query: str) -> None:
|
||||
from datetime import datetime, timezone
|
||||
|
||||
now = datetime.now(timezone.utc).isoformat()
|
||||
database.get_table("jobs").insert(
|
||||
{
|
||||
"uid": uid,
|
||||
"kind": "deepsearch",
|
||||
"status": queue.PENDING,
|
||||
"owner_kind": owner_kind,
|
||||
"owner_id": owner_id,
|
||||
"preferred_name": f"DeepSearch: {query}"[:64],
|
||||
"payload": json.dumps(payload),
|
||||
"result": "",
|
||||
"error": "",
|
||||
"retry_count": 0,
|
||||
"created_at": now,
|
||||
"started_at": "",
|
||||
"completed_at": "",
|
||||
"updated_at": now,
|
||||
"duration_ms": 0,
|
||||
"last_accessed_at": "",
|
||||
"expires_at": "",
|
||||
"bytes_in": 0,
|
||||
"bytes_out": 0,
|
||||
"item_count": 0,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@router.get("/{uid}")
|
||||
async def deepsearch_status(request: Request, uid: str):
|
||||
job = queue.get_job(uid)
|
||||
if not job or job.get("kind") != "deepsearch":
|
||||
raise not_found("Research not found")
|
||||
return JSONResponse(
|
||||
DeepsearchJobOut.model_validate(_job_payload(job)).model_dump(mode="json")
|
||||
)
|
||||
|
||||
|
||||
def _report_for(job: dict) -> dict:
|
||||
if job.get("status") != queue.DONE:
|
||||
return {}
|
||||
return job.get("result", {}).get("report", {})
|
||||
|
||||
|
||||
def _session_context(request: Request, uid: str, job: dict, session: dict) -> dict:
|
||||
report = _report_for(job)
|
||||
user = get_current_user(request)
|
||||
viewer_is_admin = is_admin(user)
|
||||
done = job.get("status") == queue.DONE
|
||||
return {
|
||||
"uid": uid,
|
||||
"status": job.get("status", ""),
|
||||
"query": report.get("query") or session.get("query"),
|
||||
"depth": int(session.get("depth") or 0),
|
||||
"max_pages": int(session.get("max_pages") or 0),
|
||||
"score": report.get("score"),
|
||||
"confidence": report.get("confidence"),
|
||||
"source_diversity": report.get("source_diversity"),
|
||||
"page_count": report.get("page_count", 0),
|
||||
"chunk_count": report.get("chunk_count", 0),
|
||||
"summary": report.get("summary", ""),
|
||||
"sources": report.get("sources", []),
|
||||
"findings": report.get("findings", []),
|
||||
"gaps": report.get("gaps", []),
|
||||
"timeline": report.get("timeline", []),
|
||||
"chat_ws_url": f"/tools/deepsearch/{uid}/chat" if done else None,
|
||||
"export_md_url": f"/tools/deepsearch/{uid}/export.md" if done else None,
|
||||
"export_json_url": f"/tools/deepsearch/{uid}/export.json" if done else None,
|
||||
"export_pdf_url": f"/tools/deepsearch/{uid}/export.pdf" if done else None,
|
||||
"viewer_is_admin": viewer_is_admin,
|
||||
"viewer_owns": _owns(request, session),
|
||||
"created_at": session.get("created_at"),
|
||||
"completed_at": session.get("completed_at") or None,
|
||||
"request": request,
|
||||
"user": user,
|
||||
"meta_robots": "noindex,nofollow",
|
||||
}
|
||||
|
||||
|
||||
@router.get("/{uid}/session")
|
||||
async def deepsearch_session(request: Request, uid: str):
|
||||
job = queue.get_job(uid)
|
||||
session = database.get_deepsearch_session(uid)
|
||||
if not job or job.get("kind") != "deepsearch" or not session:
|
||||
raise not_found("Research not found")
|
||||
queue.touch_job(uid, TOUCH_EXTEND_SECONDS)
|
||||
context = _session_context(request, uid, job, session)
|
||||
return respond(request, "tools/deepsearch_session.html", context, model=DeepsearchSessionOut)
|
||||
|
||||
|
||||
@router.post("/{uid}/pause")
|
||||
async def deepsearch_pause(request: Request, uid: str):
|
||||
return _control(request, uid, "paused")
|
||||
|
||||
|
||||
@router.post("/{uid}/resume")
|
||||
async def deepsearch_resume(request: Request, uid: str):
|
||||
return _control(request, uid, "running")
|
||||
|
||||
|
||||
@router.post("/{uid}/cancel")
|
||||
async def deepsearch_cancel(request: Request, uid: str):
|
||||
return _control(request, uid, "cancelled")
|
||||
|
||||
|
||||
def _control(request: Request, uid: str, state: str):
|
||||
session = database.get_deepsearch_session(uid)
|
||||
job = queue.get_job(uid)
|
||||
if not job or job.get("kind") != "deepsearch" or not session:
|
||||
raise not_found("Research not found")
|
||||
if not _owns(request, session):
|
||||
return JSONResponse(
|
||||
{"error": {"status": 403, "message": "Not your research session."}},
|
||||
status_code=403,
|
||||
)
|
||||
_set_control(uid, state)
|
||||
return JSONResponse({"ok": True, "state": state})
|
||||
|
||||
|
||||
def _export_report(uid: str) -> dict | None:
|
||||
job = queue.get_job(uid)
|
||||
if not job or job.get("kind") != "deepsearch" or job.get("status") != queue.DONE:
|
||||
return None
|
||||
queue.touch_job(uid, TOUCH_EXTEND_SECONDS)
|
||||
return job.get("result", {}).get("report", {})
|
||||
|
||||
|
||||
@router.get("/{uid}/export.md")
|
||||
async def deepsearch_export_md(request: Request, uid: str):
|
||||
report = _export_report(uid)
|
||||
if report is None:
|
||||
raise not_found("Report not ready")
|
||||
return Response(content=to_markdown(report), media_type="text/markdown")
|
||||
|
||||
|
||||
@router.get("/{uid}/export.json")
|
||||
async def deepsearch_export_json(request: Request, uid: str):
|
||||
report = _export_report(uid)
|
||||
if report is None:
|
||||
raise not_found("Report not ready")
|
||||
return Response(content=to_json(report), media_type="application/json")
|
||||
|
||||
|
||||
@router.get("/{uid}/export.pdf")
|
||||
async def deepsearch_export_pdf(request: Request, uid: str):
|
||||
report = _export_report(uid)
|
||||
if report is None:
|
||||
raise not_found("Report not ready")
|
||||
return Response(content=to_pdf(report), media_type="application/pdf")
|
||||
|
||||
|
||||
@router.websocket("/{uid}/ws")
|
||||
async def deepsearch_ws(websocket: WebSocket, uid: str):
|
||||
await websocket.accept()
|
||||
if not service_manager.owns_lock():
|
||||
await websocket.close(code=4013)
|
||||
return
|
||||
job = queue.get_job(uid)
|
||||
if not job or job.get("kind") != "deepsearch":
|
||||
await websocket.close(code=1008)
|
||||
return
|
||||
listener = hub.subscribe(uid)
|
||||
try:
|
||||
for frame in hub.snapshot(uid):
|
||||
await websocket.send_json(frame)
|
||||
current = queue.get_job(uid)
|
||||
if current and current.get("status") in TERMINAL_STATES:
|
||||
done = current.get("status") == queue.DONE
|
||||
await websocket.send_json(
|
||||
{
|
||||
"type": "done" if done else "failed",
|
||||
"status": current.get("status"),
|
||||
"session_url": f"/tools/deepsearch/{uid}/session",
|
||||
"error": current.get("error") or None,
|
||||
}
|
||||
)
|
||||
return
|
||||
while True:
|
||||
frame = await listener.get()
|
||||
await websocket.send_json(frame)
|
||||
if frame.get("type") in ("done", "failed"):
|
||||
break
|
||||
except WebSocketDisconnect:
|
||||
pass
|
||||
except Exception:
|
||||
logger.exception("DeepSearch websocket loop failed for %s", uid)
|
||||
finally:
|
||||
hub.unsubscribe(uid, listener)
|
||||
|
||||
|
||||
@router.websocket("/{uid}/chat")
|
||||
async def deepsearch_chat_ws(websocket: WebSocket, uid: str):
|
||||
await websocket.accept()
|
||||
if not service_manager.owns_lock():
|
||||
await websocket.close(code=4013)
|
||||
return
|
||||
job = queue.get_job(uid)
|
||||
session = database.get_deepsearch_session(uid)
|
||||
if not job or job.get("kind") != "deepsearch" or job.get("status") != queue.DONE or not session:
|
||||
await websocket.close(code=1008)
|
||||
return
|
||||
user = get_current_user(websocket)
|
||||
api_key = (
|
||||
user.get("api_key") if user else ""
|
||||
) or database.internal_gateway_key()
|
||||
chat = DeepsearchChat(_collection_for(uid), api_key)
|
||||
history: list[dict] = []
|
||||
for row in database.get_deepsearch_messages(uid):
|
||||
history.append({"role": row.get("role"), "content": row.get("content")})
|
||||
await websocket.send_json({"type": "ready", "history": history})
|
||||
from devplacepy.services.audit import record as audit
|
||||
|
||||
try:
|
||||
while True:
|
||||
raw = await websocket.receive_text()
|
||||
try:
|
||||
frame = json.loads(raw)
|
||||
except (ValueError, TypeError):
|
||||
continue
|
||||
if frame.get("type") != "input":
|
||||
continue
|
||||
try:
|
||||
form = DeepsearchChatForm.model_validate({"message": frame.get("text", "")})
|
||||
except Exception:
|
||||
await websocket.send_json({"type": "error", "text": "Empty message."})
|
||||
continue
|
||||
question = form.message
|
||||
await websocket.send_json({"type": "status", "text": "Searching the research collection"})
|
||||
database.add_deepsearch_message(generate_uid(), uid, "user", question)
|
||||
history.append({"role": "user", "content": question})
|
||||
answer = await chat.answer(question, history)
|
||||
database.add_deepsearch_message(
|
||||
generate_uid(),
|
||||
uid,
|
||||
"assistant",
|
||||
answer.text,
|
||||
json.dumps(answer.citations),
|
||||
)
|
||||
history.append({"role": "assistant", "content": answer.text})
|
||||
await websocket.send_json(
|
||||
{"type": "reply", "text": answer.text, "citations": answer.citations}
|
||||
)
|
||||
audit.record(
|
||||
websocket,
|
||||
"deepsearch.chat",
|
||||
summary=f"DeepSearch chat on session {uid[:8]}",
|
||||
metadata={"session": uid},
|
||||
links=[audit.job(uid)],
|
||||
)
|
||||
except WebSocketDisconnect:
|
||||
pass
|
||||
except Exception:
|
||||
logger.exception("DeepSearch chat loop failed for %s", uid)
|
||||
Reference in New Issue
Block a user