Basic backend for app

app/.env.example

@@ -0,0 +1,9 @@
+# API Keys (at least one required for LLM)
+ANTHROPIC_API_KEY=your-anthropic-key
+OPENAI_API_KEY=your-openai-key
+
+# Database path
+DB_PATH=/path/to/your/haiku.rag.lancedb
+
+# Optional: Ollama base URL (if using local models)
+OLLAMA_BASE_URL=http://localhost:11434

app/backend/agent.py

@@ -0,0 +1,124 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel
+from pydantic_ai import Agent, RunContext
+
+from haiku.rag.client import HaikuRAG
+from haiku.rag.config.models import AppConfig
+from haiku.rag.store.models import SearchResult
+from haiku.rag.utils import get_model
+
+if TYPE_CHECKING:
+    from haiku.rag.graph.agui.emitter import AGUIEmitter
+
+
+class ChatSessionState(BaseModel):
+    """State shared between frontend and agent via AG-UI."""
+
+    session_id: str = ""
+
+
+@dataclass
+class ChatDeps:
+    """Dependencies for chat agent."""
+
+    client: HaikuRAG
+    config: AppConfig
+    agui_emitter: "AGUIEmitter | None" = None
+    search_results: list[SearchResult] | None = None
+
+
+CHAT_SYSTEM_PROMPT = """You are a helpful research assistant powered by haiku.rag, a knowledge base system.
+
+You have access to a knowledge base of documents. Use your tools to search and answer questions.
+
+CRITICAL RULES:
+1. For greetings or casual chat: respond directly WITHOUT using any tools
+2. For substantive questions requiring information: use the search or ask tools
+3. NEVER make up information - always use tools to get facts from the knowledge base
+4. When citing sources, reference the chunk IDs from search results
+
+How to decide which tool to use:
+- "search" - When you need to find relevant documents or explore what's in the knowledge base
+- "ask" - When you have a specific question that needs a direct answer with citations
+
+Be friendly and conversational. When you use tools, summarize the key findings for the user."""
+
+
+def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
+    """Create the chat agent with search and ask tools."""
+    model = get_model(config.qa.model, config)
+
+    agent: Agent[ChatDeps, str] = Agent(
+        model,
+        deps_type=ChatDeps,
+        output_type=str,
+        instructions=CHAT_SYSTEM_PROMPT,
+    )
+
+    @agent.tool
+    async def search(
+        ctx: RunContext[ChatDeps],
+        query: str,
+        limit: int = 5,
+        document_filter: str | None = None,
+    ) -> str:
+        """Search the knowledge base for relevant documents.
+
+        Use this when you need to find documents or explore the knowledge base.
+        Returns relevant chunks with metadata.
+
+        Args:
+            query: The search query
+            limit: Maximum number of results (default 5)
+            document_filter: Optional SQL WHERE clause to filter documents (e.g. "id IN ('doc1', 'doc2')")
+        """
+        if ctx.deps.agui_emitter:
+            ctx.deps.agui_emitter.log(f"Searching: {query}")
+
+        results = await ctx.deps.client.search(
+            query, limit=limit, filter=document_filter
+        )
+        results = await ctx.deps.client.expand_context(results)
+
+        # Store for potential citation resolution
+        ctx.deps.search_results = results
+
+        if not results:
+            return "No results found for your query."
+
+        # Format results for the agent
+        parts = [r.format_for_agent() for r in results]
+        return "\n\n".join(parts)
+
+    @agent.tool
+    async def ask(
+        ctx: RunContext[ChatDeps],
+        question: str,
+        document_filter: str | None = None,
+    ) -> str:
+        """Answer a specific question using the knowledge base.
+
+        Use this for direct questions that need a focused answer with citations.
+
+        Args:
+            question: The question to answer
+            document_filter: Optional SQL WHERE clause to filter documents (e.g. "id IN ('doc1', 'doc2')")
+        """
+        if ctx.deps.agui_emitter:
+            ctx.deps.agui_emitter.log(f"Answering: {question}")
+
+        answer, citations = await ctx.deps.client.ask(question, filter=document_filter)
+
+        # Format answer with citations
+        if citations:
+            citation_list = "\n".join(
+                f"  [{i + 1}] {c.document_uri or c.document_title or 'Unknown'} (chunk: {c.chunk_id})"
+                for i, c in enumerate(citations)
+            )
+            return f"{answer}\n\nSources:\n{citation_list}"
+
+        return answer
+
+    return agent

