diff --git a/app/.env.example b/app/.env.example new file mode 100644 index 00000000..338dbd92 --- /dev/null +++ b/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 diff --git a/app/backend/agent.py b/app/backend/agent.py new file mode 100644 index 00000000..1eb7ef55 --- /dev/null +++ b/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 diff --git a/app/backend/main.py b/app/backend/main.py new file mode 100644 index 00000000..10f18f3f --- /dev/null +++ b/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, + ) diff --git a/app/backend/pyproject.toml b/app/backend/pyproject.toml new file mode 100644 index 00000000..4c95a415 --- /dev/null +++ b/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"