995d081aa1
visualize_chunk re-expanded chunks from scratch to recover their refs, which could not faithfully reproduce the original merge, scores, and clip — so a visualization could highlight different pages than the citation covered. Carry the cited items on Citation.doc_item_refs and resolve bounding boxes from them directly; re-expansion remains only as the fallback for callers with no stored context (CLI, inspector). Chat, inspector, the app endpoint, and the frontend pass the refs through.
306 lines
9.4 KiB
Python
306 lines
9.4 KiB
Python
import asyncio
|
|
import logging
|
|
import os
|
|
from contextlib import asynccontextmanager
|
|
from pathlib import Path
|
|
|
|
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.client import HaikuRAG
|
|
from haiku.rag.config import load_yaml_config
|
|
from haiku.rag.config.models import AppConfig
|
|
from haiku.rag.skills.rag import create_skill, get_agent_preamble
|
|
from haiku.rag.utils import get_model
|
|
from haiku.skills import (
|
|
SkillDeps,
|
|
SkillToolset,
|
|
run_agui_stream,
|
|
)
|
|
from haiku.skills.prompts import build_system_prompt
|
|
|
|
load_dotenv(find_dotenv(usecwd=True))
|
|
|
|
# Configure logfire (only sends data if LOGFIRE_TOKEN is present)
|
|
try:
|
|
import logfire
|
|
|
|
logfire.configure(send_to_logfire="if-token-present", console=False)
|
|
logfire.instrument_pydantic_ai()
|
|
except Exception:
|
|
pass
|
|
|
|
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
|
|
|
|
|
|
# Create skill, toolset, and agent
|
|
skill = create_skill(db_path=db_path, config=Config)
|
|
toolset = SkillToolset(skills=[skill])
|
|
|
|
agent = Agent(
|
|
get_model(Config.qa.model, Config),
|
|
instructions=build_system_prompt(
|
|
toolset.skill_catalog, preamble=get_agent_preamble(Config)
|
|
),
|
|
toolsets=[toolset],
|
|
deps_type=SkillDeps,
|
|
)
|
|
|
|
|
|
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 {}
|
|
|
|
async def event_stream():
|
|
async with run_agui_stream(
|
|
adapter, toolset=toolset, deps=SkillDeps(state=incoming_state)
|
|
) as stream:
|
|
# Emit a STATE_SNAPSHOT after RUN_STARTED so the client holds every
|
|
# namespace object before any STATE_DELTA patches into it. Without it,
|
|
# the first `add /rag/<field>/...` fails against a missing parent.
|
|
if incoming_state:
|
|
toolset.restore_state_snapshot(incoming_state)
|
|
snapshot = StateSnapshotEvent(
|
|
type=EventType.STATE_SNAPSHOT,
|
|
snapshot=toolset.build_state_snapshot(),
|
|
)
|
|
|
|
async def with_state_snapshot():
|
|
emitted = False
|
|
async for event in stream:
|
|
yield event
|
|
if not emitted and getattr(event, "type", None) == (
|
|
EventType.RUN_STARTED
|
|
):
|
|
yield snapshot
|
|
emitted = True
|
|
|
|
async for chunk in adapter.encode_stream(with_state_snapshot()):
|
|
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,
|
|
)
|