285 lines
8.7 KiB
Python
285 lines
8.7 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.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,
|
|
capabilities=[capability],
|
|
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,
|
|
)
|