mirror of
https://github.com/rennf93/roboco.git
synced 2026-08-03 07:23:24 +02:00
Task sequence and RAG LLM improvements
This commit is contained in:
@@ -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
|
||||
|
||||
+8
-10
@@ -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
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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)
|
||||
|
||||
+5
-3
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
# =============================================================================
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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),
|
||||
|
||||
+28
-2
@@ -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()
|
||||
|
||||
|
||||
+11
-8
@@ -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)",
|
||||
)
|
||||
|
||||
# ==========================================================================
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)",
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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."""
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
+544
-94
@@ -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,
|
||||
|
||||
@@ -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."""
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
@@ -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():
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user