haiku.rag/app/backend/main.py

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.compaction import (
    create_capability as create_compaction,
)
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,
    # Conversations here are multi-turn, so earlier questions are reduced to the
    # evidence they cited rather than carried whole.
    capabilities=[capability, create_compaction()],
    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,
    )