262 lines
7.6 KiB
Python
262 lines
7.6 KiB
Python
import asyncio
|
|
import logging
|
|
import os
|
|
from contextlib import asynccontextmanager
|
|
from pathlib import Path
|
|
|
|
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)
|
|
|
|
async def event_stream():
|
|
async with run_agui_stream(
|
|
adapter, toolset=toolset, deps=SkillDeps()
|
|
) as stream:
|
|
async for chunk in adapter.encode_stream(stream):
|
|
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 a chunk as base64."""
|
|
import base64
|
|
from io import BytesIO
|
|
|
|
chunk_id = request.path_params["chunk_id"]
|
|
|
|
if not db_path.exists():
|
|
return JSONResponse({"error": "Database not found"}, status_code=404)
|
|
|
|
client = await get_client()
|
|
|
|
chunk = await client.chunk_repository.get_by_id(chunk_id)
|
|
if not chunk:
|
|
return JSONResponse({"error": "Chunk not found"}, status_code=404)
|
|
|
|
images = await client.visualize_chunk(chunk)
|
|
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": chunk.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,
|
|
)
|