136 lines
3.8 KiB
Python
136 lines
3.8 KiB
Python
"""Main entry point for the haiku.rag AG-UI research assistant backend."""
|
|
|
|
import logging
|
|
import os
|
|
from contextlib import asynccontextmanager
|
|
from pathlib import Path
|
|
|
|
from agent import ResearchDeps, ResearchState, create_agent
|
|
from starlette.applications import Starlette
|
|
from starlette.middleware import Middleware
|
|
from starlette.middleware.cors import CORSMiddleware
|
|
from starlette.responses import JSONResponse
|
|
from starlette.routing import Mount, Route
|
|
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.config import Config
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Global client instance
|
|
client: HaikuRAG | None = None
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app):
|
|
"""Manage HaikuRAG client lifecycle."""
|
|
global client
|
|
|
|
# Get database path from environment or use default
|
|
db_path_str = os.getenv("DB_PATH", "haiku_rag.lancedb")
|
|
db_path = Path(db_path_str)
|
|
|
|
if not db_path.exists():
|
|
logger.error(
|
|
f"Database not found at {db_path}. Please initialize haiku.rag first."
|
|
)
|
|
logger.error("Run: haiku-rag add <path-to-documents>")
|
|
raise RuntimeError(f"Database not found: {db_path}")
|
|
|
|
logger.info(f"Initializing HaikuRAG client with database: {db_path}")
|
|
client = HaikuRAG(db_path)
|
|
|
|
logger.info("Research assistant backend ready")
|
|
logger.info(f"QA Provider: {Config.QA_PROVIDER}")
|
|
logger.info(f"QA Model: {Config.QA_MODEL}")
|
|
|
|
yield
|
|
|
|
# Cleanup
|
|
if client:
|
|
logger.info("Closing HaikuRAG client")
|
|
client.close()
|
|
|
|
|
|
# Create research agent instance
|
|
agent = create_agent()
|
|
|
|
|
|
async def health(request):
|
|
"""Health check endpoint."""
|
|
db_path_str = os.getenv("DB_PATH", "haiku_rag.lancedb")
|
|
return JSONResponse(
|
|
{
|
|
"status": "healthy",
|
|
"agent_model": str(agent.model),
|
|
"qa_provider": Config.QA_PROVIDER,
|
|
"qa_model": Config.QA_MODEL,
|
|
"ollama_base_url": Config.OLLAMA_BASE_URL,
|
|
"db_path": db_path_str,
|
|
"db_exists": Path(db_path_str).exists(),
|
|
}
|
|
)
|
|
|
|
|
|
# Create AG-UI app once with the client
|
|
# State will be managed per-session by AG-UI
|
|
ag_ui_app = None
|
|
|
|
|
|
def get_ag_ui_app():
|
|
"""Get or create AG-UI app."""
|
|
global ag_ui_app
|
|
if ag_ui_app is None and client is not None:
|
|
if client is None:
|
|
raise RuntimeError("Client not initialized")
|
|
|
|
# Create deps with shared client but new state per session
|
|
research_deps = ResearchDeps(client=client, state=ResearchState())
|
|
logger.info("Creating AG-UI app with initial state")
|
|
ag_ui_app = agent.to_ag_ui(deps=research_deps)
|
|
return ag_ui_app
|
|
|
|
|
|
# Initialize AG-UI app after client is ready in lifespan
|
|
async def agent_endpoint(scope, receive, send):
|
|
"""Proxy requests to AG-UI app."""
|
|
app = get_ag_ui_app()
|
|
if app is None:
|
|
response = JSONResponse({"error": "Client not initialized"}, status_code=503)
|
|
await response(scope, receive, send)
|
|
return
|
|
await app(scope, receive, send)
|
|
|
|
|
|
# Mount the AG-UI app at /agent and add health endpoint
|
|
app = Starlette(
|
|
routes=[
|
|
Route("/health", health),
|
|
Mount("/agent", agent_endpoint),
|
|
],
|
|
middleware=[
|
|
Middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["http://localhost:3000", "http://frontend:3000"],
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
],
|
|
lifespan=lifespan,
|
|
)
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
|
|
print("Starting haiku.rag research assistant backend...")
|
|
print(f"Agent model: {agent.model}")
|
|
print(f"QA provider: {Config.QA_PROVIDER}")
|
|
print(f"QA model: {Config.QA_MODEL}")
|
|
|
|
uvicorn.run(
|
|
"main:app",
|
|
host="0.0.0.0",
|
|
port=8000,
|
|
reload=True,
|
|
)
|