diff --git a/CLAUDE.md b/CLAUDE.md index 1c0021f7..61b0ea89 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -25,8 +25,8 @@ CEO (Renzo - Human) ### Hardware Infrastructure -- **Olares One (Powerhouse)**: Intel Ultra 9 + RTX 5090, runs Claude Code instances and AI inference -- **UGREEN NAS (Warehouse)**: 36TB RAID6, hosts PostgreSQL, Redis +- **Olares One (Powerhouse)**: Intel Ultra 9 + RTX 5090, runs Claude Code instances and AI inference - NOT YET ARRIVED +- **UGREEN NAS (Warehouse)**: 36TB RAID6, 128GB RAM, hosts PostgreSQL, Redis - **Pi Cluster (Operations)**: Monitoring, notifications, smart home ## Development Standards @@ -326,7 +326,6 @@ ROBOCO_RAG_USE_HYDE=true ROBOCO_RAG_USE_HYBRID_SEARCH=true # AI/LLM -ROBOCO_DEFAULT_LLM_MODEL=claude-opus-4-5-20251101 ROBOCO_DEFAULT_EMBEDDING_MODEL=embeddinggemma:300m ROBOCO_LOCAL_LLM_MODEL=gemma3:4b ROBOCO_LOCAL_LLM_BASE_URL=http://roboco-ollama:11434/v1 diff --git a/agents/prompts/base.md b/agents/prompts/base.md index 6110c2e6..6400568e 100644 --- a/agents/prompts/base.md +++ b/agents/prompts/base.md @@ -114,12 +114,11 @@ These tools let you inspect git state: ## Knowledge Base Tools -- `roboco_kb_search(query)` - Search code, docs, decisions -- `roboco_rag_query(question)` - AI-generated answers -- `roboco_kb_stats()` - See what's indexed -- `roboco_search_error(pattern)` - Find error solutions +- `roboco_ask_mentor(question)` - **Primary tool** - AI answers with follow-up support +- `roboco_kb_search(query)` - Raw semantic search +- `roboco_search_error(error_message)` - Find known error solutions - `roboco_check_decision(topic)` - Find past decisions -- `roboco_search_learnings(topic)` - Find team learnings +- `roboco_search_learnings(query)` - Find team learnings ## Journaling (ALL agents) @@ -147,8 +146,7 @@ Need docs updated? Create a task for your cell's Documenter. ## RAG Checkpoints -Before critical actions, verify with RAG: -- **Full workflow example**: `roboco_kb_search("{your_role} workflow")` -- **Tool parameters**: `roboco_kb_search("mcp tools")` -- **Error solutions**: `roboco_search_error(pattern)` -- **Past decisions**: `roboco_check_decision(topic)` +Before critical actions, check the knowledge base: +- `roboco_ask_mentor("How do I implement X?")` - Best practices, patterns +- `roboco_search_error(pattern)` - Known error solutions +- `roboco_check_decision(topic)` - Past architectural decisions diff --git a/alembic/versions/001_initial_schema.py b/alembic/versions/001_initial_schema.py index cfcfefcf..0b78b8a8 100644 --- a/alembic/versions/001_initial_schema.py +++ b/alembic/versions/001_initial_schema.py @@ -258,6 +258,12 @@ def upgrade() -> None: server_default="code", ), sa.Column("requires_git", sa.Boolean(), nullable=False, server_default="true"), + sa.Column( + "nature", + sa.Enum("technical", "non_technical", name="tasknature"), + nullable=False, + server_default="technical", + ), sa.Column( "project_id", postgresql.UUID(as_uuid=True), @@ -941,6 +947,7 @@ def downgrade() -> None: # Drop enums op.execute("DROP TYPE IF EXISTS worksessionstatus") op.execute("DROP TYPE IF EXISTS tasktype") + op.execute("DROP TYPE IF EXISTS tasknature") op.execute("DROP TYPE IF EXISTS handoffstatus") op.execute("DROP TYPE IF EXISTS journalentrytype") op.execute("DROP TYPE IF EXISTS notificationpriority") diff --git a/docker-compose.yml b/docker-compose.yml index 6d56c7c5..7889d28f 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -38,6 +38,55 @@ services: timeout: 5s retries: 5 + # ========================================================================== + # Ollama - Local LLM and Embedding Server + # ========================================================================== + ollama: + image: ollama/ollama:latest + container_name: roboco-ollama + restart: unless-stopped + ports: + - "11435:11434" + volumes: + - ${ROBOCO_DATA_DIR:-./data}/ollama:/root/.ollama + healthcheck: + # Use ollama CLI (guaranteed available) to check if server is responding + test: ["CMD", "ollama", "list"] + interval: 10s + timeout: 5s + retries: 5 + start_period: 10s + + # Ollama model puller - pulls required models on startup + # Uses streaming curl to wait for full model download + ollama-init: + image: curlimages/curl:latest + container_name: roboco-ollama-init + depends_on: + ollama: + condition: service_healthy + restart: "no" + entrypoint: ["/bin/sh", "-c"] + command: + - | + set -e + echo "=== Pulling embedding model (embeddinggemma:300m) ===" + # Ollama /api/pull streams JSON lines until complete - consume full stream + # Note: $$ escapes $ for docker-compose variable substitution + curl -sN http://ollama:11434/api/pull -d '{"name":"embeddinggemma:300m"}' | while read -r line; do + status=$$(echo "$$line" | grep -o '"status":"[^"]*"' | cut -d'"' -f4) + [ -n "$$status" ] && echo " $$status" + done + echo "=== Pulling LLM model (gemma3:4b) ===" + curl -sN http://ollama:11434/api/pull -d '{"name":"gemma3:4b"}' | while read -r line; do + status=$$(echo "$$line" | grep -o '"status":"[^"]*"' | cut -d'"' -f4) + [ -n "$$status" ] && echo " $$status" + done + echo "=== Verifying models are available ===" + curl -sf http://ollama:11434/api/tags | grep -q "embeddinggemma" && echo " embeddinggemma: OK" + curl -sf http://ollama:11434/api/tags | grep -q "gemma3" && echo " gemma3: OK" + echo "=== All models ready! ===" + # ========================================================================== # Agent Base Image Builder (specialized images built on-demand by orchestrator) # ========================================================================== @@ -166,6 +215,11 @@ services: # API ROBOCO_HOST: 0.0.0.0 ROBOCO_PORT: 8000 + # Ollama (use container name) + ROBOCO_LOCAL_LLM_BASE_URL: http://roboco-ollama:11434/v1 + ROBOCO_LOCAL_LLM_MODEL: gemma3:4b + ROBOCO_DEFAULT_EMBEDDING_MODEL: embeddinggemma:300m + ROBOCO_OLLAMA_BASE_URL: http://roboco-ollama:11434 # Host paths for spawning agent containers (required for Docker-in-Docker) # IMPORTANT: These must be ABSOLUTE paths on the host filesystem ROBOCO_HOST_PROJECT_DIR: ${ROBOCO_HOST_PROJECT_DIR:-/volume1/roboco} @@ -185,6 +239,10 @@ services: condition: service_healthy redis: condition: service_healthy + ollama: + condition: service_healthy + ollama-init: + condition: service_completed_successfully agent-base-image: condition: service_completed_successfully # Default agents to spawn (override in .env or command line) diff --git a/roboco/api/app.py b/roboco/api/app.py index 0985dc43..eb90e2b9 100644 --- a/roboco/api/app.py +++ b/roboco/api/app.py @@ -83,9 +83,11 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None]: app.state.transcription = _AppServices.transcription app.state.extraction = _AppServices.extraction - # Initialize Phase 3 services (RAG - optional, non-blocking) - # OptimalService.initialize() auto-indexes /docs/standards/ and /docs/workflows/ + # Initialize OptimalService (RAG) - BLOCKS until fully ready + # This ensures /health only returns 200 when RAG is operational + # Typical initialization time: 30-90 seconds (embedding + indexing) try: + logger.info("Initializing OptimalService (RAG)...") optimal_service = await get_optimal_service() app.state.optimal = optimal_service logger.info("OptimalService (RAG) initialized successfully") @@ -96,7 +98,7 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None]: ) app.state.optimal = None - logger.info("All services initialized") + logger.info("All services initialized, API ready") yield diff --git a/roboco/api/routes/optimal.py b/roboco/api/routes/optimal.py index 1437fd0c..aa510fb1 100644 --- a/roboco/api/routes/optimal.py +++ b/roboco/api/routes/optimal.py @@ -309,6 +309,18 @@ async def rag_query( status_code=status.HTTP_504_GATEWAY_TIMEOUT, detail=f"RAG query timed out after {rag_timeout}s", ) from e + except RuntimeError as e: + # Service not initialized or other runtime errors + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail=f"RAG service error: {e}", + ) from e + except Exception as e: + # Catch-all for unexpected errors + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=f"RAG query failed: {e}", + ) from e return RAGQueryResponse( answer=response.answer, @@ -448,6 +460,26 @@ async def get_single_index_stats( ) +@router.get("/stats/staleness") +async def check_staleness( + agent: CurrentAgentContext, + permissions: PermissionServiceDep, +) -> dict[str, Any]: + """ + Check if indexes are stale (source files modified after last indexing). + + Returns staleness info for file-based indexes (CODE, DOCUMENTATION). + """ + if not permissions.can_perform_kb_action(agent, KBAction.VIEW_STATS): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="Not authorized to view index staleness", + ) + + service = await get_optimal_service() + return await service.check_index_staleness() + + @router.get("/health", response_model=RAGHealthResponse) async def rag_health_check() -> RAGHealthResponse: """ @@ -481,7 +513,13 @@ async def rag_health_check() -> RAGHealthResponse: try: async with asyncio.timeout(health_timeout): embedder = await get_shared_embedder(model=settings.default_embedding_model) - test_embedding = embedder.embed_query("health check") + # Use async method if available (OllamaEmbedder), else run sync in thread + if hasattr(embedder, "aembed_query"): + test_embedding = await embedder.aembed_query("health check") + else: + test_embedding = await asyncio.to_thread( + embedder.embed_query, "health check" + ) if test_embedding and len(test_embedding) == settings.embedding_dimensions: embedding_ok = True details["embedding_model"] = settings.default_embedding_model @@ -673,6 +711,7 @@ async def reindex_all( agent: CurrentAgentContext, permissions: PermissionServiceDep, force: bool = False, + timeout_seconds: int = 300, # 5 minute default ) -> dict[str, Any]: """ Trigger re-indexing of code and documentation. @@ -683,19 +722,31 @@ async def reindex_all( Args: force: If True, reindex even if indexes aren't empty + timeout_seconds: Maximum time to wait for reindexing (default: 300s) Returns: - Count of indexed code files and documentation files + Detailed indexing report with success/failure counts """ + import asyncio + if not permissions.can_perform_kb_action(agent, KBAction.INDEX_CODE): raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail="Not authorized to trigger reindexing", ) - service = await get_optimal_service() - result = await service.auto_index_on_startup(force=force) - return {"status": "reindexed", **result} + try: + async with asyncio.timeout(timeout_seconds): + service = await get_optimal_service() + # Call the private method that returns the report + report = await service._auto_index_on_startup(force=force) + return {"status": "reindexed", **report.to_dict()} + except TimeoutError as e: + raise HTTPException( + status_code=status.HTTP_504_GATEWAY_TIMEOUT, + detail=f"Reindexing timed out after {timeout_seconds} seconds. " + "Try indexing smaller directories or increasing timeout.", + ) from e # ============================================================================= diff --git a/roboco/api/routes/tasks.py b/roboco/api/routes/tasks.py index e9db4704..e094a03f 100644 --- a/roboco/api/routes/tasks.py +++ b/roboco/api/routes/tasks.py @@ -105,6 +105,7 @@ async def create_task( assigned_to=data.assigned_to, target_date=data.target_date, estimated_complexity=data.estimated_complexity, + nature=data.nature, status=data.status, sequence=data.sequence, # Task ordering within siblings dependency_ids=data.dependency_ids, # Dependencies for claim filtering diff --git a/roboco/api/schemas/tasks.py b/roboco/api/schemas/tasks.py index f009963a..7fa0b5c5 100644 --- a/roboco/api/schemas/tasks.py +++ b/roboco/api/schemas/tasks.py @@ -10,7 +10,7 @@ from uuid import UUID from pydantic import BaseModel, Field -from roboco.models.base import Complexity, TaskStatus, Team +from roboco.models.base import Complexity, TaskNature, TaskStatus, TaskType, Team from roboco.models.session import SessionScope from roboco.utils.converters import require_uuid, to_python_uuid, to_python_uuid_list @@ -225,6 +225,19 @@ class TaskResponse(BaseModel): status: TaskStatus priority: int sequence: int # Order number within siblings + nature: TaskNature # Technical or non-technical work + + # Task Type & Git Configuration + task_type: TaskType # code, documentation, research, etc. + requires_git: bool # Whether this task requires git workflow + project_id: UUID | None = None # Project this task works on + + # Parallel Execution Tracking (for AWAITING_DOCUMENTATION phase) + docs_complete: bool = False # Documenter has finished + pr_created: bool = False # Developer has created PR + + # PM Approval Tracking (for AWAITING_PM_REVIEW phase) + pm_approvals: dict[str, bool] = {} # e.g. {'main_pm': True, 'cell_pm': True} # Ownership team: Team @@ -291,6 +304,7 @@ class TaskSummaryResponse(BaseModel): created_at: datetime updated_at: datetime | None estimated_complexity: Complexity + nature: TaskNature class Config: from_attributes = True @@ -530,6 +544,17 @@ def task_to_response(task: "TaskTable") -> TaskResponse: status=task.status, priority=task.priority, sequence=task.sequence, + nature=task.nature, + # Task Type & Git Configuration + task_type=task.task_type, + requires_git=task.requires_git, + project_id=to_python_uuid(task.project_id), + # Parallel Execution Tracking + docs_complete=task.docs_complete, + pr_created=task.pr_created, + # PM Approval Tracking + pm_approvals=task.pm_approvals or {}, + # Ownership team=task.team, created_by=require_uuid(task.created_by), assigned_to=to_python_uuid(task.assigned_to), diff --git a/roboco/bootstrap.py b/roboco/bootstrap.py index 8aafcd28..47144f3a 100644 --- a/roboco/bootstrap.py +++ b/roboco/bootstrap.py @@ -6,8 +6,10 @@ Data constants are in roboco/seeds/, database operations in roboco/db/seed.py. """ import asyncio +from http import HTTPStatus from pathlib import Path +import httpx import structlog import uvicorn @@ -42,6 +44,29 @@ async def _run_api_server() -> None: await server.serve() +async def _wait_for_api_ready(max_wait: int = 120) -> None: + """ + Wait for the API server to be ready to accept connections. + + The lifespan does document indexing which can take 30+ seconds, + so we poll the health endpoint instead of using a fixed sleep. + """ + api_url = f"http://127.0.0.1:{settings.port}/health" + waited = 0 + while waited < max_wait: + try: + async with httpx.AsyncClient(timeout=2.0) as client: + resp = await client.get(api_url) + if resp.status_code == HTTPStatus.OK: + logger.info("API server ready", waited_seconds=waited) + return + except Exception: + pass + await asyncio.sleep(2) + waited += 2 + logger.warning("API server not ready after timeout, starting orchestrator anyway") + + async def main( skip_db: bool = False, skip_orchestrator: bool = False, @@ -100,8 +125,9 @@ async def main( api_task = asyncio.create_task(_run_api_server()) logger.info("API server starting", host=settings.host, port=settings.port) - # Wait a moment for API to start before orchestrator begins polling - await asyncio.sleep(2) + # Wait for API to actually be ready (not just a fixed sleep) + # The lifespan does document indexing which can take 30+ seconds + await _wait_for_api_ready() await orchestrator.start() diff --git a/roboco/config.py b/roboco/config.py index 7444f42d..ebe91ca6 100644 --- a/roboco/config.py +++ b/roboco/config.py @@ -165,24 +165,27 @@ class Settings(BaseSettings): openai_api_key: str | None = None # For embeddings # Default models - default_llm_model: str = "claude-opus-4-5-20251101" default_embedding_model: str = Field( - default="BAAI/bge-base-en-v1.5", - description="Embedding model", + default="embeddinggemma:300m", + description="Embedding model (Ollama model name)", ) embedding_dimensions: int = Field( default=768, - description="Embedding dimensions (768 for BGE-base)", + description="Embedding dimensions (768 for embeddinggemma/BGE-base)", ) # Local LLM for RAG (HyDE, reranking, etc.) local_llm_model: str = Field( - default="qwen3:8b", - description="Local LLM model for HyDE and RAG operations", + default="gemma3:4b", + description="Local LLM for HyDE/RAG (non-thinking models are faster)", ) local_llm_base_url: str = Field( - default="http://192.168.50.111:11434/v1", - description="Base URL for local LLM (Ollama)", + default="http://roboco-ollama:11434/v1", + description="Base URL for local LLM (Ollama OpenAI-compat API)", + ) + ollama_base_url: str = Field( + default="http://roboco-ollama:11434", + description="Base URL for Ollama native API (embeddings, model mgmt)", ) # ========================================================================== diff --git a/roboco/db/tables.py b/roboco/db/tables.py index 30d6a6a1..a76b7587 100644 --- a/roboco/db/tables.py +++ b/roboco/db/tables.py @@ -38,6 +38,7 @@ from roboco.models.base import ( NotificationPriority, NotificationType, SessionStatus, + TaskNature, TaskStatus, TaskType, Team, @@ -137,6 +138,9 @@ class TaskTable(Base): task_type: Mapped[TaskType] = mapped_column( Enum(TaskType), nullable=False, default=TaskType.CODE ) + nature: Mapped[TaskNature] = mapped_column( + Enum(TaskNature), nullable=False, default=TaskNature.TECHNICAL + ) requires_git: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True) # Project & Branch (set by PM during setup) diff --git a/roboco/mcp/optimal_server.py b/roboco/mcp/optimal_server.py index 1675d176..4be90a5c 100644 --- a/roboco/mcp/optimal_server.py +++ b/roboco/mcp/optimal_server.py @@ -88,15 +88,22 @@ def _register_search_tools(mcp: FastMCP, client: ApiClient) -> None: "SEARCH_FAILED", "Failed to search knowledge base", {"api_error": resp.text}, + hint="Try roboco_ask_mentor(question) for AI-synthesized answers.", ) result = resp.json() - return { + total = result.get("total", 0) + response: dict[str, Any] = { "status": "success", "query": query, - "total": result.get("total", 0), + "total": total, "results": result.get("results", []), } + if total == 0: + response["hint"] = ( + "No results. Try roboco_ask_mentor(question) for better answers." + ) + return response @mcp.tool() async def roboco_rag_query( @@ -108,11 +115,12 @@ def _register_search_tools(mcp: FastMCP, client: ApiClient) -> None: """ RAG query - get an AI-generated answer using knowledge base context. - Use this when you need an answer synthesized from the knowledge base, - not just search results. Good for questions like: - - "How does authentication work in this codebase?" - - "What's the pattern for error handling?" - - "What decisions were made about the database schema?" + NOTE: For most questions, prefer `roboco_ask_mentor` instead! + The mentor searches multiple indexes and supports follow-up questions. + + Use this simpler tool only when you need: + - A quick answer from a specific index type + - To filter by project or task_id Args: query: Natural language question @@ -132,22 +140,33 @@ def _register_search_tools(mcp: FastMCP, client: ApiClient) -> None: if task_id: payload["task_id"] = task_id - resp = await client.post("/optimal/rag/query", json=payload) + # RAG queries can take longer due to LLM call - use 65s timeout + resp = await client.post("/optimal/rag/query", json=payload, timeout=65.0) if not resp.ok: return format_error_response( "RAG_FAILED", "Failed to query RAG", {"api_error": resp.text}, + hint="Try roboco_ask_mentor(question) instead - it's more robust.", ) result = resp.json() - return { + answer = result.get("answer", "") + context_used = result.get("context_used", 0) + response: dict[str, Any] = { "status": "success", "query": query, - "answer": result.get("answer", ""), + "answer": answer, "citations": result.get("citations", []), - "context_used": result.get("context_used", 0), + "context_used": context_used, } + # Guide to mentor for better results + if context_used == 0 or "couldn't find" in answer.lower(): + response["hint"] = ( + "Limited results. roboco_ask_mentor(question) searches more sources " + "and supports follow-up questions." + ) + return response @mcp.tool() async def roboco_kb_stats() -> dict[str, Any]: @@ -336,6 +355,9 @@ def _register_mentor_tools(mcp: FastMCP, client: ApiClient) -> None: """ Ask the organizational knowledge base for help. + THIS IS THE PRIMARY TOOL for knowledge base questions. + Use this instead of roboco_rag_query for most questions. + This is a conversational interface - you can ask follow-up questions by providing the conversation_id from a previous response. @@ -344,6 +366,7 @@ def _register_mentor_tools(mcp: FastMCP, client: ApiClient) -> None: - Past architectural decisions - Team learnings and reflections - Codebase patterns + - Known error solutions Args: question: Your question (natural language) @@ -369,7 +392,8 @@ def _register_mentor_tools(mcp: FastMCP, client: ApiClient) -> None: if domain: payload["domain"] = domain - resp = await client.post("/optimal/mentor/ask", json=payload) + # Mentor uses LLM - allow 65s timeout + resp = await client.post("/optimal/mentor/ask", json=payload, timeout=65.0) if not resp.ok: return format_error_response( "MENTOR_FAILED", @@ -421,12 +445,20 @@ def _register_error_tools(mcp: FastMCP, client: ApiClient) -> None: ) result = resp.json() - return { + solutions_found = len(result.get("results", [])) + response: dict[str, Any] = { "status": "success", "error_message": error_message, - "solutions_found": len(result.get("results", [])), + "solutions_found": solutions_found, "results": result.get("results", []), } + if solutions_found == 0: + response["hint"] = ( + "No known solutions. Try roboco_ask_mentor(f'How do I fix: {error}') " + "for guidance. If you solve it, use roboco_record_error_solution() " + "to help future agents." + ) + return response @mcp.tool() async def roboco_record_error_solution( @@ -806,12 +838,20 @@ def _register_learning_tools(mcp: FastMCP, client: ApiClient) -> None: ) result = resp.json() - return { + total = result.get("total", 0) + response: dict[str, Any] = { "status": "success", "query": query, - "total": result.get("total", 0), + "total": total, "results": result.get("results", []), } + if total == 0: + response["hint"] = ( + "No learnings found. Try roboco_ask_mentor(question) for broader " + "knowledge. If you learn something useful, use " + "roboco_record_learning() to share it." + ) + return response def _register_index_management_tools(mcp: FastMCP, client: ApiClient) -> None: diff --git a/roboco/mcp/schemas/__init__.py b/roboco/mcp/schemas/__init__.py index 2028f059..f88478ff 100644 --- a/roboco/mcp/schemas/__init__.py +++ b/roboco/mcp/schemas/__init__.py @@ -184,6 +184,9 @@ class TaskCreateInput(BaseModel): complexity: str = Field( default="medium", description="Complexity: low, medium, high, critical" ) + nature: str = Field( + default="technical", description="Task nature: technical, non_technical" + ) status: str = Field( default="backlog", description="Status: 'backlog' (default) or 'pending' (ready for work)", diff --git a/roboco/mcp/task_server.py b/roboco/mcp/task_server.py index c25b5a1b..6ae5ee87 100644 --- a/roboco/mcp/task_server.py +++ b/roboco/mcp/task_server.py @@ -648,6 +648,7 @@ def _register_pm_tools(mcp: FastMCP, client: ApiClient, agent_id: str) -> None: data: TaskCreateInput with: - title, description, acceptance_criteria, team (required) - parent_task_id, assigned_to, priority, status (optional) + - nature: Task nature (technical, non_technical) - sequence: Order within siblings (0 = default) - dependency_ids: Task IDs that must complete first diff --git a/roboco/mcp/tasks/handlers/management.py b/roboco/mcp/tasks/handlers/management.py index 731f399c..fd5809dc 100644 --- a/roboco/mcp/tasks/handlers/management.py +++ b/roboco/mcp/tasks/handlers/management.py @@ -257,6 +257,7 @@ def _build_task_payload(input_data: TaskCreateInput) -> dict[str, Any]: "team": input_data.team, "priority": input_data.priority, "estimated_complexity": input_data.complexity, + "nature": input_data.nature, "status": input_data.status, # Always included, defaults to "backlog" "sequence": input_data.sequence, # Task ordering (lower = first) } diff --git a/roboco/models/__init__.py b/roboco/models/__init__.py index 500e8a6c..f9ac4629 100644 --- a/roboco/models/__init__.py +++ b/roboco/models/__init__.py @@ -26,6 +26,7 @@ from roboco.models.base import ( SessionStatus, SubstituteReason, # Enums + TaskNature, TaskStatus, Team, ) @@ -143,6 +144,7 @@ __all__ = [ "SubstituteReason", "Task", "TaskCreate", + "TaskNature", "TaskPlan", "TaskStatus", "TaskUpdate", diff --git a/roboco/models/base.py b/roboco/models/base.py index 23c19fc5..fdd5ea3a 100644 --- a/roboco/models/base.py +++ b/roboco/models/base.py @@ -47,6 +47,13 @@ class TaskType(str, Enum): ADMINISTRATIVE = "administrative" # No git +class TaskNature(str, Enum): + """Task nature classification - technical vs non-technical work.""" + + TECHNICAL = "technical" + NON_TECHNICAL = "non_technical" + + class Complexity(str, Enum): """Task complexity levels.""" diff --git a/roboco/models/task.py b/roboco/models/task.py index f106f8b0..2550d69e 100644 --- a/roboco/models/task.py +++ b/roboco/models/task.py @@ -15,6 +15,7 @@ from pydantic import Field from roboco.models.base import ( Complexity, RobocoBase, + TaskNature, TaskStatus, TaskType, Team, @@ -158,6 +159,9 @@ class Task(TimestampMixin): task_type: TaskType = Field( default=TaskType.CODE, description="Type of task (code, research, etc.)" ) + nature: TaskNature = Field( + default=TaskNature.TECHNICAL, description="Technical or non-technical work" + ) requires_git: bool = Field( default=True, description="Whether this task requires git workflow" ) @@ -290,6 +294,7 @@ class TaskCreate(RobocoBase): # Git configuration task_type: TaskType = TaskType.CODE + nature: TaskNature = TaskNature.TECHNICAL requires_git: bool = True project_id: UUID | None = None @@ -311,6 +316,7 @@ class TaskUpdate(RobocoBase): # Git fields task_type: TaskType | None = None + nature: TaskNature | None = None requires_git: bool | None = None project_id: UUID | None = None branch_name: str | None = None @@ -348,5 +354,6 @@ class TaskCreateRequest: # Git configuration task_type: TaskType = field(default=TaskType.CODE) + nature: TaskNature = field(default=TaskNature.TECHNICAL) requires_git: bool = True project_id: UUID | None = None diff --git a/roboco/services/optimal.py b/roboco/services/optimal.py index 95ef8423..dd661ae6 100644 --- a/roboco/services/optimal.py +++ b/roboco/services/optimal.py @@ -9,6 +9,7 @@ The service uses a plugin-based architecture where each index type is handled by a specialized plugin that implements the BaseIndexPlugin interface. """ +from dataclasses import dataclass, field from pathlib import Path from typing import Any @@ -48,6 +49,70 @@ from roboco.services.optimal_brain.indexes.reviews import ( logger = structlog.get_logger() +@dataclass +class IndexingReport: + """ + Detailed report of indexing operation results. + + Provides visibility into what was indexed successfully vs what failed, + enabling proper error handling and recovery. + """ + + index_type: str + total_attempted: int = 0 + successful: int = 0 + failed: int = 0 + skipped: int = 0 # Already indexed or filtered out + failed_sources: list[tuple[str, str]] = field(default_factory=list) + duration_seconds: float = 0.0 + + @property + def success_rate(self) -> float: + """Percentage of attempted items that succeeded.""" + if self.total_attempted == 0: + return 100.0 + return (self.successful / self.total_attempted) * 100 + + @property + def has_failures(self) -> bool: + """True if any items failed.""" + return self.failed > 0 + + def to_dict(self) -> dict[str, Any]: + """Convert to dictionary for API responses.""" + return { + "index_type": self.index_type, + "total_attempted": self.total_attempted, + "successful": self.successful, + "failed": self.failed, + "skipped": self.skipped, + "success_rate": round(self.success_rate, 1), + "has_failures": self.has_failures, + "failed_sources": self.failed_sources[:10], # Limit for API + "duration_seconds": round(self.duration_seconds, 2), + } + + +@dataclass +class AutoIndexReport: + """Combined report for auto-indexing on startup.""" + + code: IndexingReport | None = None + documentation: IndexingReport | None = None + overall_success: bool = True + warnings: list[str] = field(default_factory=list) + + def to_dict(self) -> dict[str, Any]: + """Convert to dictionary for API responses.""" + docs = self.documentation.to_dict() if self.documentation else None + return { + "code": self.code.to_dict() if self.code else None, + "documentation": docs, + "overall_success": self.overall_success, + "warnings": self.warnings, + } + + # Plugin registry mapping IndexType to plugin class PLUGIN_REGISTRY: dict[IndexType, type[BaseIndexPlugin]] = { IndexType.CODE: CodeIndexPlugin, @@ -87,28 +152,112 @@ class OptimalService: self._initialized = False self._plugins: dict[IndexType, BaseIndexPlugin] = {} self._prompt_templates: dict[str, dict[str, Any]] = {} + self._indexing_task: Any = None # Background indexing task async def initialize(self) -> None: - """Initialize all knowledge base indexes.""" + """ + Initialize all knowledge base indexes. + + Uses graceful degradation - if a plugin fails to initialize, log the error + and continue with remaining plugins. The service will function with + reduced capabilities rather than completely failing. + """ if self._initialized: return logger.info("Initializing OptimalService with plugin architecture") + # Track initialization results for reporting + initialized_count = 0 + failed_plugins: list[tuple[IndexType, str]] = [] + + import asyncio + + # Per-plugin initialization timeout (embedding validation can be slow) + plugin_init_timeout = 30.0 + # Create and initialize plugins for each index type for index_type, plugin_class in PLUGIN_REGISTRY.items(): - plugin = plugin_class() - await plugin.initialize() - self._plugins[index_type] = plugin - logger.info(f"Initialized {index_type.value} plugin") + try: + plugin = plugin_class() + async with asyncio.timeout(plugin_init_timeout): + await plugin.initialize() + self._plugins[index_type] = plugin + initialized_count += 1 + logger.info(f"Initialized {index_type.value} plugin") + except TimeoutError: + error_msg = f"Plugin init timed out ({plugin_init_timeout}s)" + failed_plugins.append((index_type, error_msg)) + logger.error( + "Plugin initialization timeout - continuing with degraded mode", + index_type=index_type.value, + timeout=plugin_init_timeout, + ) + except Exception as e: + # Log error but continue with other plugins + error_msg = str(e) + failed_plugins.append((index_type, error_msg)) + logger.error( + "Failed to initialize plugin - continuing with degraded mode", + index_type=index_type.value, + error=error_msg, + ) - self._initialized = True - logger.info("OptimalService initialization complete") + # Service is initialized if at least one plugin succeeded + if initialized_count > 0: + self._initialized = True + logger.info( + "OptimalService initialization complete", + initialized=initialized_count, + failed=len(failed_plugins), + ) + else: + # All plugins failed - this is a critical error + raise RuntimeError( + f"OptimalService failed to initialize any plugins. " + f"Errors: {failed_plugins}" + ) - # Auto-index code and documentation on startup - await self._auto_index_on_startup() + # Report failed plugins for debugging + if failed_plugins: + logger.warning( + "Some index plugins failed to initialize", + failed=[f"{idx.value}: {err}" for idx, err in failed_plugins], + ) - async def _auto_index_on_startup(self) -> None: + # Auto-index code and documentation on startup (truly non-blocking) + # Run in background so API can start accepting requests immediately + self._indexing_task = asyncio.create_task(self._auto_index_on_startup_safe()) + + async def _auto_index_on_startup_safe(self) -> None: + """ + Safe wrapper for auto-indexing that catches all errors. + + Runs auto-indexing in background without blocking API startup. + Logs errors but doesn't crash the service if Ollama is unavailable. + """ + try: + report = await self._auto_index_on_startup() + if report.warnings: + logger.warning( + "Auto-indexing completed with warnings", + warnings=report.warnings, + ) + else: + logger.info( + "Auto-indexing completed successfully", + code_indexed=report.code.successful if report.code else 0, + docs_indexed=report.documentation.successful + if report.documentation + else 0, + ) + except Exception as e: + logger.error( + "Auto-indexing failed - service operational but indexes may be empty", + error=str(e), + ) + + async def _auto_index_on_startup(self, force: bool = False) -> AutoIndexReport: """ Auto-index code and documentation on startup. @@ -117,32 +266,55 @@ class OptimalService: - /docs/standards/ - Coding, security, workflow standards - /docs/workflows/ - Agent workflow documentation - This ensures agents can search for code, standards, and workflows - immediately after startup. - """ - await self._auto_index_code() - await self._auto_index_docs() + Args: + force: If True, reindex even if indexes already have content - async def _auto_index_code(self) -> None: + Returns: + AutoIndexReport with detailed results for each index type + """ + report = AutoIndexReport() + + # Index code + code_report = await self._auto_index_code(force=force) + report.code = code_report + if code_report and code_report.has_failures: + report.warnings.append(f"Code indexing had {code_report.failed} failures") + + # Index documentation + docs_report = await self._auto_index_docs(force=force) + report.documentation = docs_report + if docs_report and docs_report.has_failures: + report.warnings.append( + f"Documentation indexing had {docs_report.failed} failures" + ) + + # Overall success if at least something was indexed + total_successful = (code_report.successful if code_report else 0) + ( + docs_report.successful if docs_report else 0 + ) + report.overall_success = total_successful > 0 or not report.warnings + + return report + + async def _auto_index_code(self, force: bool = False) -> IndexingReport | None: """Auto-index source code files on startup.""" - # Find the roboco source directory (the Python package, not the repo root) - # We want to index roboco/ package, NOT the entire repo (which has .venv) + import time + + report = IndexingReport(index_type="code") + start_time = time.time() + + # Find the roboco source directory possible_code_roots = [ - Path("/app/roboco"), # Docker: /app is repo root, roboco/ is package - Path(__file__).parent.parent, # Local: optimal.py -> services -> roboco - Path.cwd() / "roboco", # Local: cwd/roboco + Path("/app/roboco"), # Docker + Path(__file__).parent.parent, # Local + Path.cwd() / "roboco", ] code_root = None for path in possible_code_roots: - # Check for __init__.py to confirm it's a Python package (not repo root) init_file = path / "__init__.py" if path.exists() and path.is_dir() and init_file.exists(): code_root = path - logger.debug( - "Found code package directory", - path=str(path), - ) break if code_root is None: @@ -150,44 +322,57 @@ class OptimalService: "Code directory not found", searched_paths=[str(p) for p in possible_code_roots], ) - return + return None - # Check if code index is empty - code_plugin = self._get_plugin(IndexType.CODE) - code_count = await code_plugin.count() + # Check if code index is empty (unless force=True) + if not force: + code_plugin = self._get_plugin(IndexType.CODE) + code_count = await code_plugin.count() + if code_count > 0: + logger.info( + "Code index already populated", + chunk_count=code_count, + skipping=True, + ) + report.skipped = code_count + report.duration_seconds = time.time() - start_time + return report - if code_count > 0: - logger.info( - "Code index already populated", - chunk_count=code_count, - skipping=True, - ) - return - - logger.info( - "Auto-indexing source code", - directory=str(code_root), - ) + logger.info("Auto-indexing source code", directory=str(code_root)) try: - count = await self.index_code([str(code_root)], project="roboco") + from roboco.services.optimal_brain.indexes.code import MAX_AUTO_INDEX_FILES + + count = await self.index_code( + [str(code_root)], + project="roboco", + max_files=MAX_AUTO_INDEX_FILES, + ) + report.successful = count + report.total_attempted = count logger.info("Code auto-indexing complete", files_indexed=count) except Exception as e: - logger.warning("Code auto-indexing failed", error=str(e)) + error_msg = str(e) + report.failed = 1 + report.total_attempted = 1 + report.failed_sources.append((str(code_root), error_msg)) + logger.error("Code auto-indexing failed", error=error_msg) - async def _auto_index_docs(self) -> None: - """ - Auto-index documentation directories on startup. + report.duration_seconds = time.time() - start_time + return report - Indexes: - - /docs/standards/ - Coding, security, workflow standards - - /docs/workflows/ - Agent workflow documentation - """ - # Find the docs directory relative to the project root + async def _auto_index_docs(self, force: bool = False) -> IndexingReport | None: + """Auto-index documentation directories on startup.""" + import time + + report = IndexingReport(index_type="documentation") + start_time = time.time() + + # Find the docs directory possible_docs_roots = [ - Path("/app/docs"), # Docker absolute path - Path(__file__).parent.parent.parent / "docs", # roboco/docs (local) - Path.cwd() / "docs", # Current working directory + Path("/app/docs"), + Path(__file__).parent.parent.parent / "docs", + Path.cwd() / "docs", ] docs_root = None @@ -201,61 +386,75 @@ class OptimalService: "Docs directory not found", searched_paths=[str(p) for p in possible_docs_roots], ) - return + return None - # Directories to auto-index (relative to docs root) + # Directories to auto-index auto_index_dirs = ["standards", "workflows"] for subdir in auto_index_dirs: target_dir = docs_root / subdir if not target_dir.exists(): - logger.debug(f"Auto-index directory not found: {target_dir}") continue - await self._index_docs_directory(target_dir, subdir) + subdir_report = await self._index_docs_directory( + target_dir, subdir, _force=force + ) + # Aggregate results + report.total_attempted += subdir_report.total_attempted + report.successful += subdir_report.successful + report.failed += subdir_report.failed + report.skipped += subdir_report.skipped + report.failed_sources.extend(subdir_report.failed_sources) - async def _index_docs_directory(self, directory: Path, name: str) -> None: + report.duration_seconds = time.time() - start_time + return report + + async def _index_docs_directory( + self, directory: Path, name: str, _force: bool = False + ) -> IndexingReport: """Index all markdown files in a documentation directory.""" + report = IndexingReport(index_type=f"docs/{name}") + md_files = list(directory.rglob("*.md")) if not md_files: logger.info(f"No files found to index in {name}/", path=str(directory)) - return + return report + report.total_attempted = len(md_files) logger.info( f"Auto-indexing {name} files", directory=str(directory), file_count=len(md_files), ) - # Index each file - total_indexed = 0 + # Index each file with individual error tracking for md_file in md_files: try: if name == "standards": - count = await self.index_standards_file(str(md_file)) - total_indexed += count + await self.index_standards_file(str(md_file)) + report.successful += 1 else: - # For workflows, index as documentation await self.index_documentation([str(md_file)]) - total_indexed += 1 + report.successful += 1 - logger.debug( - f"Indexed {name} file", - file=str(md_file), - items_indexed=count if name == "standards" else 1, - ) + logger.debug(f"Indexed {name} file", file=str(md_file)) except Exception as e: + error_msg = str(e) + report.failed += 1 + report.failed_sources.append((str(md_file), error_msg)) logger.warning( f"Failed to index {name} file", file=str(md_file), - error=str(e), + error=error_msg, ) logger.info( f"{name.capitalize()} auto-indexing complete", - files_processed=len(md_files), - total_indexed=total_indexed, + successful=report.successful, + failed=report.failed, + total=report.total_attempted, ) + return report async def close(self) -> None: """Cleanup resources.""" @@ -286,11 +485,20 @@ class OptimalService: self, sources: list[str], project: str | None = None, + max_files: int | None = None, ) -> int: - """Index code files/directories and track in database.""" + """Index code files/directories and track in database. + + Args: + sources: List of file paths, directories, or glob patterns + project: Optional project identifier for filtering + max_files: Optional limit on files to index (for auto-indexing) + """ plugin = self._get_plugin(IndexType.CODE) if isinstance(plugin, CodeIndexPlugin): - count, indexed_files = await plugin.index_sources(sources, project) + count, indexed_files = await plugin.index_sources( + sources, project, max_files=max_files + ) # Batch track all indexed files using repository docs_to_track = [ @@ -644,7 +852,8 @@ class OptimalService: """ Query the knowledge base with RAG. - Retrieves relevant context and generates an answer. + Aggregates results from all indexes and generates a synthesized answer. + Skips empty indexes to avoid unnecessary LLM calls. """ if not self._initialized: raise RuntimeError("OptimalService not initialized") @@ -653,25 +862,71 @@ class OptimalService: context.index_types if context and context.index_types else list(IndexType) ) + logger.info( + "RAG query starting", query=query[:50], num_indexes=len(index_types) + ) + + # Aggregate citations and answers from all non-empty indexes + all_citations: list[SearchResult] = [] + best_answer: str = "" + for index_type in index_types: plugin = self._plugins.get(index_type) - if plugin: - try: - answer, citations = await plugin.ask(query=query, top_k=top_k) - if answer: - return RAGResponse( - answer=answer, - citations=citations, - query=query, - context_used=len(citations), - ) - except Exception as e: - logger.warning( - "RAG query failed for index", - index_type=index_type.value, - error=str(e), - ) + if not plugin: + continue + try: + # Skip empty indexes to save LLM calls + count = await plugin.count() + if count == 0: + logger.debug( + "Skipping empty index", + index_type=index_type.value, + ) + continue + + answer, citations = await plugin.ask(query=query, top_k=top_k) + all_citations.extend(citations) + + logger.info( + "RAG query index result", + index_type=index_type.value, + has_answer=bool(answer), + num_citations=len(citations), + ) + + # Keep the first real LLM answer we get + if answer and not best_answer: + best_answer = answer + + except Exception as e: + logger.warning( + "RAG query failed for index", + index_type=index_type.value, + error=str(e), + ) + + # Sort all citations by score + all_citations.sort(key=lambda r: r.score, reverse=True) + top_citations = all_citations[: top_k * 2] + + # If we have citations but no LLM answer, synthesize one + if top_citations and not best_answer: + logger.info( + "Synthesizing answer from aggregated citations", + num_citations=len(top_citations), + ) + best_answer = await self._synthesize_from_citations(query, top_citations) + + if best_answer: + return RAGResponse( + answer=best_answer, + citations=top_citations, + query=query, + context_used=len(top_citations), + ) + + logger.warning("RAG query found no answers in any index") return RAGResponse( answer="I couldn't find relevant information to answer your question.", citations=[], @@ -679,6 +934,106 @@ class OptimalService: context_used=0, ) + async def _synthesize_from_citations( + self, + query: str, + citations: list[SearchResult], + ) -> str: + """ + Synthesize an answer from aggregated citations using the local LLM. + + Called when individual indexes returned citations but no LLM answer + (e.g., due to timeouts or errors). Includes retry logic for transient + failures. + """ + import asyncio + + import httpx + + from roboco.config import settings + + if not citations: + return "" + + # Group citations by index type for context + context_parts: list[str] = [] + for citation in citations[:10]: # Limit context size + idx_type = citation.index_type + source_type = idx_type.value if idx_type else "unknown" + context_parts.append(f"[{source_type}] {citation.content}") + + context = "\n\n---\n\n".join(context_parts) + + prompt = ( + f"Based on the following context from the knowledge base, " + f"answer the question concisely.\n\n" + f"Context:\n{context}\n\n" + f"Question: {query}\n\n" + f"Answer:" + ) + + # Retry configuration + max_retries = 3 + retry_delay_base = 0.5 + + async with httpx.AsyncClient(timeout=30.0) as client: + for attempt in range(max_retries): + try: + resp = await client.post( + f"{settings.local_llm_base_url}/chat/completions", + json={ + "model": settings.local_llm_model, + "messages": [{"role": "user", "content": prompt}], + "max_tokens": 1024, + }, + ) + if resp.is_success: + data = resp.json() + answer: str = data["choices"][0]["message"]["content"] + return answer + elif resp.status_code >= httpx.codes.INTERNAL_SERVER_ERROR: + # Server error - retry + logger.warning( + "Synthesis LLM server error, retrying", + status=resp.status_code, + attempt=attempt + 1, + ) + else: + # Client error - don't retry + logger.warning( + "Synthesis LLM call failed", + status=resp.status_code, + ) + break + + except httpx.TimeoutException: + logger.warning( + "Synthesis LLM timeout, retrying", + attempt=attempt + 1, + ) + except httpx.ConnectError as e: + logger.warning( + "Synthesis LLM connection error, retrying", + attempt=attempt + 1, + error=str(e), + ) + except Exception as e: + logger.warning("Synthesis failed", error=str(e)) + break + + # Exponential backoff before retry + if attempt < max_retries - 1: + delay = retry_delay_base * (2**attempt) + await asyncio.sleep(delay) + + # Fallback: return a simple summary of top citations + if citations: + summary = "Based on the knowledge base:\n\n" + for i, c in enumerate(citations[:3], 1): + summary += f"{i}. {c.content[:300]}...\n\n" + return summary + return "" + # ========================================================================= # SPECIALIZED SEARCH (Optimal Brain) # ========================================================================= @@ -861,6 +1216,101 @@ class OptimalService: return stats + async def check_index_staleness( + self, + index_type: IndexType | None = None, + ) -> dict[str, Any]: + """ + Check if indexes are stale (source files modified after last indexing). + + This helps detect when a reindex is needed because files have changed. + + Args: + index_type: Specific index to check, or None for CODE and DOCUMENTATION + + Returns: + Dict with staleness info per index type + """ + from datetime import UTC, datetime + + from sqlalchemy import func, select + + from roboco.db import get_db_context + from roboco.db.tables import IndexedDocumentTable + + result: dict[str, Any] = {"stale_indexes": [], "details": {}} + + # Only check CODE and DOCUMENTATION (file-based indexes) + indexes_to_check = ( + [index_type] if index_type else [IndexType.CODE, IndexType.DOCUMENTATION] + ) + + async with get_db_context() as session: + for idx_type in indexes_to_check: + if idx_type not in self._plugins: + continue + + # Get last indexed time + last_indexed_query = ( + select(func.max(IndexedDocumentTable.indexed_at)) + .select_from(IndexedDocumentTable) + .where(IndexedDocumentTable.index_type == idx_type.value) + ) + last_indexed_result = await session.execute(last_indexed_query) + last_indexed = last_indexed_result.scalar() + + if last_indexed is None: + # Never indexed + result["stale_indexes"].append(idx_type.value) + result["details"][idx_type.value] = { + "status": "never_indexed", + "last_indexed": None, + "recommendation": "Run /kb/reindex to index this content", + } + continue + + # Get indexed source paths + sources_query = ( + select(IndexedDocumentTable.source) + .where(IndexedDocumentTable.index_type == idx_type.value) + .distinct() + ) + sources_result = await session.execute(sources_query) + indexed_sources = [row[0] for row in sources_result.fetchall()] + + # Check if any source files are newer than last_indexed + stale_files: list[str] = [] + for source in indexed_sources[:100]: # Limit check to 100 files + source_path = Path(source) + if source_path.exists(): + try: + mtime = datetime.fromtimestamp( + source_path.stat().st_mtime, tz=UTC + ) + if mtime > last_indexed: + stale_files.append(source) + except OSError: + pass # Skip files we can't stat + + if stale_files: + result["stale_indexes"].append(idx_type.value) + result["details"][idx_type.value] = { + "status": "stale", + "last_indexed": last_indexed.isoformat(), + "stale_file_count": len(stale_files), + "stale_files_sample": stale_files[:5], + "recommendation": "Run /kb/reindex?force=true to update", + } + else: + result["details"][idx_type.value] = { + "status": "current", + "last_indexed": last_indexed.isoformat(), + "indexed_sources_count": len(indexed_sources), + } + + result["needs_reindex"] = len(result["stale_indexes"]) > 0 + return result + async def auto_index_on_startup( self, code_sources: list[str] | None = None, diff --git a/roboco/services/optimal_brain/indexes/base.py b/roboco/services/optimal_brain/indexes/base.py index 6532c501..b297f389 100644 --- a/roboco/services/optimal_brain/indexes/base.py +++ b/roboco/services/optimal_brain/indexes/base.py @@ -12,7 +12,7 @@ from typing import Any, cast import structlog from piragi import AsyncRagi -from piragi.types import Document +from piragi.types import Citation, Document from roboco.config import settings from roboco.models.optimal import IndexType, SearchResult @@ -32,9 +32,9 @@ class IndexConfig: use_hyde: bool = True use_hybrid_search: bool = True use_cross_encoder: bool = False - embedding_model: str = "BAAI/bge-base-en-v1.5" - llm_model: str = "llama3.2" - llm_base_url: str = "http://192.168.50.111:11434/v1" + embedding_model: str = "embeddinggemma:300m" + llm_model: str = "gemma3:4b" + llm_base_url: str = "http://roboco-ollama:11434/v1" @classmethod def from_settings(cls, index_type: IndexType) -> "IndexConfig": @@ -183,6 +183,52 @@ class BaseIndexPlugin(ABC): config["embedding"]["api_key"] = "not-needed" return config + async def _validate_embedding_dimensions(self, embedder: Any) -> None: + """ + Validate that embedder produces expected dimensions. + + Catches dimension mismatches early (e.g., wrong model configured) + instead of failing silently during vector search. + """ + import asyncio + + from roboco.config import settings + + expected_dim = settings.embedding_dimensions + + try: + # Use async method if available + if hasattr(embedder, "aembed_query"): + test_embedding = await embedder.aembed_query("dimension test") + else: + test_embedding = await asyncio.to_thread( + embedder.embed_query, "dimension test" + ) + + actual_dim = len(test_embedding) + + if actual_dim != expected_dim: + raise RuntimeError( + f"Embedding dimension mismatch for {self.index_type.value}: " + f"model produces {actual_dim} dimensions, " + f"but settings.embedding_dimensions={expected_dim}. " + f"Update ROBOCO_EMBEDDING_DIMENSIONS or use correct model." + ) + + logger.debug( + "Embedding dimension validated", + index_type=self.index_type.value, + dimensions=actual_dim, + ) + except RuntimeError: + raise + except Exception as e: + logger.warning( + "Could not validate embedding dimensions", + index_type=self.index_type.value, + error=str(e), + ) + async def initialize(self) -> None: """Initialize the index backend.""" if self._initialized: @@ -200,9 +246,11 @@ class BaseIndexPlugin(ABC): model=self.config.embedding_model, ) + # Validate embedding dimensions match configuration + # This catches mismatches early instead of failing silently during search + await self._validate_embedding_dimensions(shared_embedder) + # Create store with correct vector dimension for embedding model - # Piragi's factory defaults to 768 for PostgresStore, matching - # the embedding model which produces 768-dimensional embeddings store = self._create_store_with_dimension() # Use config with dummy embedding URL to prevent model loading @@ -541,17 +589,27 @@ class BaseIndexPlugin(ABC): import asyncio try: - # Run retrieve in thread pool to avoid blocking event loop - # piragi's retrieve() calls embedder.embed_chunks() synchronously ragi_sync = self.ragi._sync + embedder = ragi_sync.embedder - def _do_retrieve() -> list: - # Embed the query - query_embedding = ragi_sync.embedder.embed_query(query) - # Search the store - return ragi_sync.store.search(query_embedding, top_k=top_k) + # Use async embedding if available (OllamaEmbedder has aembed_query) + if hasattr(embedder, "aembed_query"): + query_embedding = await embedder.aembed_query(query) + else: + # Fallback to sync in thread for SentenceTransformers + query_embedding = await asyncio.to_thread(embedder.embed_query, query) + + # Store search is sync - run in thread + def _do_search() -> list[Citation]: + results: list[Citation] = ragi_sync.store.search( + query_embedding, + top_k=top_k, + min_chunk_length=0, + ) + return results + + chunks = await asyncio.to_thread(_do_search) - chunks = await asyncio.to_thread(_do_retrieve) results = [] for chunk in chunks: # Apply filters if provided @@ -579,6 +637,16 @@ class BaseIndexPlugin(ABC): ) return [] + def _fallback_answer(self, _search_results: list[SearchResult]) -> str: + """ + Return empty to let OptimalService continue searching other indexes. + + Previously this returned a formatted string of search results, but that + caused query() to early-return and skip remaining indexes. Now we return + empty and let the service-level aggregation handle fallback synthesis. + """ + return "" + async def ask( self, query: str, @@ -592,57 +660,103 @@ class BaseIndexPlugin(ABC): top_k: Number of context chunks to use Returns: - Tuple of (answer, citations) + Tuple of (answer, citations). Returns ("", []) on failure + to allow OptimalService to continue to next index. """ + import asyncio + import httpx + # Per-index timeout to prevent one slow index from blocking everything + INDEX_TIMEOUT = 15.0 + + search_results: list[SearchResult] = [] + try: - # First get context using our properly async search - search_results = await self.search(query, top_k=top_k) - - if not search_results: - return "No relevant context found.", [] - - # Build context for LLM - context_texts = [r.content for r in search_results] - context = "\n\n---\n\n".join(context_texts) - - # Build prompt - prompt = ( - f"Based on the following context, answer the question.\n\n" - f"Context:\n{context}\n\n" - f"Question: {query}\n\n" - f"Answer:" - ) - - # Call LLM via Ollama API (async HTTP) - async with httpx.AsyncClient(timeout=60.0) as client: - resp = await client.post( - f"{self.config.llm_base_url}/chat/completions", - json={ - "model": self.config.llm_model, - "messages": [{"role": "user", "content": prompt}], - "max_tokens": 1024, - }, + async with asyncio.timeout(INDEX_TIMEOUT): + # First get context using our properly async search + logger.info( + "ask() starting search", + index_type=self.index_type.value, + query=query[:50], ) - if resp.is_success: - data = resp.json() - answer_text = data["choices"][0]["message"]["content"] - return answer_text, search_results - else: - logger.warning( - "LLM call failed", - status=resp.status_code, - error=resp.text, + search_results = await self.search(query, top_k=top_k) + + logger.info( + "ask() search completed", + index_type=self.index_type.value, + num_results=len(search_results), + ) + + if not search_results: + # Return empty to continue to next index + return "", [] + + # Build context for LLM + context_texts = [r.content for r in search_results] + context = "\n\n---\n\n".join(context_texts) + + # Build prompt + prompt = ( + f"Based on the following context, answer the question.\n\n" + f"Context:\n{context}\n\n" + f"Question: {query}\n\n" + f"Answer:" + ) + + # Call LLM via Ollama API (async HTTP) + llm_url = f"{self.config.llm_base_url}/chat/completions" + logger.info( + "ask() calling LLM", + index_type=self.index_type.value, + llm_url=llm_url, + model=self.config.llm_model, + ) + async with httpx.AsyncClient(timeout=60.0) as client: + resp = await client.post( + llm_url, + json={ + "model": self.config.llm_model, + "messages": [{"role": "user", "content": prompt}], + "max_tokens": 1024, + }, ) - return "", search_results + if resp.is_success: + data = resp.json() + answer_text = data["choices"][0]["message"]["content"] + return answer_text, search_results + else: + logger.warning( + "LLM call failed in ask", + index_type=self.index_type.value, + status=resp.status_code, + error=resp.text[:200] if resp.text else "no error text", + ) + # Return empty to let service aggregate and synthesize + return "", search_results + + except TimeoutError: + logger.warning( + "Index ask() timed out", + index_type=self.index_type.value, + timeout=INDEX_TIMEOUT, + ) + # Return search results even on timeout - service can aggregate them + return "", search_results + except httpx.TimeoutException: + logger.warning( + "LLM HTTP call timed out in ask", + index_type=self.index_type.value, + ) + return "", search_results except Exception as e: logger.warning( "RAG query failed", index_type=self.index_type.value, error=str(e), ) - return "", [] + # Return whatever search results we have for aggregation + return "", search_results async def count(self) -> int: """Get the number of documents in the index.""" diff --git a/roboco/services/optimal_brain/indexes/code.py b/roboco/services/optimal_brain/indexes/code.py index e41b7910..73a373f9 100644 --- a/roboco/services/optimal_brain/indexes/code.py +++ b/roboco/services/optimal_brain/indexes/code.py @@ -4,8 +4,16 @@ Code Index Plugin Handles indexing and searching code files and repositories. Uses a simple line-based chunking strategy instead of sentence-based, since code doesn't have natural sentence boundaries. + +Features: +- Incremental indexing: only re-embeds files that changed (via content hashing) +- Priority-based ordering: indexes models, services, api first +- Line-based chunking: respects code structure """ +import hashlib +import json +import re from pathlib import Path from typing import Any @@ -18,25 +26,147 @@ from roboco.services.optimal_brain.indexes.base import BaseIndexPlugin, IngestRe logger = structlog.get_logger() +class FileHashRegistry: + """ + Tracks file content hashes for incremental indexing. + + Persists hashes to disk so only changed files are re-embedded. + """ + + def __init__(self, cache_file: Path | None = None): + """Initialize with optional cache file path.""" + self._cache_file = cache_file or Path(".piragi/file_hashes.json") + self._hashes: dict[str, str] = {} + self._load() + + def _load(self) -> None: + """Load cached hashes from disk.""" + if self._cache_file.exists(): + try: + self._hashes = json.loads(self._cache_file.read_text()) + logger.debug("Loaded file hash cache", count=len(self._hashes)) + except Exception as e: + logger.warning("Failed to load file hash cache", error=str(e)) + self._hashes = {} + + def _save(self) -> None: + """Persist hashes to disk.""" + try: + self._cache_file.parent.mkdir(parents=True, exist_ok=True) + self._cache_file.write_text(json.dumps(self._hashes, indent=2)) + except Exception as e: + logger.warning("Failed to save file hash cache", error=str(e)) + + @staticmethod + def _hash_content(content: str) -> str: + """Generate hash for file content.""" + return hashlib.md5(content.encode(), usedforsecurity=False).hexdigest() + + def has_changed(self, file_path: str, content: str) -> bool: + """Check if file content has changed since last index.""" + current_hash = self._hash_content(content) + stored_hash = self._hashes.get(file_path) + return stored_hash != current_hash + + def update(self, file_path: str, content: str) -> None: + """Update stored hash for a file.""" + self._hashes[file_path] = self._hash_content(content) + + def update_batch(self, files: list[tuple[str, str]]) -> None: + """Update hashes for multiple files and save.""" + for file_path, content in files: + self._hashes[file_path] = self._hash_content(content) + self._save() + + def remove(self, file_path: str) -> None: + """Remove a file from the registry.""" + self._hashes.pop(file_path, None) + + @property + def count(self) -> int: + """Number of tracked files.""" + return len(self._hashes) + + +# Patterns for detecting natural code boundaries +CODE_BOUNDARY_PATTERNS = [ + re.compile(r"^\s*(def|async def|class)\s+\w+"), # Python + re.compile(r"^\s*(function|const|let|var)\s+\w+"), # JS/TS + re.compile(r"^\s*(func|type|struct)\s+\w+"), # Go + re.compile(r"^\s*(fn|impl|struct|enum)\s+"), # Rust + re.compile(r"^\s*(public|private|protected|static)?\s*(void|int)"), # Java +] + + +def _is_boundary_line(line: str) -> bool: + """Check if line starts a logical code block (function, class, etc).""" + stripped = line.lstrip() + if not stripped: + return False + return any(pattern.match(stripped) for pattern in CODE_BOUNDARY_PATTERNS) + + +def _score_break_point(line: str) -> int: + """Score how good a line is as a chunk break point (higher = better).""" + stripped = line.strip() + if not stripped: + return 3 # Empty line - great break point + if stripped in ("}", "end", "]"): + return 2 # Closing brace/keyword + if stripped.endswith(";") or stripped.endswith(":"): + return 1 # Statement end + return 0 + + +def _find_best_break(lines: list[str], next_line: str) -> int: + """Find the best break point index in a list of lines.""" + # If next line is a boundary (function/class), break at end + if _is_boundary_line(next_line): + return len(lines) + + best_break = len(lines) + best_score = -1 + + # Check last ~5 lines for good break points + start = max(0, len(lines) - 5) + for j in range(start, len(lines)): + score = _score_break_point(lines[j]) + if score > best_score: + best_score = score + best_break = j + 1 + + return best_break + + +def _calc_overlap(lines: list[str], max_overlap: int) -> list[str]: + """Calculate overlap lines from end of a chunk.""" + overlap_lines: list[str] = [] + overlap_size = 0 + for line in reversed(lines): + if overlap_size + len(line) + 1 > max_overlap: + break + overlap_lines.insert(0, line) + overlap_size += len(line) + 1 + return overlap_lines + + def chunk_code( content: str, source: str, chunk_size: int = 1500, - chunk_overlap: int = 200, + chunk_overlap: int = 150, ) -> list[Chunk]: """ - Chunk code using a simple line-based strategy. + Chunk code using a smart boundary-aware strategy. - Unlike prose, code doesn't have sentence boundaries. This chunker: - - Splits by lines (respects code structure) - - Uses character-based sizes (not token-based, simpler and faster) - - Tries to break at blank lines or function/class boundaries + Prefers breaking at function/class boundaries, falls back to blank lines. + Uses reduced overlap since embedding cache avoids redundant work. Args: content: Source code content source: Source file path/URI chunk_size: Target chunk size in characters (~375 tokens) - chunk_overlap: Overlap between chunks in characters + chunk_overlap: Overlap between chunks (reduced from 200) Returns: List of Chunk objects @@ -46,53 +176,59 @@ def chunk_code( lines = content.split("\n") chunks: list[Chunk] = [] - current_chunk_lines: list[str] = [] + current_lines: list[str] = [] current_size = 0 - chunk_index = 0 + min_chunk = 300 # Minimum chunk size for line in lines: - line_size = len(line) + 1 # +1 for newline + line_size = len(line) + 1 - # Check if adding this line would exceed chunk size - if current_size + line_size > chunk_size and current_chunk_lines: - # Save current chunk - chunk_text = "\n".join(current_chunk_lines) + # Check if we need to break + should_break = current_size + line_size > chunk_size and current_lines + + if should_break: + best_break = _find_best_break(current_lines, line) + break_lines = current_lines[:best_break] + break_size = sum(len(ln) + 1 for ln in break_lines) + + if break_size >= min_chunk: + chunks.append( + Chunk( + text="\n".join(break_lines), + source=source, + chunk_index=len(chunks), + metadata={}, + ) + ) + remaining = current_lines[best_break:] + overlap = _calc_overlap(break_lines, chunk_overlap) + current_lines = overlap + remaining + current_size = sum(len(ln) + 1 for ln in current_lines) + + current_lines.append(line) + current_size += line_size + + # Handle remaining content + if current_lines: + chunk_text = "\n".join(current_lines) + if len(chunk_text.strip()) >= min_chunk or not chunks: chunks.append( Chunk( text=chunk_text, source=source, - chunk_index=chunk_index, + chunk_index=len(chunks), metadata={}, ) ) - chunk_index += 1 - - # Calculate overlap: keep last N characters worth of lines - overlap_lines: list[str] = [] - overlap_size = 0 - for prev_line in reversed(current_chunk_lines): - if overlap_size + len(prev_line) + 1 > chunk_overlap: - break - overlap_lines.insert(0, prev_line) - overlap_size += len(prev_line) + 1 - - current_chunk_lines = overlap_lines - current_size = overlap_size - - current_chunk_lines.append(line) - current_size += line_size - - # Don't forget the last chunk - if current_chunk_lines: - chunk_text = "\n".join(current_chunk_lines) - chunks.append( - Chunk( - text=chunk_text, + elif chunks: + # Merge small trailing content into last chunk + last = chunks[-1] + chunks[-1] = Chunk( + text=last.text + "\n" + chunk_text, source=source, - chunk_index=chunk_index, - metadata={}, + chunk_index=last.chunk_index, + metadata=last.metadata, ) - ) return chunks @@ -121,8 +257,9 @@ CODE_EXTENSIONS = { ".toml": "toml", } -# Directories to skip during indexing +# Directories to skip during indexing (non-essential for understanding codebase) SKIP_DIRECTORIES = { + # VCS and build ".git", ".venv", "venv", @@ -132,6 +269,7 @@ SKIP_DIRECTORIES = { ".next", "dist", "build", + # Cache directories ".mypy_cache", ".pytest_cache", ".ruff_cache", @@ -139,8 +277,34 @@ SKIP_DIRECTORIES = { ".tox", "eggs", "*.egg-info", + # Tests and migrations (not critical for auto-indexing) + "tests", + "test", + "alembic", + "migrations", + "seeds", + "fixtures", + "conftest", } +# Priority directories to index first (most important for understanding codebase) +# Order matters: index these subdirectories first +PRIORITY_DIRECTORIES = [ + "agents", + "api", + "db", + "enforcement", + "events", + "llm", + "mcp", + "models", + "runtime", + "services", +] + +# Maximum files to auto-index on startup (prevent timeout) +MAX_AUTO_INDEX_FILES = 100 + class CodeIndexPlugin(BaseIndexPlugin): """ @@ -150,8 +314,23 @@ class CodeIndexPlugin(BaseIndexPlugin): - Source code files (Python, TypeScript, etc.) - Repository directories - Glob patterns for selective indexing + + Features: + - Incremental indexing: skips unchanged files based on content hash + - Priority ordering: indexes models, services, api first + - Batch embedding: processes all files together for efficiency """ + def __init__(self) -> None: + super().__init__() + self._hash_registry: FileHashRegistry | None = None + + def _get_hash_registry(self) -> FileHashRegistry: + """Lazy-load the hash registry.""" + if self._hash_registry is None: + self._hash_registry = FileHashRegistry() + return self._hash_registry + @property def index_type(self) -> IndexType: return IndexType.CODE @@ -175,84 +354,141 @@ class CodeIndexPlugin(BaseIndexPlugin): file_path = kwargs.get("file_path", doc_id or "unknown") return f"roboco://code/{file_path}" + def _is_valid_code_file(self, file_path: Path) -> bool: + """Check if file is a valid code file for indexing.""" + if not file_path.is_file(): + return False + if file_path.suffix not in CODE_EXTENSIONS: + return False + return not any(skip in file_path.parts for skip in SKIP_DIRECTORIES) + + def _collect_files_from_source(self, source: str) -> list[Path]: + """Collect code files from a source path, directory, or glob.""" + source_path = Path(source) + + # Expand glob patterns and directories + if "*" in source: + candidates = list(Path().glob(source)) + elif source_path.is_dir(): + candidates = list(source_path.rglob("*")) + elif source_path.exists(): + candidates = [source_path] + else: + logger.warning(f"Source not found: {source}") + return [] + + return [f for f in candidates if self._is_valid_code_file(f)] + async def index_sources( self, sources: list[str], project: str | None = None, + max_files: int | None = None, + force_reindex: bool = False, ) -> tuple[int, list[dict[str, Any]]]: """ - Index code files/directories with batch embedding for performance. + Index code files/directories with incremental embedding. - Uses batch processing to embed all files together instead of one-by-one, - achieving 10-15x speedup on large codebases. + Uses content hashing to skip unchanged files, dramatically reducing + re-indexing time. Only files with changed content are re-embedded. + + Files are sorted by priority: models, services, api, db, etc. come first. + During auto-indexing, limits to MAX_AUTO_INDEX_FILES to prevent timeouts. Args: sources: List of file paths, directories, or glob patterns project: Optional project identifier for filtering + max_files: Optional limit on files to index (for auto-indexing) + force_reindex: If True, re-index all files regardless of hash Returns: Tuple of (count, indexed_files) where indexed_files contains metadata for each file indexed (for database tracking) """ - # Step 1: Collect all files and their contents - files_data: list[dict[str, Any]] = [] - + # Step 1: Collect all files + all_files: list[Path] = [] for source in sources: - source_path = Path(source) + all_files.extend(self._collect_files_from_source(source)) - # Expand glob patterns and directories - if "*" in source: - files = list(Path().glob(source)) - elif source_path.is_dir(): - files = [ - f - for f in source_path.rglob("*") - if f.suffix in CODE_EXTENSIONS - and not any(skip in f.parts for skip in SKIP_DIRECTORIES) - ] - elif source_path.exists(): - files = [source_path] - else: - logger.warning(f"Source not found: {source}") - continue - - logger.info(f"Found {len(files)} code files to index in {source}") - - for file_path in files: - if not file_path.is_file(): - continue - if file_path.suffix not in CODE_EXTENSIONS: - continue - if any(skip in file_path.parts for skip in SKIP_DIRECTORIES): - continue - - try: - content = file_path.read_text(encoding="utf-8") - language = CODE_EXTENSIONS.get(file_path.suffix) - - files_data.append( - { - "content": content, - "file_path": file_path, - "language": language, - "project": project or "default", - } - ) - except Exception as e: - logger.warning( - "Failed to read code file", - file=str(file_path), - error=str(e), - ) - - if not files_data: + if not all_files: return 0, [] - logger.info(f"Batch processing {len(files_data)} code files") + # Step 2: Sort by priority (models, services, api, etc. first) + def priority_key(path: Path) -> tuple[int, str]: + parts = path.parts + for idx, priority_dir in enumerate(PRIORITY_DIRECTORIES): + if priority_dir in parts: + return (idx, str(path)) + return (len(PRIORITY_DIRECTORIES), str(path)) + + all_files.sort(key=priority_key) + + # Step 3: Apply file limit if specified + if max_files and len(all_files) > max_files: + logger.info( + f"Limiting to {max_files} files (found {len(all_files)})", + priority_dirs=PRIORITY_DIRECTORIES[:4], + ) + all_files = all_files[:max_files] + + logger.info(f"Found {len(all_files)} code files to check") + + # Step 4: Read file contents and filter by hash (incremental indexing) + hash_registry = self._get_hash_registry() + files_data: list[dict[str, Any]] = [] + skipped_count = 0 + + for file_path in all_files: + try: + content = file_path.read_text(encoding="utf-8") + file_key = str(file_path.absolute()) + + # Skip unchanged files unless force_reindex is True + file_changed = hash_registry.has_changed(file_key, content) + if not force_reindex and not file_changed: + skipped_count += 1 + continue + + language = CODE_EXTENSIONS.get(file_path.suffix) + files_data.append( + { + "content": content, + "file_path": file_path, + "language": language, + "project": project or "default", + } + ) + except Exception as e: + logger.warning( + "Failed to read code file", + file=str(file_path), + error=str(e), + ) + + if skipped_count > 0: + logger.info( + f"Incremental indexing: {skipped_count} unchanged files skipped, " + f"{len(files_data)} files to embed" + ) + + if not files_data: + logger.info("No files need re-indexing (all unchanged)") + return 0, [] + + logger.info(f"Batch processing {len(files_data)} changed code files") results = await self._ingest_code_batch(files_data) count = sum(1 for r in results if r.success) + # Update hashes for successfully indexed files + if count > 0: + indexed_hashes: list[tuple[str, str]] = [ + (str(data["file_path"].absolute()), data["content"]) + for data, result in zip(files_data, results, strict=True) + if result.success + ] + hash_registry.update_batch(indexed_hashes) + # Build indexed_files list for database tracking indexed_files = [ { @@ -265,7 +501,10 @@ class CodeIndexPlugin(BaseIndexPlugin): for data in files_data ] - logger.info(f"Batch indexing complete: {count} files indexed") + logger.info( + f"Batch indexing complete: {count} files indexed, " + f"{skipped_count} unchanged files skipped" + ) return count, indexed_files async def _ingest_code_batch( diff --git a/roboco/services/optimal_brain/mentor.py b/roboco/services/optimal_brain/mentor.py index 0babe80c..9ab1a9ac 100644 --- a/roboco/services/optimal_brain/mentor.py +++ b/roboco/services/optimal_brain/mentor.py @@ -11,8 +11,10 @@ import uuid from datetime import UTC, datetime from typing import Any, cast +import httpx import structlog +from roboco.config import settings from roboco.models.optimal import ( IndexType, MentorConversation, @@ -295,15 +297,23 @@ class MentorService: async def _synthesize_answer( self, - _question: str, + question: str, sources: list[SearchResult], - _conversation_context: str, + conversation_context: str, ) -> str: """ - Synthesize an answer from search results. + Synthesize an answer from search results using an LLM. - For now, this uses a simple template. In production, this would - use an LLM to generate a coherent answer. + Uses the local Ollama LLM to generate a coherent answer based on + the retrieved context from the knowledge base. + + Args: + question: The user's question + sources: Retrieved search results from the knowledge base + conversation_context: Previous conversation turns for context + + Returns: + LLM-generated answer synthesized from the sources """ if not sources: return ( @@ -311,48 +321,108 @@ class MentorService: "Try rephrasing your question or asking about a different topic." ) - # Build answer from top sources - answer_parts = [] - - # Group sources by type + # Build context from sources, grouped by type for clarity + context_parts = [] by_type: dict[IndexType, list[SearchResult]] = {} - for source in sources[:5]: + for source in sources[:10]: # Use top 10 sources if source.index_type not in by_type: by_type[source.index_type] = [] by_type[source.index_type].append(source) - # Standards first - if IndexType.STANDARDS in by_type: - answer_parts.append("**Standards & Guidelines:**") - for s in by_type[IndexType.STANDARDS][:2]: - answer_parts.append(f"- {s.content[:200]}...") + # Format each source type + type_labels = { + IndexType.STANDARDS: "Standards & Guidelines", + IndexType.DECISIONS: "Past Decisions", + IndexType.LEARNINGS: "Team Learnings", + IndexType.JOURNALS: "Agent Journals", + IndexType.CODE: "Code References", + IndexType.REVIEWS: "Code Reviews", + IndexType.DOCUMENTATION: "Documentation", + IndexType.ERRORS: "Error Patterns", + IndexType.CONVERSATIONS: "Conversations", + } - # Decisions - if IndexType.DECISIONS in by_type: - answer_parts.append("\n**Past Decisions:**") - for s in by_type[IndexType.DECISIONS][:2]: - answer_parts.append(f"- {s.content[:200]}...") + for index_type, results in by_type.items(): + label = type_labels.get(index_type, index_type.value) + context_parts.append(f"## {label}") + for r in results[:3]: # Max 3 per type + # Include source reference and content + context_parts.append(f"[Source: {r.source}]") + context_parts.append(r.content) + context_parts.append("") # Blank line separator - # Learnings - if IndexType.LEARNINGS in by_type or IndexType.JOURNALS in by_type: - answer_parts.append("\n**Team Learnings:**") - learnings = by_type.get(IndexType.LEARNINGS, []) + by_type.get( - IndexType.JOURNALS, [] + context = "\n".join(context_parts) + + # Build the prompt + system_prompt = ( + "You are a knowledgeable mentor for a software development team. " + "Answer questions based on the provided context from the knowledge base. " + "Be concise but thorough. If the context doesn't fully answer the " + "question, say so and provide what information you can. " + "Reference specific sources when relevant." + ) + + user_prompt_parts = [] + + if conversation_context: + user_prompt_parts.append( + f"Previous conversation:\n{conversation_context}\n" ) - for s in learnings[:2]: - answer_parts.append(f"- {s.content[:200]}...") - # Code patterns - if IndexType.CODE in by_type: - answer_parts.append("\n**Code References:**") - for s in by_type[IndexType.CODE][:2]: - answer_parts.append(f"- See: {s.source}") + user_prompt_parts.append(f"Knowledge base context:\n{context}\n") + user_prompt_parts.append(f"Question: {question}") - if not answer_parts: - first_content = sources[0].content[:500] - return f"Based on my search, here's what I found:\n\n{first_content}" + user_prompt = "\n".join(user_prompt_parts) - return "\n".join(answer_parts) + # Call LLM via Ollama OpenAI-compatible API + try: + async with httpx.AsyncClient(timeout=60.0) as client: + response = await client.post( + f"{settings.local_llm_base_url}/chat/completions", + json={ + "model": settings.local_llm_model, + "messages": [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": user_prompt}, + ], + "max_tokens": 1024, + "temperature": 0.7, + }, + ) + + if response.is_success: + data = response.json() + answer: str = data["choices"][0]["message"]["content"] + return answer + else: + logger.warning( + "LLM call failed in mentor", + status=response.status_code, + error=response.text[:200], + ) + # Fall back to simple extraction + return self._fallback_answer(sources) + + except httpx.TimeoutException: + logger.warning("LLM call timed out in mentor") + return self._fallback_answer(sources) + except Exception as e: + logger.warning("LLM call failed in mentor", error=str(e)) + return self._fallback_answer(sources) + + def _fallback_answer(self, sources: list[SearchResult]) -> str: + """Generate a fallback answer when LLM is unavailable.""" + if not sources: + return "Unable to generate an answer at this time." + + # Simple extraction of top content + parts = ["Based on the knowledge base, here's what I found:\n"] + for source in sources[:3]: + parts.append(f"**From {source.index_type.value}:**") + parts.append(source.content[:300]) + parts.append("") + + return "\n".join(parts) async def get_conversation_history( self, conversation_id: str diff --git a/roboco/services/optimal_brain/ollama_embedder.py b/roboco/services/optimal_brain/ollama_embedder.py new file mode 100644 index 00000000..45e15f9b --- /dev/null +++ b/roboco/services/optimal_brain/ollama_embedder.py @@ -0,0 +1,669 @@ +""" +Ollama Embedder + +Provides embedding generation using Ollama's native API. +Drop-in replacement for piragi's EmbeddingGenerator when using Ollama models. + +Features: +- Parallel batch processing for faster embedding +- Content-based caching to avoid re-embedding +- Connection pooling for efficiency +- Retry logic for transient failures +""" + +import asyncio +import hashlib +import time +from collections.abc import Callable +from typing import Any + +import httpx +from piragi.types import Chunk + +from roboco.config import settings +from roboco.logging import get_logger + +logger = get_logger(__name__) + +# Retry configuration +MAX_RETRIES = 3 +RETRY_DELAY_BASE = 0.5 # seconds, exponential backoff + +# Parallel processing configuration +MAX_CONCURRENT_BATCHES = 4 # Number of batches to process in parallel +DEFAULT_BATCH_SIZE = 16 # Smaller batches for CPU-friendly embedding + + +class OllamaEmbedderError(Exception): + """Base exception for Ollama embedder errors.""" + + pass + + +class OllamaConnectionError(OllamaEmbedderError): + """Raised when Ollama server is unreachable.""" + + pass + + +class OllamaModelError(OllamaEmbedderError): + """Raised when the embedding model is unavailable or returns invalid data.""" + + pass + + +class EmbeddingCache: + """ + LRU cache for embeddings keyed by content hash. + + Avoids re-computing embeddings for identical content. + """ + + def __init__(self, max_size: int = 10000): + self._cache: dict[str, list[float]] = {} + self._access_order: list[str] = [] + self._max_size = max_size + self._hits = 0 + self._misses = 0 + + @staticmethod + def _hash_content(content: str) -> str: + """Generate hash for content.""" + return hashlib.md5(content.encode(), usedforsecurity=False).hexdigest() + + def get(self, content: str) -> list[float] | None: + """Get cached embedding by content.""" + key = self._hash_content(content) + if key in self._cache: + self._hits += 1 + # Move to end (most recently used) + self._access_order.remove(key) + self._access_order.append(key) + return self._cache[key] + self._misses += 1 + return None + + def put(self, content: str, embedding: list[float]) -> None: + """Cache embedding for content.""" + key = self._hash_content(content) + if key in self._cache: + return # Already cached + + # Evict oldest if at capacity + while len(self._cache) >= self._max_size: + oldest = self._access_order.pop(0) + del self._cache[oldest] + + self._cache[key] = embedding + self._access_order.append(key) + + def get_many(self, contents: list[str]) -> tuple[list[int], list[list[float]]]: + """ + Get cached embeddings for multiple contents. + + Returns: + Tuple of (indices of cached items, their embeddings) + """ + cached_indices = [] + cached_embeddings = [] + for i, content in enumerate(contents): + emb = self.get(content) + if emb is not None: + cached_indices.append(i) + cached_embeddings.append(emb) + return cached_indices, cached_embeddings + + def put_many(self, contents: list[str], embeddings: list[list[float]]) -> None: + """Cache multiple embeddings.""" + for content, emb in zip(contents, embeddings, strict=True): + self.put(content, emb) + + @property + def stats(self) -> dict[str, Any]: + """Get cache statistics.""" + total = self._hits + self._misses + hit_rate = (self._hits / total * 100) if total > 0 else 0 + return { + "size": len(self._cache), + "max_size": self._max_size, + "hits": self._hits, + "misses": self._misses, + "hit_rate": f"{hit_rate:.1f}%", + } + + +class OllamaEmbedder: + """ + Embedding generator using Ollama's native API. + + Features: + - Parallel batch processing (configurable concurrency) + - Content-based caching (avoids re-embedding identical content) + - Connection pooling for efficiency + - Retry logic with exponential backoff + """ + + def __init__( + self, + model: str = "embeddinggemma:300m", + base_url: str | None = None, + timeout: float = 120.0, + max_concurrent: int = MAX_CONCURRENT_BATCHES, + cache_size: int = 10000, + ): + """Initialize Ollama embedder. + + Args: + model: Ollama model name for embeddings + base_url: Ollama API base URL (default: from settings) + timeout: Request timeout in seconds (default 120s for CPU embedding) + max_concurrent: Max concurrent batch requests (default 4) + cache_size: Max cached embeddings (default 10000) + """ + self.model = model + self.base_url = base_url or settings.ollama_base_url + self.timeout = timeout + self.max_concurrent = max_concurrent + self._dimensions: int | None = None + self._cache = EmbeddingCache(max_size=cache_size) + # Reusable clients for connection pooling + self._sync_client: httpx.Client | None = None + self._async_client: httpx.AsyncClient | None = None + # Semaphore for limiting concurrent requests + self._semaphore: asyncio.Semaphore | None = None + + def _get_sync_client(self) -> httpx.Client: + """Get or create sync HTTP client with connection pooling.""" + if self._sync_client is None or self._sync_client.is_closed: + timeout = httpx.Timeout( + connect=10.0, + read=self.timeout, + write=30.0, + pool=10.0, + ) + self._sync_client = httpx.Client( + timeout=timeout, + limits=httpx.Limits( + max_connections=self.max_concurrent * 2, + max_keepalive_connections=self.max_concurrent, + ), + ) + return self._sync_client + + async def _get_async_client(self) -> httpx.AsyncClient: + """Get or create async HTTP client with connection pooling.""" + if self._async_client is None or self._async_client.is_closed: + timeout = httpx.Timeout( + connect=10.0, + read=self.timeout, + write=30.0, + pool=10.0, + ) + self._async_client = httpx.AsyncClient( + timeout=timeout, + limits=httpx.Limits( + max_connections=self.max_concurrent * 2, + max_keepalive_connections=self.max_concurrent, + ), + ) + return self._async_client + + def _get_semaphore(self) -> asyncio.Semaphore: + """Get or create semaphore for limiting concurrent requests.""" + if self._semaphore is None: + self._semaphore = asyncio.Semaphore(self.max_concurrent) + return self._semaphore + + def close(self) -> None: + """Close HTTP clients and release resources.""" + if self._sync_client and not self._sync_client.is_closed: + self._sync_client.close() + self._sync_client = None + + async def aclose(self) -> None: + """Async close HTTP clients and release resources.""" + if self._async_client and not self._async_client.is_closed: + await self._async_client.aclose() + self._async_client = None + self.close() + + @property + def dimensions(self) -> int: + """Get embedding dimensions (cached after first call).""" + if self._dimensions is None: + test_embedding = self.embed_query("test") + self._dimensions = len(test_embedding) + return self._dimensions + + def set_dimensions(self, dim: int) -> None: + """Pre-set dimensions to avoid blocking call.""" + self._dimensions = dim + + async def get_dimensions_async(self) -> int: + """Async-friendly way to get embedding dimensions.""" + if self._dimensions is None: + test_embedding = await self.aembed_query("test") + self._dimensions = len(test_embedding) + return self._dimensions + + @property + def cache_stats(self) -> dict[str, Any]: + """Get embedding cache statistics.""" + return self._cache.stats + + def _handle_embed_response( + self, response: httpx.Response, input_count: int = 1 + ) -> list[list[float]]: + """Validate and extract embeddings from API response.""" + if not response.is_success: + error_text = response.text[:200] if response.text else "Unknown error" + if response.status_code == httpx.codes.NOT_FOUND: + raise OllamaModelError( + f"Model '{self.model}' not found. " + f"Run 'ollama pull {self.model}' to download it." + ) + raise OllamaEmbedderError( + f"Ollama API error {response.status_code}: {error_text}" + ) + + try: + data = response.json() + except Exception as e: + raise OllamaEmbedderError(f"Invalid JSON response: {e}") from e + + embeddings: list[list[float]] | None = data.get("embeddings") + if not embeddings: + raise OllamaModelError( + f"No embeddings returned for model '{self.model}'. " + "The model may not support embeddings." + ) + + if len(embeddings) != input_count: + raise OllamaModelError( + f"Expected {input_count} embeddings, got {len(embeddings)}" + ) + + for i, emb in enumerate(embeddings): + if not emb or not isinstance(emb, list): + raise OllamaModelError(f"Invalid embedding at index {i}") + + return embeddings + + def embed_query( + self, + query: str, + task_instruction: str | None = None, + ) -> list[float]: + """Generate embedding for a single query with retry logic.""" + _ = task_instruction + + # Check cache first + cached = self._cache.get(query) + if cached is not None: + return cached + + client = self._get_sync_client() + last_error: Exception | None = None + + for attempt in range(MAX_RETRIES): + try: + response = client.post( + f"{self.base_url}/api/embed", + json={"model": self.model, "input": query}, + ) + embeddings = self._handle_embed_response(response, input_count=1) + result = embeddings[0] + self._cache.put(query, result) + return result + + except httpx.ConnectError as e: + last_error = OllamaConnectionError( + f"Cannot connect to Ollama at {self.base_url}: {e}" + ) + except httpx.TimeoutException as e: + last_error = OllamaConnectionError(f"Ollama request timed out: {e}") + except (OllamaModelError, OllamaEmbedderError): + raise + except Exception as e: + last_error = OllamaEmbedderError(f"Unexpected error: {e}") + + if attempt < MAX_RETRIES - 1: + delay = RETRY_DELAY_BASE * (2**attempt) + logger.warning( + "Ollama embed_query retry", + attempt=attempt + 1, + delay=delay, + error=str(last_error), + ) + time.sleep(delay) + + raise last_error or OllamaEmbedderError("Max retries exceeded") + + def _embed_batch_sync( + self, client: httpx.Client, batch: list[str], batch_index: int + ) -> list[list[float]]: + """Embed a single batch synchronously with retry logic.""" + last_error: Exception | None = None + + for attempt in range(MAX_RETRIES): + try: + response = client.post( + f"{self.base_url}/api/embed", + json={"model": self.model, "input": batch}, + ) + return self._handle_embed_response(response, input_count=len(batch)) + + except httpx.ConnectError as e: + last_error = OllamaConnectionError( + f"Cannot connect to Ollama at {self.base_url}: {e}" + ) + except httpx.TimeoutException as e: + last_error = OllamaConnectionError(f"Ollama request timed out: {e}") + except (OllamaModelError, OllamaEmbedderError): + raise + except Exception as e: + last_error = OllamaEmbedderError(f"Unexpected error: {e}") + + if attempt < MAX_RETRIES - 1: + delay = RETRY_DELAY_BASE * (2**attempt) + logger.warning( + "Ollama embed_documents retry", + attempt=attempt + 1, + batch_index=batch_index, + delay=delay, + error=str(last_error), + ) + time.sleep(delay) + + raise last_error or OllamaEmbedderError("Max retries exceeded") + + def embed_documents( + self, + documents: list[str], + task_instruction: str | None = None, + batch_size: int = DEFAULT_BATCH_SIZE, + ) -> list[list[float]]: + """Generate embeddings for multiple documents (sequential, uses cache).""" + _ = task_instruction + + if not documents: + return [] + + # Check cache for all documents + result_embeddings: list[list[float] | None] = [None] * len(documents) + uncached_indices: list[int] = [] + uncached_docs: list[str] = [] + + for i, doc in enumerate(documents): + cached = self._cache.get(doc) + if cached is not None: + result_embeddings[i] = cached + else: + uncached_indices.append(i) + uncached_docs.append(doc) + + if uncached_indices: + logger.info( + "Embedding cache stats", + cached=len(documents) - len(uncached_indices), + uncached=len(uncached_indices), + total=len(documents), + ) + + if not uncached_docs: + return [e for e in result_embeddings if e is not None] + + # Embed uncached documents in batches + client = self._get_sync_client() + new_embeddings: list[list[float]] = [] + + for i in range(0, len(uncached_docs), batch_size): + batch = uncached_docs[i : i + batch_size] + embeddings = self._embed_batch_sync(client, batch, batch_index=i) + new_embeddings.extend(embeddings) + # Cache new embeddings + for doc, emb in zip(batch, embeddings, strict=True): + self._cache.put(doc, emb) + + # Merge cached and new embeddings + for idx, emb in zip(uncached_indices, new_embeddings, strict=True): + result_embeddings[idx] = emb + + return [e for e in result_embeddings if e is not None] + + async def _embed_batch_async( + self, + batch: list[str], + batch_index: int, + ) -> list[list[float]]: + """Embed a single batch with semaphore-limited concurrency.""" + semaphore = self._get_semaphore() + client = await self._get_async_client() + + async with semaphore: + last_error: Exception | None = None + + for attempt in range(MAX_RETRIES): + try: + logger.debug( + "Parallel embed batch", + batch_index=batch_index, + batch_size=len(batch), + attempt=attempt, + ) + response = await client.post( + f"{self.base_url}/api/embed", + json={"model": self.model, "input": batch}, + ) + return self._handle_embed_response(response, input_count=len(batch)) + + except httpx.ConnectError as e: + last_error = OllamaConnectionError( + f"Cannot connect to Ollama at {self.base_url}: {e}" + ) + except httpx.TimeoutException as e: + last_error = OllamaConnectionError(f"Ollama request timed out: {e}") + except (OllamaModelError, OllamaEmbedderError): + raise + except Exception as e: + last_error = OllamaEmbedderError(f"Unexpected error: {e}") + + if attempt < MAX_RETRIES - 1: + delay = RETRY_DELAY_BASE * (2**attempt) + logger.warning( + "Parallel embed batch retry", + batch_index=batch_index, + attempt=attempt + 1, + delay=delay, + error=str(last_error), + ) + await asyncio.sleep(delay) + + raise last_error or OllamaEmbedderError("Max retries exceeded") + + async def aembed_documents_parallel( + self, + documents: list[str], + batch_size: int = DEFAULT_BATCH_SIZE, + ) -> list[list[float]]: + """ + Embed documents in parallel batches for maximum throughput. + + Processes multiple batches concurrently (limited by max_concurrent). + Uses caching to skip already-embedded content. + """ + if not documents: + return [] + + # Check cache for all documents + result_embeddings: list[list[float] | None] = [None] * len(documents) + uncached_indices: list[int] = [] + uncached_docs: list[str] = [] + + for i, doc in enumerate(documents): + cached = self._cache.get(doc) + if cached is not None: + result_embeddings[i] = cached + else: + uncached_indices.append(i) + uncached_docs.append(doc) + + cache_hits = len(documents) - len(uncached_indices) + if cache_hits > 0: + logger.info( + "Embedding cache hits", + cached=cache_hits, + uncached=len(uncached_indices), + total=len(documents), + hit_rate=f"{cache_hits / len(documents) * 100:.1f}%", + ) + + if not uncached_docs: + return [e for e in result_embeddings if e is not None] + + # Split uncached docs into batches + batches: list[list[str]] = [] + for i in range(0, len(uncached_docs), batch_size): + batches.append(uncached_docs[i : i + batch_size]) + + logger.info( + "Parallel embedding starting", + total_docs=len(uncached_docs), + batches=len(batches), + batch_size=batch_size, + max_concurrent=self.max_concurrent, + ) + + start_time = time.time() + + # Process all batches in parallel (semaphore limits concurrency) + tasks = [self._embed_batch_async(batch, i) for i, batch in enumerate(batches)] + batch_results = await asyncio.gather(*tasks) + + # Flatten results and cache + new_embeddings: list[list[float]] = [] + doc_idx = 0 + for batch, embeddings in zip(batches, batch_results, strict=True): + for doc, emb in zip(batch, embeddings, strict=True): + new_embeddings.append(emb) + self._cache.put(doc, emb) + doc_idx += len(batch) + + elapsed = time.time() - start_time + docs_per_sec = len(uncached_docs) / elapsed if elapsed > 0 else 0 + + logger.info( + "Parallel embedding complete", + docs=len(uncached_docs), + elapsed=f"{elapsed:.1f}s", + docs_per_sec=f"{docs_per_sec:.1f}", + ) + + # Merge cached and new embeddings + for idx, emb in zip(uncached_indices, new_embeddings, strict=True): + result_embeddings[idx] = emb + + return [e for e in result_embeddings if e is not None] + + def embed_chunks( + self, + chunks: list[Chunk], + on_progress: Callable[[str], None] | None = None, + ) -> list[Chunk]: + """Generate embeddings for chunks using parallel processing.""" + if not chunks: + return chunks + + texts = [chunk.text for chunk in chunks] + + # Use async parallel embedding via event loop + try: + loop = asyncio.get_event_loop() + if loop.is_running(): + # Already in async context - use sync fallback + embeddings = self.embed_documents(texts, batch_size=DEFAULT_BATCH_SIZE) + else: + embeddings = loop.run_until_complete( + self.aembed_documents_parallel(texts) + ) + except RuntimeError: + # No event loop - create one + embeddings = asyncio.run(self.aembed_documents_parallel(texts)) + + for chunk, embedding in zip(chunks, embeddings, strict=True): + chunk.embedding = embedding + + if on_progress: + on_progress(f"Embedded {len(chunks)} chunks") + + return chunks + + async def aembed_chunks( + self, + chunks: list[Chunk], + on_progress: Callable[[str], None] | None = None, + ) -> list[Chunk]: + """Async version of embed_chunks with parallel processing.""" + if not chunks: + return chunks + + texts = [chunk.text for chunk in chunks] + embeddings = await self.aembed_documents_parallel(texts) + + for chunk, embedding in zip(chunks, embeddings, strict=True): + chunk.embedding = embedding + + if on_progress: + on_progress(f"Embedded {len(chunks)} chunks") + + return chunks + + async def aembed_query(self, query: str) -> list[float]: + """Async version of embed_query with retry logic and caching.""" + # Check cache first + cached = self._cache.get(query) + if cached is not None: + return cached + + client = await self._get_async_client() + last_error: Exception | None = None + + for attempt in range(MAX_RETRIES): + try: + response = await client.post( + f"{self.base_url}/api/embed", + json={"model": self.model, "input": query}, + ) + embeddings = self._handle_embed_response(response, input_count=1) + result = embeddings[0] + self._cache.put(query, result) + return result + + except httpx.ConnectError as e: + last_error = OllamaConnectionError( + f"Cannot connect to Ollama at {self.base_url}: {e}" + ) + except httpx.TimeoutException as e: + last_error = OllamaConnectionError(f"Ollama request timed out: {e}") + except (OllamaModelError, OllamaEmbedderError): + raise + except Exception as e: + last_error = OllamaEmbedderError(f"Unexpected error: {e}") + + if attempt < MAX_RETRIES - 1: + delay = RETRY_DELAY_BASE * (2**attempt) + logger.warning( + "Ollama aembed_query retry", + attempt=attempt + 1, + delay=delay, + error=str(last_error), + ) + await asyncio.sleep(delay) + + raise last_error or OllamaEmbedderError("Max retries exceeded") + + async def aembed_documents( + self, documents: list[str], batch_size: int = DEFAULT_BATCH_SIZE + ) -> list[list[float]]: + """Async embed_documents - delegates to parallel implementation.""" + return await self.aembed_documents_parallel(documents, batch_size=batch_size) diff --git a/roboco/services/optimal_brain/shared_embedder.py b/roboco/services/optimal_brain/shared_embedder.py index c8cca40e..1e4e6926 100644 --- a/roboco/services/optimal_brain/shared_embedder.py +++ b/roboco/services/optimal_brain/shared_embedder.py @@ -1,25 +1,68 @@ """ Shared Embedder Singleton -Provides a single EmbeddingGenerator instance shared across all index plugins -to avoid loading the SentenceTransformer model 9 times (~3s each = 27s startup). +Provides a single embedder instance shared across all index plugins. +Supports both Ollama models (embeddinggemma, etc.) and SentenceTransformers (BGE, etc.). """ import asyncio -from typing import TYPE_CHECKING +from collections.abc import Callable +from typing import TYPE_CHECKING, Protocol, Union +from piragi.types import Chunk + +from roboco.config import settings from roboco.logging import get_logger if TYPE_CHECKING: from piragi.embeddings import EmbeddingGenerator + from roboco.services.optimal_brain.ollama_embedder import OllamaEmbedder + logger = get_logger(__name__) +class Embedder(Protocol): + """Protocol for embedder interface.""" + + def embed_query( + self, query: str, task_instruction: str | None = None + ) -> list[float]: ... + + def embed_documents( + self, + documents: list[str], + task_instruction: str | None = None, + batch_size: int = 32, + ) -> list[list[float]]: ... + + def embed_chunks( + self, + chunks: list[Chunk], + on_progress: Callable[[str], None] | None = None, + ) -> list[Chunk]: ... + + +# Known Ollama embedding models +OLLAMA_EMBEDDING_MODELS = { + "embeddinggemma", + "nomic-embed-text", + "mxbai-embed-large", + "all-minilm", + "snowflake-arctic-embed", +} + + +def _is_ollama_model(model: str) -> bool: + """Check if model name is an Ollama embedding model.""" + model_base = model.split(":")[0].lower() + return model_base in OLLAMA_EMBEDDING_MODELS + + class _SharedEmbedderHolder: """Holder class for shared embedder state (avoids global statement).""" - instance: "EmbeddingGenerator | None" = None + instance: Union["EmbeddingGenerator", "OllamaEmbedder", None] = None lock: asyncio.Lock | None = None @classmethod @@ -31,21 +74,22 @@ class _SharedEmbedderHolder: async def get_shared_embedder( - model: str = "BAAI/bge-base-en-v1.5", + model: str | None = None, device: str | None = None, timeout: float = 60.0, -) -> "EmbeddingGenerator": +) -> Embedder: """Get or create the shared embedder instance. Thread-safe singleton that loads the model only once. + Automatically selects Ollama or SentenceTransformers based on model name. Args: - model: Embedding model name (default: BAAI/bge-base-en-v1.5) - device: Device to use (None = auto-detect) + model: Embedding model name (default: from settings) + device: Device to use for SentenceTransformers (None = auto-detect) timeout: Max seconds to wait for model loading (default: 60) Returns: - Shared EmbeddingGenerator instance + Shared embedder instance (OllamaEmbedder or EmbeddingGenerator) Raises: TimeoutError: If model loading takes too long @@ -55,50 +99,97 @@ async def get_shared_embedder( return _SharedEmbedderHolder.instance async with _SharedEmbedderHolder.get_lock(): - # Double-check after acquiring lock - if _SharedEmbedderHolder.instance is not None: - return _SharedEmbedderHolder.instance + # Double-check after acquiring lock (another coroutine may have created it) + if _SharedEmbedderHolder.instance is None: + model = model or settings.default_embedding_model - logger.info( - "Creating shared embedder (one-time load)", - model=model, - device=device or "auto", - ) - - # Import here to avoid circular imports and defer heavy import - from piragi.embeddings import EmbeddingGenerator - - # Run model loading in thread to not block event loop - def _create_embedder() -> "EmbeddingGenerator": - return EmbeddingGenerator( - model=model, - device=device, - batch_size=32, - ) - - try: - async with asyncio.timeout(timeout): - _SharedEmbedderHolder.instance = await asyncio.to_thread( - _create_embedder + # Use Ollama for Ollama models, SentenceTransformers otherwise + if _is_ollama_model(model): + logger.info( + "Creating shared Ollama embedder", + model=model, + base_url=settings.ollama_base_url, ) - except TimeoutError: - logger.error( - "Embedder initialization timed out", - model=model, - timeout=timeout, - ) - raise TimeoutError( - f"Embedding model loading timed out after {timeout}s. " - "This may indicate network issues or corrupted model cache." - ) from None - except Exception as e: - logger.error("Embedder initialization failed", model=model, error=str(e)) - raise RuntimeError(f"Failed to load embedding model: {e}") from e - logger.info("Shared embedder created successfully") + from roboco.services.optimal_brain.ollama_embedder import OllamaEmbedder + + _SharedEmbedderHolder.instance = OllamaEmbedder( + model=model, + base_url=settings.ollama_base_url, + ) + else: + logger.info( + "Creating shared SentenceTransformers embedder", + model=model, + device=device or "auto", + ) + + # Import here to avoid circular imports and defer heavy import + from piragi.embeddings import EmbeddingGenerator + + # Run model loading in thread to not block event loop + def _create_embedder() -> "EmbeddingGenerator": + return EmbeddingGenerator( + model=model, + device=device, + batch_size=32, + ) + + try: + async with asyncio.timeout(timeout): + _SharedEmbedderHolder.instance = await asyncio.to_thread( + _create_embedder + ) + except TimeoutError: + logger.error( + "Embedder initialization timed out", + model=model, + timeout=timeout, + ) + raise TimeoutError( + f"Embedding model loading timed out after {timeout}s. " + "This may indicate network issues or corrupted model cache." + ) from None + except Exception as e: + logger.error( + "Embedder initialization failed", model=model, error=str(e) + ) + raise RuntimeError(f"Failed to load embedding model: {e}") from e + + # Validate embedder implements required protocol methods + assert _SharedEmbedderHolder.instance is not None + _validate_embedder_protocol(_SharedEmbedderHolder.instance, model) + + logger.info("Shared embedder created successfully", model=model) + + # At this point instance is guaranteed to be set (or exception raised) + assert _SharedEmbedderHolder.instance is not None return _SharedEmbedderHolder.instance +def _validate_embedder_protocol(embedder: Embedder, model: str) -> None: + """ + Validate that embedder implements the required protocol methods. + + Checks at creation time rather than failing during first use. + + Args: + embedder: The embedder instance to validate + model: Model name for error messages + + Raises: + RuntimeError: If embedder is missing required methods + """ + required_methods = ["embed_query", "embed_documents", "embed_chunks"] + missing = [m for m in required_methods if not callable(getattr(embedder, m, None))] + + if missing: + raise RuntimeError( + f"Embedder for model '{model}' is missing required methods: {missing}. " + f"Embedder type: {type(embedder).__name__}" + ) + + async def close_shared_embedder() -> None: """Release the shared embedder resources.""" async with _SharedEmbedderHolder.get_lock(): diff --git a/roboco/services/task.py b/roboco/services/task.py index 4b62949a..de1e7625 100644 --- a/roboco/services/task.py +++ b/roboco/services/task.py @@ -207,6 +207,7 @@ class TaskService(BaseService): parent_task_id=req.parent_task_id, target_date=req.target_date, estimated_complexity=req.estimated_complexity, + nature=req.nature, status=req.status if req.status else TaskStatus.PENDING, sequence=req.sequence, # Task ordering within siblings dependency_ids=req.dependency_ids, # Task IDs that must complete first