281 lines
9 KiB
Python
281 lines
9 KiB
Python
import logging
|
|
import os
|
|
from pathlib import Path
|
|
|
|
from dotenv import find_dotenv, load_dotenv
|
|
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.agents.chat import (
|
|
AGUI_STATE_KEY,
|
|
ChatDeps,
|
|
ChatSessionState,
|
|
QAResponse,
|
|
create_chat_agent,
|
|
)
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.config import load_yaml_config
|
|
from haiku.rag.config.models import AppConfig
|
|
|
|
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}")
|
|
|
|
# Create the chat agent
|
|
chat_agent = create_chat_agent(Config)
|
|
|
|
# Client cache for proper lifecycle
|
|
_client_cache: dict[str, HaikuRAG] = {}
|
|
|
|
|
|
def get_client(effective_db_path: Path) -> HaikuRAG:
|
|
"""Get or create cached client."""
|
|
path_key = str(effective_db_path)
|
|
if path_key not in _client_cache:
|
|
_client_cache[path_key] = HaikuRAG(
|
|
db_path=effective_db_path, config=Config, create=True
|
|
)
|
|
return _client_cache[path_key]
|
|
|
|
|
|
async def stream_chat(request: Request) -> Response:
|
|
"""Chat streaming endpoint with AG-UI protocol."""
|
|
body = await request.body()
|
|
|
|
# Parse request to build run_input
|
|
accept = request.headers.get("accept", SSE_CONTENT_TYPE)
|
|
run_input = AGUIAdapter.build_run_input(body)
|
|
|
|
# Restore session state from incoming AG-UI state
|
|
session_state = ChatSessionState(session_id="") # New session: empty session_id
|
|
state = getattr(run_input, "state", None)
|
|
if state and AGUI_STATE_KEY in state:
|
|
chat_state = state[AGUI_STATE_KEY]
|
|
if chat_state and chat_state.get("session_id"):
|
|
session_state = ChatSessionState(
|
|
session_id=chat_state["session_id"],
|
|
qa_history=[
|
|
QAResponse(**qa) for qa in chat_state.get("qa_history", [])
|
|
],
|
|
document_filter=chat_state.get("document_filter", []),
|
|
initial_context=chat_state.get("initial_context"),
|
|
citation_registry=chat_state.get("citation_registry", {}),
|
|
)
|
|
logger.info(
|
|
f"Incoming state: session={session_state.session_id[:8]}, "
|
|
f"qa_history={len(session_state.qa_history)}, "
|
|
f"citations={len(session_state.citation_registry)}"
|
|
)
|
|
else:
|
|
logger.info("Incoming state: new session")
|
|
else:
|
|
logger.info("Incoming state: new session")
|
|
|
|
deps = ChatDeps(
|
|
client=get_client(db_path),
|
|
config=Config,
|
|
session_state=session_state,
|
|
state_key=AGUI_STATE_KEY,
|
|
)
|
|
|
|
# Use AGUIAdapter for streaming
|
|
adapter = AGUIAdapter(agent=chat_agent, run_input=run_input, accept=accept)
|
|
event_stream = adapter.run_stream(deps=deps)
|
|
|
|
async def logged_event_stream():
|
|
async for event in event_stream:
|
|
event_type = getattr(event, "type", None)
|
|
if event_type and "state" in str(event_type).lower():
|
|
delta = getattr(event, "delta", None)
|
|
snapshot = getattr(event, "snapshot", None)
|
|
if delta is not None:
|
|
logger.info(f"Outgoing StateDeltaEvent: {len(delta)} ops")
|
|
for op in delta:
|
|
# Extract key from path like /haiku.rag.chat/qa_history/0
|
|
parts = op["path"].split("/")
|
|
key = "/".join(parts[2:]) if len(parts) > 2 else op["path"]
|
|
logger.info(f" {op['op']} {key}")
|
|
elif snapshot is not None:
|
|
chat_state = snapshot.get(AGUI_STATE_KEY, {})
|
|
sid = chat_state.get("session_id", "")[:8] if chat_state else ""
|
|
qa_len = len(chat_state.get("qa_history", [])) if chat_state else 0
|
|
reg_len = (
|
|
len(chat_state.get("citation_registry", {}))
|
|
if chat_state
|
|
else 0
|
|
)
|
|
logger.info(
|
|
f"Outgoing StateSnapshotEvent: session={sid}, "
|
|
f"qa={qa_len}, keys={reg_len}"
|
|
)
|
|
yield event
|
|
|
|
sse_event_stream = adapter.encode_stream(logged_event_stream())
|
|
|
|
return StreamingResponse(
|
|
sse_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",
|
|
"agent_model": str(chat_agent.model),
|
|
"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 = get_client(db_path)
|
|
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,
|
|
}
|
|
)
|
|
|
|
client = get_client(db_path)
|
|
stats = client.store.get_stats()
|
|
|
|
return JSONResponse(
|
|
{
|
|
"exists": True,
|
|
"path": str(db_path),
|
|
"documents": stats.get("documents", {}).get("num_rows", 0),
|
|
"chunks": stats.get("chunks", {}).get("num_rows", 0),
|
|
"documents_bytes": stats.get("documents", {}).get("total_bytes", 0),
|
|
"chunks_bytes": stats.get("chunks", {}).get("total_bytes", 0),
|
|
"has_vector_index": stats.get("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 = get_client(db_path)
|
|
|
|
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,
|
|
}
|
|
)
|
|
|
|
|
|
# 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=["*"],
|
|
)
|
|
],
|
|
)
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
|
|
uvicorn.run(
|
|
"main:app",
|
|
host="0.0.0.0",
|
|
port=8000,
|
|
reload=True,
|
|
)
|