app/backend/main.py

@@ -0,0 +1,205 @@
+import logging
+import os
+from pathlib import Path
+
+from agent import ChatDeps, ChatSessionState, create_chat_agent
+from anyio import create_memory_object_stream, create_task_group
+from anyio.streams.memory import MemoryObjectSendStream
+from dotenv import load_dotenv
+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, 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.graph.agui.emitter import AGUIEmitter
+from haiku.rag.graph.agui.server import RunAgentInput, format_sse_event
+
+load_dotenv()
+
+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) -> StreamingResponse:
+    """Chat streaming endpoint with AG-UI protocol."""
+    body = await request.json()
+    logger.info(f"Received request: {list(body.keys())}")
+    input_data = RunAgentInput(**body)
+
+    user_message = ""
+    if input_data.messages:
+        user_message = input_data.messages[-1].get("content", "")
+
+    send_stream, receive_stream = create_memory_object_stream[str]()
+
+    async def run_agent_with_streaming(
+        send_stream: MemoryObjectSendStream[str],
+    ) -> None:
+        """Execute agent and forward events to stream."""
+        async with send_stream:
+            try:
+                # Create emitter for streaming
+                emitter: AGUIEmitter = AGUIEmitter(
+                    thread_id=input_data.thread_id,
+                    run_id=input_data.run_id,
+                    use_deltas=True,
+                )
+
+                # Get client
+                effective_db_path = db_path
+                if input_data.config and input_data.config.get("db_path"):
+                    effective_db_path = Path(input_data.config["db_path"])
+                client = get_client(effective_db_path)
+
+                # Create deps
+                deps = ChatDeps(
+                    client=client,
+                    config=Config,
+                    agui_emitter=emitter,
+                )
+
+                # Start run with empty state
+                initial_state = ChatSessionState(
+                    session_id=input_data.thread_id or "",
+                )
+                emitter.start_run(initial_state=initial_state)
+
+                # Forward events
+                async def forward_events():
+                    async for event in emitter:
+                        event_type = event.get("type")
+                        logger.debug(f"AG-UI event: {event_type}")
+                        await send_stream.send(format_sse_event(event))
+
+                # Run agent and forward concurrently
+                async with create_task_group() as tg:
+                    tg.start_soon(forward_events)
+
+                    result = await chat_agent.run(user_message, deps=deps)
+                    emitter.log(result.output)
+                    emitter.finish_run(result.output)
+                    await emitter.close()
+
+            except Exception as e:
+                logger.exception("Error executing agent")
+                try:
+                    await send_stream.send(
+                        format_sse_event({"type": "RUN_ERROR", "message": str(e)})
+                    )
+                except Exception:
+                    pass
+
+    async def event_generator():
+        """Generate SSE events."""
+        async with create_task_group() as tg:
+            tg.start_soon(run_agent_with_streaming, send_stream)
+            async with receive_stream:
+                async for event_str in receive_stream:
+                    yield event_str
+
+    return StreamingResponse(
+        event_generator(),
+        media_type="text/event-stream",
+        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
+            ]
+        }
+    )
+
+
+# Create Starlette app
+app = Starlette(
+    routes=[
+        Route("/v1/chat/stream", stream_chat, methods=["POST"]),
+        Route("/api/documents", list_documents, methods=["GET"]),
+        Route("/health", health_check, methods=["GET"]),
+    ],
+    middleware=[
+        Middleware(
+            CORSMiddleware,
+            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,
+    )

app/backend/pyproject.toml

@@ -0,0 +1,25 @@
+[project]
+name = "haiku-rag-app"
+version = "0.1.0"
+description = "Conversational RAG application with haiku.rag"
+requires-python = ">=3.12"
+dependencies = [
+    "starlette>=0.50.0",
+    "uvicorn[standard]>=0.40.0",
+    "pydantic-ai-slim[ag-ui,anthropic,openai]>=1.39.0",
+    "python-dotenv>=1.2.1",
+    "haiku.rag-slim[agui]>=0.23.1",
+]
+
+[dependency-groups]
+dev = ["pyright>=1.1.407", "ruff>=0.14.10"]
+
+[tool.hatch.metadata]
+allow-direct-references = true
+
+[tool.hatch.build.targets.wheel]
+packages = ["."]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"