771ac9c96c
The README feature list and the overview stopped at the analysis capability. Both examples and the app backend composed agents without the capabilities the documentation recommends alongside an evidence capability. custom_agent.py ran each input as an independent agent run, so it needed a state dict and a carried history before compaction could mean anything there: without state the evidence record is empty, and earlier evidence would reduce to receipts retaining nothing.
294 lines
9.1 KiB
Python
294 lines
9.1 KiB
Python
import asyncio
|
|
import logging
|
|
import os
|
|
from contextlib import asynccontextmanager
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from ag_ui.core import EventType, StateSnapshotEvent
|
|
from dotenv import find_dotenv, load_dotenv
|
|
from pydantic_ai import Agent
|
|
from pydantic_ai.ui import SSE_CONTENT_TYPE
|
|
from pydantic_ai.ui.ag_ui import AGUIAdapter
|
|
from starlette.applications import Starlette
|
|
from starlette.middleware import Middleware
|
|
from starlette.middleware.cors import CORSMiddleware
|
|
from starlette.requests import Request
|
|
from starlette.responses import JSONResponse, Response, StreamingResponse
|
|
from starlette.routing import Route
|
|
|
|
from haiku.rag.capabilities.compaction import (
|
|
create_capability as create_compaction,
|
|
)
|
|
from haiku.rag.capabilities.policy import (
|
|
create_capability as create_citation_policy,
|
|
)
|
|
from haiku.rag.capabilities.rag import AGENT_PREAMBLE, RAGState, create_capability
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.config import load_yaml_config
|
|
from haiku.rag.config.models import AppConfig
|
|
from haiku.rag.telemetry import configure as configure_telemetry
|
|
from haiku.rag.utils import get_model
|
|
|
|
load_dotenv(find_dotenv(usecwd=True))
|
|
|
|
configure_telemetry(service_name="haiku-rag-app")
|
|
|
|
logging.basicConfig(
|
|
level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Load config
|
|
config_path = Path("/app/haiku.rag.yaml")
|
|
if config_path.exists():
|
|
yaml_data = load_yaml_config(config_path)
|
|
Config = AppConfig.model_validate(yaml_data)
|
|
else:
|
|
Config = AppConfig()
|
|
|
|
# Get DB path from environment
|
|
db_path_str = os.getenv("DB_PATH", "haiku_rag.lancedb")
|
|
db_path = Path(db_path_str)
|
|
|
|
logger.info(f"Database path: {db_path}")
|
|
logger.info(f"QA Provider: {Config.qa.model.provider}, Model: {Config.qa.model.name}")
|
|
|
|
# Only HaikuRAG client is a singleton (expensive to create)
|
|
_client: HaikuRAG | None = None
|
|
_client_lock = asyncio.Lock()
|
|
|
|
|
|
async def get_client() -> HaikuRAG:
|
|
"""Get or create the cached client.
|
|
|
|
Guarded by a lock because the first request after startup can race with
|
|
itself: two concurrent callers would both pass the None check, each build
|
|
and enter a HaikuRAG, and the loser would leak its LanceDB connection.
|
|
"""
|
|
global _client
|
|
if _client is None:
|
|
async with _client_lock:
|
|
if _client is None:
|
|
client = HaikuRAG(db_path=db_path, config=Config, create=True)
|
|
await client.__aenter__()
|
|
_client = client
|
|
return _client
|
|
|
|
|
|
@dataclass
|
|
class AppDeps:
|
|
state: dict[str, Any] = field(default_factory=dict)
|
|
|
|
|
|
capability = create_capability(db_path=db_path, config=Config, defer_loading=False)
|
|
|
|
agent = Agent(
|
|
get_model(Config.qa.model, Config),
|
|
instructions=AGENT_PREAMBLE,
|
|
# Conversations here are multi-turn, so earlier questions are reduced to the
|
|
# evidence they cited rather than carried whole, and every answer declares
|
|
# what grounds it so the UI can show citations for all of them.
|
|
capabilities=[capability, create_compaction(), create_citation_policy()],
|
|
deps_type=AppDeps,
|
|
)
|
|
|
|
|
|
async def stream_chat(request: Request) -> Response:
|
|
"""Chat streaming endpoint with AG-UI protocol."""
|
|
body = await request.body()
|
|
accept = request.headers.get("accept", SSE_CONTENT_TYPE)
|
|
run_input = AGUIAdapter.build_run_input(body)
|
|
|
|
adapter = AGUIAdapter(agent=agent, run_input=run_input, accept=accept)
|
|
|
|
incoming_state = run_input.state if isinstance(run_input.state, dict) else {}
|
|
incoming_state.setdefault("rag", RAGState().model_dump(mode="json"))
|
|
deps = AppDeps(state=incoming_state)
|
|
|
|
async def event_stream():
|
|
async def with_final_state():
|
|
async for event in adapter.run_stream(deps=deps):
|
|
if getattr(event, "type", None) == EventType.RUN_FINISHED:
|
|
yield StateSnapshotEvent(
|
|
type=EventType.STATE_SNAPSHOT,
|
|
snapshot=deps.state,
|
|
)
|
|
yield event
|
|
|
|
async for chunk in adapter.encode_stream(with_final_state()):
|
|
yield chunk
|
|
|
|
return StreamingResponse(
|
|
event_stream(),
|
|
media_type=accept,
|
|
headers={
|
|
"Cache-Control": "no-cache",
|
|
"Connection": "keep-alive",
|
|
"X-Accel-Buffering": "no",
|
|
},
|
|
)
|
|
|
|
|
|
async def health_check(_: Request) -> JSONResponse:
|
|
"""Health check endpoint."""
|
|
return JSONResponse(
|
|
{
|
|
"status": "healthy",
|
|
"qa_provider": Config.qa.model.provider,
|
|
"qa_model": Config.qa.model.name,
|
|
"db_path": str(db_path),
|
|
"db_exists": db_path.exists(),
|
|
}
|
|
)
|
|
|
|
|
|
async def list_documents(_: Request) -> JSONResponse:
|
|
"""List all documents in the database."""
|
|
if not db_path.exists():
|
|
return JSONResponse({"documents": [], "error": "Database not found"})
|
|
|
|
client = await get_client()
|
|
docs = await client.document_repository.list_all()
|
|
return JSONResponse(
|
|
{
|
|
"documents": [
|
|
{"id": doc.id, "title": doc.title, "uri": doc.uri} for doc in docs
|
|
]
|
|
}
|
|
)
|
|
|
|
|
|
async def db_info(_: Request) -> JSONResponse:
|
|
"""Get database info and statistics."""
|
|
if not db_path.exists():
|
|
return JSONResponse(
|
|
{
|
|
"exists": False,
|
|
"path": str(db_path),
|
|
"documents": 0,
|
|
"chunks": 0,
|
|
}
|
|
)
|
|
|
|
from haiku.rag.store.engine import get_database_stats
|
|
|
|
client = await get_client()
|
|
stats = await get_database_stats(client.store.db)
|
|
|
|
return JSONResponse(
|
|
{
|
|
"exists": True,
|
|
"path": str(db_path),
|
|
"documents": stats["documents"].get("num_rows", 0),
|
|
"chunks": stats["chunks"].get("num_rows", 0),
|
|
"documents_bytes": stats["documents"].get("total_bytes", 0),
|
|
"chunks_bytes": stats["chunks"].get("total_bytes", 0),
|
|
"has_vector_index": stats["chunks"].get("has_vector_index", False),
|
|
}
|
|
)
|
|
|
|
|
|
async def visualize_chunk(request: Request) -> JSONResponse:
|
|
"""Return visual grounding images for one or more chunks as base64.
|
|
|
|
The path param accepts comma-separated chunk ids (a merged citation's
|
|
constituent chunks). The optional ``refs`` query param is a JSON-encoded
|
|
list of the citation's ``doc_item_refs`` — the exact items the model saw —
|
|
so the highlight matches the cited content instead of re-expanding.
|
|
"""
|
|
import base64
|
|
import json
|
|
from io import BytesIO
|
|
|
|
chunk_id = request.path_params["chunk_id"]
|
|
|
|
refs: list[str] | None = None
|
|
refs_param = request.query_params.get("refs")
|
|
if refs_param:
|
|
try:
|
|
parsed = json.loads(refs_param)
|
|
except ValueError:
|
|
parsed = None
|
|
if isinstance(parsed, list):
|
|
refs = [str(x) for x in parsed]
|
|
|
|
if not db_path.exists():
|
|
return JSONResponse({"error": "Database not found"}, status_code=404)
|
|
|
|
client = await get_client()
|
|
|
|
chunks = []
|
|
for cid in chunk_id.split(","):
|
|
chunk = await client.chunk_repository.get_by_id(cid)
|
|
if chunk:
|
|
chunks.append(chunk)
|
|
if not chunks:
|
|
return JSONResponse({"error": "Chunk not found"}, status_code=404)
|
|
|
|
images = await client.visualize_chunk(chunks, refs)
|
|
if not images:
|
|
return JSONResponse({"images": [], "message": "No visual grounding available"})
|
|
|
|
base64_images = []
|
|
for img in images:
|
|
buffer = BytesIO()
|
|
img.save(buffer, format="PNG")
|
|
buffer.seek(0)
|
|
base64_images.append(base64.b64encode(buffer.read()).decode("utf-8"))
|
|
|
|
return JSONResponse(
|
|
{
|
|
"images": base64_images,
|
|
"chunk_id": chunk_id,
|
|
"document_uri": chunks[0].document_uri,
|
|
}
|
|
)
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(_app: Starlette):
|
|
"""Shut down the cached HaikuRAG client cleanly on app exit.
|
|
|
|
Awaits any in-flight background vacuum tasks and closes the LanceDB
|
|
connection. Without this, vacuum tasks are cancelled abruptly and the
|
|
connection is never closed on process shutdown.
|
|
"""
|
|
yield
|
|
global _client
|
|
if _client is not None:
|
|
await _client.__aexit__(None, None, None)
|
|
_client = None
|
|
|
|
|
|
# Create Starlette app
|
|
app = Starlette(
|
|
routes=[
|
|
Route("/v1/chat/stream", stream_chat, methods=["POST"]),
|
|
Route("/api/documents", list_documents, methods=["GET"]),
|
|
Route("/api/info", db_info, methods=["GET"]),
|
|
Route("/api/visualize/{chunk_id}", visualize_chunk, methods=["GET"]),
|
|
Route("/health", health_check, methods=["GET"]),
|
|
],
|
|
middleware=[
|
|
Middleware(
|
|
CORSMiddleware, # type: ignore[invalid-argument-type]
|
|
allow_origins=["http://localhost:3000", "http://frontend:3000"],
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
],
|
|
lifespan=lifespan,
|
|
)
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
|
|
uvicorn.run(
|
|
"main:app",
|
|
host="0.0.0.0",
|
|
port=8000,
|
|
reload=True,
|
|
)
|