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
|
### Hardware Infrastructure
|
||||||
|
|
||||||
- **Olares One (Powerhouse)**: Intel Ultra 9 + RTX 5090, runs Claude Code instances and AI inference
|
- **Olares One (Powerhouse)**: Intel Ultra 9 + RTX 5090, runs Claude Code instances and AI inference - NOT YET ARRIVED
|
||||||
- **UGREEN NAS (Warehouse)**: 36TB RAID6, hosts PostgreSQL, Redis
|
- **UGREEN NAS (Warehouse)**: 36TB RAID6, 128GB RAM, hosts PostgreSQL, Redis
|
||||||
- **Pi Cluster (Operations)**: Monitoring, notifications, smart home
|
- **Pi Cluster (Operations)**: Monitoring, notifications, smart home
|
||||||
|
|
||||||
## Development Standards
|
## Development Standards
|
||||||
@@ -326,7 +326,6 @@ ROBOCO_RAG_USE_HYDE=true
|
|||||||
ROBOCO_RAG_USE_HYBRID_SEARCH=true
|
ROBOCO_RAG_USE_HYBRID_SEARCH=true
|
||||||
|
|
||||||
# AI/LLM
|
# AI/LLM
|
||||||
ROBOCO_DEFAULT_LLM_MODEL=claude-opus-4-5-20251101
|
|
||||||
ROBOCO_DEFAULT_EMBEDDING_MODEL=embeddinggemma:300m
|
ROBOCO_DEFAULT_EMBEDDING_MODEL=embeddinggemma:300m
|
||||||
ROBOCO_LOCAL_LLM_MODEL=gemma3:4b
|
ROBOCO_LOCAL_LLM_MODEL=gemma3:4b
|
||||||
ROBOCO_LOCAL_LLM_BASE_URL=http://roboco-ollama:11434/v1
|
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
|
## Knowledge Base Tools
|
||||||
|
|
||||||
- `roboco_kb_search(query)` - Search code, docs, decisions
|
- `roboco_ask_mentor(question)` - **Primary tool** - AI answers with follow-up support
|
||||||
- `roboco_rag_query(question)` - AI-generated answers
|
- `roboco_kb_search(query)` - Raw semantic search
|
||||||
- `roboco_kb_stats()` - See what's indexed
|
- `roboco_search_error(error_message)` - Find known error solutions
|
||||||
- `roboco_search_error(pattern)` - Find error solutions
|
|
||||||
- `roboco_check_decision(topic)` - Find past decisions
|
- `roboco_check_decision(topic)` - Find past decisions
|
||||||
- `roboco_search_learnings(topic)` - Find team learnings
|
- `roboco_search_learnings(query)` - Find team learnings
|
||||||
|
|
||||||
## Journaling (ALL agents)
|
## Journaling (ALL agents)
|
||||||
|
|
||||||
@@ -147,8 +146,7 @@ Need docs updated? Create a task for your cell's Documenter.
|
|||||||
|
|
||||||
## RAG Checkpoints
|
## RAG Checkpoints
|
||||||
|
|
||||||
Before critical actions, verify with RAG:
|
Before critical actions, check the knowledge base:
|
||||||
- **Full workflow example**: `roboco_kb_search("{your_role} workflow")`
|
- `roboco_ask_mentor("How do I implement X?")` - Best practices, patterns
|
||||||
- **Tool parameters**: `roboco_kb_search("mcp tools")`
|
- `roboco_search_error(pattern)` - Known error solutions
|
||||||
- **Error solutions**: `roboco_search_error(pattern)`
|
- `roboco_check_decision(topic)` - Past architectural decisions
|
||||||
- **Past decisions**: `roboco_check_decision(topic)`
|
|
||||||
|
|||||||
@@ -258,6 +258,12 @@ def upgrade() -> None:
|
|||||||
server_default="code",
|
server_default="code",
|
||||||
),
|
),
|
||||||
sa.Column("requires_git", sa.Boolean(), nullable=False, server_default="true"),
|
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(
|
sa.Column(
|
||||||
"project_id",
|
"project_id",
|
||||||
postgresql.UUID(as_uuid=True),
|
postgresql.UUID(as_uuid=True),
|
||||||
@@ -941,6 +947,7 @@ def downgrade() -> None:
|
|||||||
# Drop enums
|
# Drop enums
|
||||||
op.execute("DROP TYPE IF EXISTS worksessionstatus")
|
op.execute("DROP TYPE IF EXISTS worksessionstatus")
|
||||||
op.execute("DROP TYPE IF EXISTS tasktype")
|
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 handoffstatus")
|
||||||
op.execute("DROP TYPE IF EXISTS journalentrytype")
|
op.execute("DROP TYPE IF EXISTS journalentrytype")
|
||||||
op.execute("DROP TYPE IF EXISTS notificationpriority")
|
op.execute("DROP TYPE IF EXISTS notificationpriority")
|
||||||
|
|||||||
@@ -38,6 +38,55 @@ services:
|
|||||||
timeout: 5s
|
timeout: 5s
|
||||||
retries: 5
|
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)
|
# Agent Base Image Builder (specialized images built on-demand by orchestrator)
|
||||||
# ==========================================================================
|
# ==========================================================================
|
||||||
@@ -166,6 +215,11 @@ services:
|
|||||||
# API
|
# API
|
||||||
ROBOCO_HOST: 0.0.0.0
|
ROBOCO_HOST: 0.0.0.0
|
||||||
ROBOCO_PORT: 8000
|
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)
|
# Host paths for spawning agent containers (required for Docker-in-Docker)
|
||||||
# IMPORTANT: These must be ABSOLUTE paths on the host filesystem
|
# IMPORTANT: These must be ABSOLUTE paths on the host filesystem
|
||||||
ROBOCO_HOST_PROJECT_DIR: ${ROBOCO_HOST_PROJECT_DIR:-/volume1/roboco}
|
ROBOCO_HOST_PROJECT_DIR: ${ROBOCO_HOST_PROJECT_DIR:-/volume1/roboco}
|
||||||
@@ -185,6 +239,10 @@ services:
|
|||||||
condition: service_healthy
|
condition: service_healthy
|
||||||
redis:
|
redis:
|
||||||
condition: service_healthy
|
condition: service_healthy
|
||||||
|
ollama:
|
||||||
|
condition: service_healthy
|
||||||
|
ollama-init:
|
||||||
|
condition: service_completed_successfully
|
||||||
agent-base-image:
|
agent-base-image:
|
||||||
condition: service_completed_successfully
|
condition: service_completed_successfully
|
||||||
# Default agents to spawn (override in .env or command line)
|
# 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.transcription = _AppServices.transcription
|
||||||
app.state.extraction = _AppServices.extraction
|
app.state.extraction = _AppServices.extraction
|
||||||
|
|
||||||
# Initialize Phase 3 services (RAG - optional, non-blocking)
|
# Initialize OptimalService (RAG) - BLOCKS until fully ready
|
||||||
# OptimalService.initialize() auto-indexes /docs/standards/ and /docs/workflows/
|
# This ensures /health only returns 200 when RAG is operational
|
||||||
|
# Typical initialization time: 30-90 seconds (embedding + indexing)
|
||||||
try:
|
try:
|
||||||
|
logger.info("Initializing OptimalService (RAG)...")
|
||||||
optimal_service = await get_optimal_service()
|
optimal_service = await get_optimal_service()
|
||||||
app.state.optimal = optimal_service
|
app.state.optimal = optimal_service
|
||||||
logger.info("OptimalService (RAG) initialized successfully")
|
logger.info("OptimalService (RAG) initialized successfully")
|
||||||
@@ -96,7 +98,7 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None]:
|
|||||||
)
|
)
|
||||||
app.state.optimal = None
|
app.state.optimal = None
|
||||||
|
|
||||||
logger.info("All services initialized")
|
logger.info("All services initialized, API ready")
|
||||||
|
|
||||||
yield
|
yield
|
||||||
|
|
||||||
|
|||||||
@@ -309,6 +309,18 @@ async def rag_query(
|
|||||||
status_code=status.HTTP_504_GATEWAY_TIMEOUT,
|
status_code=status.HTTP_504_GATEWAY_TIMEOUT,
|
||||||
detail=f"RAG query timed out after {rag_timeout}s",
|
detail=f"RAG query timed out after {rag_timeout}s",
|
||||||
) from e
|
) 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(
|
return RAGQueryResponse(
|
||||||
answer=response.answer,
|
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)
|
@router.get("/health", response_model=RAGHealthResponse)
|
||||||
async def rag_health_check() -> RAGHealthResponse:
|
async def rag_health_check() -> RAGHealthResponse:
|
||||||
"""
|
"""
|
||||||
@@ -481,7 +513,13 @@ async def rag_health_check() -> RAGHealthResponse:
|
|||||||
try:
|
try:
|
||||||
async with asyncio.timeout(health_timeout):
|
async with asyncio.timeout(health_timeout):
|
||||||
embedder = await get_shared_embedder(model=settings.default_embedding_model)
|
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:
|
if test_embedding and len(test_embedding) == settings.embedding_dimensions:
|
||||||
embedding_ok = True
|
embedding_ok = True
|
||||||
details["embedding_model"] = settings.default_embedding_model
|
details["embedding_model"] = settings.default_embedding_model
|
||||||
@@ -673,6 +711,7 @@ async def reindex_all(
|
|||||||
agent: CurrentAgentContext,
|
agent: CurrentAgentContext,
|
||||||
permissions: PermissionServiceDep,
|
permissions: PermissionServiceDep,
|
||||||
force: bool = False,
|
force: bool = False,
|
||||||
|
timeout_seconds: int = 300, # 5 minute default
|
||||||
) -> dict[str, Any]:
|
) -> dict[str, Any]:
|
||||||
"""
|
"""
|
||||||
Trigger re-indexing of code and documentation.
|
Trigger re-indexing of code and documentation.
|
||||||
@@ -683,19 +722,31 @@ async def reindex_all(
|
|||||||
|
|
||||||
Args:
|
Args:
|
||||||
force: If True, reindex even if indexes aren't empty
|
force: If True, reindex even if indexes aren't empty
|
||||||
|
timeout_seconds: Maximum time to wait for reindexing (default: 300s)
|
||||||
|
|
||||||
Returns:
|
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):
|
if not permissions.can_perform_kb_action(agent, KBAction.INDEX_CODE):
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code=status.HTTP_403_FORBIDDEN,
|
status_code=status.HTTP_403_FORBIDDEN,
|
||||||
detail="Not authorized to trigger reindexing",
|
detail="Not authorized to trigger reindexing",
|
||||||
)
|
)
|
||||||
|
|
||||||
service = await get_optimal_service()
|
try:
|
||||||
result = await service.auto_index_on_startup(force=force)
|
async with asyncio.timeout(timeout_seconds):
|
||||||
return {"status": "reindexed", **result}
|
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,
|
assigned_to=data.assigned_to,
|
||||||
target_date=data.target_date,
|
target_date=data.target_date,
|
||||||
estimated_complexity=data.estimated_complexity,
|
estimated_complexity=data.estimated_complexity,
|
||||||
|
nature=data.nature,
|
||||||
status=data.status,
|
status=data.status,
|
||||||
sequence=data.sequence, # Task ordering within siblings
|
sequence=data.sequence, # Task ordering within siblings
|
||||||
dependency_ids=data.dependency_ids, # Dependencies for claim filtering
|
dependency_ids=data.dependency_ids, # Dependencies for claim filtering
|
||||||
|
|||||||
@@ -10,7 +10,7 @@ from uuid import UUID
|
|||||||
|
|
||||||
from pydantic import BaseModel, Field
|
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.models.session import SessionScope
|
||||||
from roboco.utils.converters import require_uuid, to_python_uuid, to_python_uuid_list
|
from roboco.utils.converters import require_uuid, to_python_uuid, to_python_uuid_list
|
||||||
|
|
||||||
@@ -225,6 +225,19 @@ class TaskResponse(BaseModel):
|
|||||||
status: TaskStatus
|
status: TaskStatus
|
||||||
priority: int
|
priority: int
|
||||||
sequence: int # Order number within siblings
|
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
|
# Ownership
|
||||||
team: Team
|
team: Team
|
||||||
@@ -291,6 +304,7 @@ class TaskSummaryResponse(BaseModel):
|
|||||||
created_at: datetime
|
created_at: datetime
|
||||||
updated_at: datetime | None
|
updated_at: datetime | None
|
||||||
estimated_complexity: Complexity
|
estimated_complexity: Complexity
|
||||||
|
nature: TaskNature
|
||||||
|
|
||||||
class Config:
|
class Config:
|
||||||
from_attributes = True
|
from_attributes = True
|
||||||
@@ -530,6 +544,17 @@ def task_to_response(task: "TaskTable") -> TaskResponse:
|
|||||||
status=task.status,
|
status=task.status,
|
||||||
priority=task.priority,
|
priority=task.priority,
|
||||||
sequence=task.sequence,
|
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,
|
team=task.team,
|
||||||
created_by=require_uuid(task.created_by),
|
created_by=require_uuid(task.created_by),
|
||||||
assigned_to=to_python_uuid(task.assigned_to),
|
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
|
import asyncio
|
||||||
|
from http import HTTPStatus
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
|
import httpx
|
||||||
import structlog
|
import structlog
|
||||||
import uvicorn
|
import uvicorn
|
||||||
|
|
||||||
@@ -42,6 +44,29 @@ async def _run_api_server() -> None:
|
|||||||
await server.serve()
|
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(
|
async def main(
|
||||||
skip_db: bool = False,
|
skip_db: bool = False,
|
||||||
skip_orchestrator: bool = False,
|
skip_orchestrator: bool = False,
|
||||||
@@ -100,8 +125,9 @@ async def main(
|
|||||||
api_task = asyncio.create_task(_run_api_server())
|
api_task = asyncio.create_task(_run_api_server())
|
||||||
logger.info("API server starting", host=settings.host, port=settings.port)
|
logger.info("API server starting", host=settings.host, port=settings.port)
|
||||||
|
|
||||||
# Wait a moment for API to start before orchestrator begins polling
|
# Wait for API to actually be ready (not just a fixed sleep)
|
||||||
await asyncio.sleep(2)
|
# The lifespan does document indexing which can take 30+ seconds
|
||||||
|
await _wait_for_api_ready()
|
||||||
|
|
||||||
await orchestrator.start()
|
await orchestrator.start()
|
||||||
|
|
||||||
|
|||||||
+11
-8
@@ -165,24 +165,27 @@ class Settings(BaseSettings):
|
|||||||
openai_api_key: str | None = None # For embeddings
|
openai_api_key: str | None = None # For embeddings
|
||||||
|
|
||||||
# Default models
|
# Default models
|
||||||
default_llm_model: str = "claude-opus-4-5-20251101"
|
|
||||||
default_embedding_model: str = Field(
|
default_embedding_model: str = Field(
|
||||||
default="BAAI/bge-base-en-v1.5",
|
default="embeddinggemma:300m",
|
||||||
description="Embedding model",
|
description="Embedding model (Ollama model name)",
|
||||||
)
|
)
|
||||||
embedding_dimensions: int = Field(
|
embedding_dimensions: int = Field(
|
||||||
default=768,
|
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 for RAG (HyDE, reranking, etc.)
|
||||||
local_llm_model: str = Field(
|
local_llm_model: str = Field(
|
||||||
default="qwen3:8b",
|
default="gemma3:4b",
|
||||||
description="Local LLM model for HyDE and RAG operations",
|
description="Local LLM for HyDE/RAG (non-thinking models are faster)",
|
||||||
)
|
)
|
||||||
local_llm_base_url: str = Field(
|
local_llm_base_url: str = Field(
|
||||||
default="http://192.168.50.111:11434/v1",
|
default="http://roboco-ollama:11434/v1",
|
||||||
description="Base URL for local LLM (Ollama)",
|
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,
|
NotificationPriority,
|
||||||
NotificationType,
|
NotificationType,
|
||||||
SessionStatus,
|
SessionStatus,
|
||||||
|
TaskNature,
|
||||||
TaskStatus,
|
TaskStatus,
|
||||||
TaskType,
|
TaskType,
|
||||||
Team,
|
Team,
|
||||||
@@ -137,6 +138,9 @@ class TaskTable(Base):
|
|||||||
task_type: Mapped[TaskType] = mapped_column(
|
task_type: Mapped[TaskType] = mapped_column(
|
||||||
Enum(TaskType), nullable=False, default=TaskType.CODE
|
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)
|
requires_git: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True)
|
||||||
|
|
||||||
# Project & Branch (set by PM during setup)
|
# Project & Branch (set by PM during setup)
|
||||||
|
|||||||
@@ -88,15 +88,22 @@ def _register_search_tools(mcp: FastMCP, client: ApiClient) -> None:
|
|||||||
"SEARCH_FAILED",
|
"SEARCH_FAILED",
|
||||||
"Failed to search knowledge base",
|
"Failed to search knowledge base",
|
||||||
{"api_error": resp.text},
|
{"api_error": resp.text},
|
||||||
|
hint="Try roboco_ask_mentor(question) for AI-synthesized answers.",
|
||||||
)
|
)
|
||||||
|
|
||||||
result = resp.json()
|
result = resp.json()
|
||||||
return {
|
total = result.get("total", 0)
|
||||||
|
response: dict[str, Any] = {
|
||||||
"status": "success",
|
"status": "success",
|
||||||
"query": query,
|
"query": query,
|
||||||
"total": result.get("total", 0),
|
"total": total,
|
||||||
"results": result.get("results", []),
|
"results": result.get("results", []),
|
||||||
}
|
}
|
||||||
|
if total == 0:
|
||||||
|
response["hint"] = (
|
||||||
|
"No results. Try roboco_ask_mentor(question) for better answers."
|
||||||
|
)
|
||||||
|
return response
|
||||||
|
|
||||||
@mcp.tool()
|
@mcp.tool()
|
||||||
async def roboco_rag_query(
|
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.
|
RAG query - get an AI-generated answer using knowledge base context.
|
||||||
|
|
||||||
Use this when you need an answer synthesized from the knowledge base,
|
NOTE: For most questions, prefer `roboco_ask_mentor` instead!
|
||||||
not just search results. Good for questions like:
|
The mentor searches multiple indexes and supports follow-up questions.
|
||||||
- "How does authentication work in this codebase?"
|
|
||||||
- "What's the pattern for error handling?"
|
Use this simpler tool only when you need:
|
||||||
- "What decisions were made about the database schema?"
|
- A quick answer from a specific index type
|
||||||
|
- To filter by project or task_id
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
query: Natural language question
|
query: Natural language question
|
||||||
@@ -132,22 +140,33 @@ def _register_search_tools(mcp: FastMCP, client: ApiClient) -> None:
|
|||||||
if task_id:
|
if task_id:
|
||||||
payload["task_id"] = 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:
|
if not resp.ok:
|
||||||
return format_error_response(
|
return format_error_response(
|
||||||
"RAG_FAILED",
|
"RAG_FAILED",
|
||||||
"Failed to query RAG",
|
"Failed to query RAG",
|
||||||
{"api_error": resp.text},
|
{"api_error": resp.text},
|
||||||
|
hint="Try roboco_ask_mentor(question) instead - it's more robust.",
|
||||||
)
|
)
|
||||||
|
|
||||||
result = resp.json()
|
result = resp.json()
|
||||||
return {
|
answer = result.get("answer", "")
|
||||||
|
context_used = result.get("context_used", 0)
|
||||||
|
response: dict[str, Any] = {
|
||||||
"status": "success",
|
"status": "success",
|
||||||
"query": query,
|
"query": query,
|
||||||
"answer": result.get("answer", ""),
|
"answer": answer,
|
||||||
"citations": result.get("citations", []),
|
"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()
|
@mcp.tool()
|
||||||
async def roboco_kb_stats() -> dict[str, Any]:
|
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.
|
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
|
This is a conversational interface - you can ask follow-up questions
|
||||||
by providing the conversation_id from a previous response.
|
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
|
- Past architectural decisions
|
||||||
- Team learnings and reflections
|
- Team learnings and reflections
|
||||||
- Codebase patterns
|
- Codebase patterns
|
||||||
|
- Known error solutions
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
question: Your question (natural language)
|
question: Your question (natural language)
|
||||||
@@ -369,7 +392,8 @@ def _register_mentor_tools(mcp: FastMCP, client: ApiClient) -> None:
|
|||||||
if domain:
|
if domain:
|
||||||
payload["domain"] = 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:
|
if not resp.ok:
|
||||||
return format_error_response(
|
return format_error_response(
|
||||||
"MENTOR_FAILED",
|
"MENTOR_FAILED",
|
||||||
@@ -421,12 +445,20 @@ def _register_error_tools(mcp: FastMCP, client: ApiClient) -> None:
|
|||||||
)
|
)
|
||||||
|
|
||||||
result = resp.json()
|
result = resp.json()
|
||||||
return {
|
solutions_found = len(result.get("results", []))
|
||||||
|
response: dict[str, Any] = {
|
||||||
"status": "success",
|
"status": "success",
|
||||||
"error_message": error_message,
|
"error_message": error_message,
|
||||||
"solutions_found": len(result.get("results", [])),
|
"solutions_found": solutions_found,
|
||||||
"results": result.get("results", []),
|
"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()
|
@mcp.tool()
|
||||||
async def roboco_record_error_solution(
|
async def roboco_record_error_solution(
|
||||||
@@ -806,12 +838,20 @@ def _register_learning_tools(mcp: FastMCP, client: ApiClient) -> None:
|
|||||||
)
|
)
|
||||||
|
|
||||||
result = resp.json()
|
result = resp.json()
|
||||||
return {
|
total = result.get("total", 0)
|
||||||
|
response: dict[str, Any] = {
|
||||||
"status": "success",
|
"status": "success",
|
||||||
"query": query,
|
"query": query,
|
||||||
"total": result.get("total", 0),
|
"total": total,
|
||||||
"results": result.get("results", []),
|
"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:
|
def _register_index_management_tools(mcp: FastMCP, client: ApiClient) -> None:
|
||||||
|
|||||||
@@ -184,6 +184,9 @@ class TaskCreateInput(BaseModel):
|
|||||||
complexity: str = Field(
|
complexity: str = Field(
|
||||||
default="medium", description="Complexity: low, medium, high, critical"
|
default="medium", description="Complexity: low, medium, high, critical"
|
||||||
)
|
)
|
||||||
|
nature: str = Field(
|
||||||
|
default="technical", description="Task nature: technical, non_technical"
|
||||||
|
)
|
||||||
status: str = Field(
|
status: str = Field(
|
||||||
default="backlog",
|
default="backlog",
|
||||||
description="Status: 'backlog' (default) or 'pending' (ready for work)",
|
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:
|
data: TaskCreateInput with:
|
||||||
- title, description, acceptance_criteria, team (required)
|
- title, description, acceptance_criteria, team (required)
|
||||||
- parent_task_id, assigned_to, priority, status (optional)
|
- parent_task_id, assigned_to, priority, status (optional)
|
||||||
|
- nature: Task nature (technical, non_technical)
|
||||||
- sequence: Order within siblings (0 = default)
|
- sequence: Order within siblings (0 = default)
|
||||||
- dependency_ids: Task IDs that must complete first
|
- 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,
|
"team": input_data.team,
|
||||||
"priority": input_data.priority,
|
"priority": input_data.priority,
|
||||||
"estimated_complexity": input_data.complexity,
|
"estimated_complexity": input_data.complexity,
|
||||||
|
"nature": input_data.nature,
|
||||||
"status": input_data.status, # Always included, defaults to "backlog"
|
"status": input_data.status, # Always included, defaults to "backlog"
|
||||||
"sequence": input_data.sequence, # Task ordering (lower = first)
|
"sequence": input_data.sequence, # Task ordering (lower = first)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -26,6 +26,7 @@ from roboco.models.base import (
|
|||||||
SessionStatus,
|
SessionStatus,
|
||||||
SubstituteReason,
|
SubstituteReason,
|
||||||
# Enums
|
# Enums
|
||||||
|
TaskNature,
|
||||||
TaskStatus,
|
TaskStatus,
|
||||||
Team,
|
Team,
|
||||||
)
|
)
|
||||||
@@ -143,6 +144,7 @@ __all__ = [
|
|||||||
"SubstituteReason",
|
"SubstituteReason",
|
||||||
"Task",
|
"Task",
|
||||||
"TaskCreate",
|
"TaskCreate",
|
||||||
|
"TaskNature",
|
||||||
"TaskPlan",
|
"TaskPlan",
|
||||||
"TaskStatus",
|
"TaskStatus",
|
||||||
"TaskUpdate",
|
"TaskUpdate",
|
||||||
|
|||||||
@@ -47,6 +47,13 @@ class TaskType(str, Enum):
|
|||||||
ADMINISTRATIVE = "administrative" # No git
|
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):
|
class Complexity(str, Enum):
|
||||||
"""Task complexity levels."""
|
"""Task complexity levels."""
|
||||||
|
|
||||||
|
|||||||
@@ -15,6 +15,7 @@ from pydantic import Field
|
|||||||
from roboco.models.base import (
|
from roboco.models.base import (
|
||||||
Complexity,
|
Complexity,
|
||||||
RobocoBase,
|
RobocoBase,
|
||||||
|
TaskNature,
|
||||||
TaskStatus,
|
TaskStatus,
|
||||||
TaskType,
|
TaskType,
|
||||||
Team,
|
Team,
|
||||||
@@ -158,6 +159,9 @@ class Task(TimestampMixin):
|
|||||||
task_type: TaskType = Field(
|
task_type: TaskType = Field(
|
||||||
default=TaskType.CODE, description="Type of task (code, research, etc.)"
|
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(
|
requires_git: bool = Field(
|
||||||
default=True, description="Whether this task requires git workflow"
|
default=True, description="Whether this task requires git workflow"
|
||||||
)
|
)
|
||||||
@@ -290,6 +294,7 @@ class TaskCreate(RobocoBase):
|
|||||||
|
|
||||||
# Git configuration
|
# Git configuration
|
||||||
task_type: TaskType = TaskType.CODE
|
task_type: TaskType = TaskType.CODE
|
||||||
|
nature: TaskNature = TaskNature.TECHNICAL
|
||||||
requires_git: bool = True
|
requires_git: bool = True
|
||||||
project_id: UUID | None = None
|
project_id: UUID | None = None
|
||||||
|
|
||||||
@@ -311,6 +316,7 @@ class TaskUpdate(RobocoBase):
|
|||||||
|
|
||||||
# Git fields
|
# Git fields
|
||||||
task_type: TaskType | None = None
|
task_type: TaskType | None = None
|
||||||
|
nature: TaskNature | None = None
|
||||||
requires_git: bool | None = None
|
requires_git: bool | None = None
|
||||||
project_id: UUID | None = None
|
project_id: UUID | None = None
|
||||||
branch_name: str | None = None
|
branch_name: str | None = None
|
||||||
@@ -348,5 +354,6 @@ class TaskCreateRequest:
|
|||||||
|
|
||||||
# Git configuration
|
# Git configuration
|
||||||
task_type: TaskType = field(default=TaskType.CODE)
|
task_type: TaskType = field(default=TaskType.CODE)
|
||||||
|
nature: TaskNature = field(default=TaskNature.TECHNICAL)
|
||||||
requires_git: bool = True
|
requires_git: bool = True
|
||||||
project_id: UUID | None = None
|
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.
|
by a specialized plugin that implements the BaseIndexPlugin interface.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
from dataclasses import dataclass, field
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
@@ -48,6 +49,70 @@ from roboco.services.optimal_brain.indexes.reviews import (
|
|||||||
logger = structlog.get_logger()
|
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 mapping IndexType to plugin class
|
||||||
PLUGIN_REGISTRY: dict[IndexType, type[BaseIndexPlugin]] = {
|
PLUGIN_REGISTRY: dict[IndexType, type[BaseIndexPlugin]] = {
|
||||||
IndexType.CODE: CodeIndexPlugin,
|
IndexType.CODE: CodeIndexPlugin,
|
||||||
@@ -87,28 +152,112 @@ class OptimalService:
|
|||||||
self._initialized = False
|
self._initialized = False
|
||||||
self._plugins: dict[IndexType, BaseIndexPlugin] = {}
|
self._plugins: dict[IndexType, BaseIndexPlugin] = {}
|
||||||
self._prompt_templates: dict[str, dict[str, Any]] = {}
|
self._prompt_templates: dict[str, dict[str, Any]] = {}
|
||||||
|
self._indexing_task: Any = None # Background indexing task
|
||||||
|
|
||||||
async def initialize(self) -> None:
|
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:
|
if self._initialized:
|
||||||
return
|
return
|
||||||
|
|
||||||
logger.info("Initializing OptimalService with plugin architecture")
|
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
|
# Create and initialize plugins for each index type
|
||||||
for index_type, plugin_class in PLUGIN_REGISTRY.items():
|
for index_type, plugin_class in PLUGIN_REGISTRY.items():
|
||||||
plugin = plugin_class()
|
try:
|
||||||
await plugin.initialize()
|
plugin = plugin_class()
|
||||||
self._plugins[index_type] = plugin
|
async with asyncio.timeout(plugin_init_timeout):
|
||||||
logger.info(f"Initialized {index_type.value} plugin")
|
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
|
# Service is initialized if at least one plugin succeeded
|
||||||
logger.info("OptimalService initialization complete")
|
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
|
# Report failed plugins for debugging
|
||||||
await self._auto_index_on_startup()
|
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.
|
Auto-index code and documentation on startup.
|
||||||
|
|
||||||
@@ -117,32 +266,55 @@ class OptimalService:
|
|||||||
- /docs/standards/ - Coding, security, workflow standards
|
- /docs/standards/ - Coding, security, workflow standards
|
||||||
- /docs/workflows/ - Agent workflow documentation
|
- /docs/workflows/ - Agent workflow documentation
|
||||||
|
|
||||||
This ensures agents can search for code, standards, and workflows
|
Args:
|
||||||
immediately after startup.
|
force: If True, reindex even if indexes already have content
|
||||||
"""
|
|
||||||
await self._auto_index_code()
|
|
||||||
await self._auto_index_docs()
|
|
||||||
|
|
||||||
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."""
|
"""Auto-index source code files on startup."""
|
||||||
# Find the roboco source directory (the Python package, not the repo root)
|
import time
|
||||||
# We want to index roboco/ package, NOT the entire repo (which has .venv)
|
|
||||||
|
report = IndexingReport(index_type="code")
|
||||||
|
start_time = time.time()
|
||||||
|
|
||||||
|
# Find the roboco source directory
|
||||||
possible_code_roots = [
|
possible_code_roots = [
|
||||||
Path("/app/roboco"), # Docker: /app is repo root, roboco/ is package
|
Path("/app/roboco"), # Docker
|
||||||
Path(__file__).parent.parent, # Local: optimal.py -> services -> roboco
|
Path(__file__).parent.parent, # Local
|
||||||
Path.cwd() / "roboco", # Local: cwd/roboco
|
Path.cwd() / "roboco",
|
||||||
]
|
]
|
||||||
|
|
||||||
code_root = None
|
code_root = None
|
||||||
for path in possible_code_roots:
|
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"
|
init_file = path / "__init__.py"
|
||||||
if path.exists() and path.is_dir() and init_file.exists():
|
if path.exists() and path.is_dir() and init_file.exists():
|
||||||
code_root = path
|
code_root = path
|
||||||
logger.debug(
|
|
||||||
"Found code package directory",
|
|
||||||
path=str(path),
|
|
||||||
)
|
|
||||||
break
|
break
|
||||||
|
|
||||||
if code_root is None:
|
if code_root is None:
|
||||||
@@ -150,44 +322,57 @@ class OptimalService:
|
|||||||
"Code directory not found",
|
"Code directory not found",
|
||||||
searched_paths=[str(p) for p in possible_code_roots],
|
searched_paths=[str(p) for p in possible_code_roots],
|
||||||
)
|
)
|
||||||
return
|
return None
|
||||||
|
|
||||||
# Check if code index is empty
|
# Check if code index is empty (unless force=True)
|
||||||
code_plugin = self._get_plugin(IndexType.CODE)
|
if not force:
|
||||||
code_count = await code_plugin.count()
|
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("Auto-indexing source code", directory=str(code_root))
|
||||||
logger.info(
|
|
||||||
"Code index already populated",
|
|
||||||
chunk_count=code_count,
|
|
||||||
skipping=True,
|
|
||||||
)
|
|
||||||
return
|
|
||||||
|
|
||||||
logger.info(
|
|
||||||
"Auto-indexing source code",
|
|
||||||
directory=str(code_root),
|
|
||||||
)
|
|
||||||
|
|
||||||
try:
|
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)
|
logger.info("Code auto-indexing complete", files_indexed=count)
|
||||||
except Exception as e:
|
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:
|
report.duration_seconds = time.time() - start_time
|
||||||
"""
|
return report
|
||||||
Auto-index documentation directories on startup.
|
|
||||||
|
|
||||||
Indexes:
|
async def _auto_index_docs(self, force: bool = False) -> IndexingReport | None:
|
||||||
- /docs/standards/ - Coding, security, workflow standards
|
"""Auto-index documentation directories on startup."""
|
||||||
- /docs/workflows/ - Agent workflow documentation
|
import time
|
||||||
"""
|
|
||||||
# Find the docs directory relative to the project root
|
report = IndexingReport(index_type="documentation")
|
||||||
|
start_time = time.time()
|
||||||
|
|
||||||
|
# Find the docs directory
|
||||||
possible_docs_roots = [
|
possible_docs_roots = [
|
||||||
Path("/app/docs"), # Docker absolute path
|
Path("/app/docs"),
|
||||||
Path(__file__).parent.parent.parent / "docs", # roboco/docs (local)
|
Path(__file__).parent.parent.parent / "docs",
|
||||||
Path.cwd() / "docs", # Current working directory
|
Path.cwd() / "docs",
|
||||||
]
|
]
|
||||||
|
|
||||||
docs_root = None
|
docs_root = None
|
||||||
@@ -201,61 +386,75 @@ class OptimalService:
|
|||||||
"Docs directory not found",
|
"Docs directory not found",
|
||||||
searched_paths=[str(p) for p in possible_docs_roots],
|
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"]
|
auto_index_dirs = ["standards", "workflows"]
|
||||||
|
|
||||||
for subdir in auto_index_dirs:
|
for subdir in auto_index_dirs:
|
||||||
target_dir = docs_root / subdir
|
target_dir = docs_root / subdir
|
||||||
if not target_dir.exists():
|
if not target_dir.exists():
|
||||||
logger.debug(f"Auto-index directory not found: {target_dir}")
|
|
||||||
continue
|
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."""
|
"""Index all markdown files in a documentation directory."""
|
||||||
|
report = IndexingReport(index_type=f"docs/{name}")
|
||||||
|
|
||||||
md_files = list(directory.rglob("*.md"))
|
md_files = list(directory.rglob("*.md"))
|
||||||
if not md_files:
|
if not md_files:
|
||||||
logger.info(f"No files found to index in {name}/", path=str(directory))
|
logger.info(f"No files found to index in {name}/", path=str(directory))
|
||||||
return
|
return report
|
||||||
|
|
||||||
|
report.total_attempted = len(md_files)
|
||||||
logger.info(
|
logger.info(
|
||||||
f"Auto-indexing {name} files",
|
f"Auto-indexing {name} files",
|
||||||
directory=str(directory),
|
directory=str(directory),
|
||||||
file_count=len(md_files),
|
file_count=len(md_files),
|
||||||
)
|
)
|
||||||
|
|
||||||
# Index each file
|
# Index each file with individual error tracking
|
||||||
total_indexed = 0
|
|
||||||
for md_file in md_files:
|
for md_file in md_files:
|
||||||
try:
|
try:
|
||||||
if name == "standards":
|
if name == "standards":
|
||||||
count = await self.index_standards_file(str(md_file))
|
await self.index_standards_file(str(md_file))
|
||||||
total_indexed += count
|
report.successful += 1
|
||||||
else:
|
else:
|
||||||
# For workflows, index as documentation
|
|
||||||
await self.index_documentation([str(md_file)])
|
await self.index_documentation([str(md_file)])
|
||||||
total_indexed += 1
|
report.successful += 1
|
||||||
|
|
||||||
logger.debug(
|
logger.debug(f"Indexed {name} file", file=str(md_file))
|
||||||
f"Indexed {name} file",
|
|
||||||
file=str(md_file),
|
|
||||||
items_indexed=count if name == "standards" else 1,
|
|
||||||
)
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
|
error_msg = str(e)
|
||||||
|
report.failed += 1
|
||||||
|
report.failed_sources.append((str(md_file), error_msg))
|
||||||
logger.warning(
|
logger.warning(
|
||||||
f"Failed to index {name} file",
|
f"Failed to index {name} file",
|
||||||
file=str(md_file),
|
file=str(md_file),
|
||||||
error=str(e),
|
error=error_msg,
|
||||||
)
|
)
|
||||||
|
|
||||||
logger.info(
|
logger.info(
|
||||||
f"{name.capitalize()} auto-indexing complete",
|
f"{name.capitalize()} auto-indexing complete",
|
||||||
files_processed=len(md_files),
|
successful=report.successful,
|
||||||
total_indexed=total_indexed,
|
failed=report.failed,
|
||||||
|
total=report.total_attempted,
|
||||||
)
|
)
|
||||||
|
return report
|
||||||
|
|
||||||
async def close(self) -> None:
|
async def close(self) -> None:
|
||||||
"""Cleanup resources."""
|
"""Cleanup resources."""
|
||||||
@@ -286,11 +485,20 @@ class OptimalService:
|
|||||||
self,
|
self,
|
||||||
sources: list[str],
|
sources: list[str],
|
||||||
project: str | None = None,
|
project: str | None = None,
|
||||||
|
max_files: int | None = None,
|
||||||
) -> int:
|
) -> 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)
|
plugin = self._get_plugin(IndexType.CODE)
|
||||||
if isinstance(plugin, CodeIndexPlugin):
|
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
|
# Batch track all indexed files using repository
|
||||||
docs_to_track = [
|
docs_to_track = [
|
||||||
@@ -644,7 +852,8 @@ class OptimalService:
|
|||||||
"""
|
"""
|
||||||
Query the knowledge base with RAG.
|
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:
|
if not self._initialized:
|
||||||
raise RuntimeError("OptimalService not initialized")
|
raise RuntimeError("OptimalService not initialized")
|
||||||
@@ -653,25 +862,71 @@ class OptimalService:
|
|||||||
context.index_types if context and context.index_types else list(IndexType)
|
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:
|
for index_type in index_types:
|
||||||
plugin = self._plugins.get(index_type)
|
plugin = self._plugins.get(index_type)
|
||||||
if plugin:
|
if not plugin:
|
||||||
try:
|
continue
|
||||||
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),
|
|
||||||
)
|
|
||||||
|
|
||||||
|
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(
|
return RAGResponse(
|
||||||
answer="I couldn't find relevant information to answer your question.",
|
answer="I couldn't find relevant information to answer your question.",
|
||||||
citations=[],
|
citations=[],
|
||||||
@@ -679,6 +934,106 @@ class OptimalService:
|
|||||||
context_used=0,
|
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)
|
# SPECIALIZED SEARCH (Optimal Brain)
|
||||||
# =========================================================================
|
# =========================================================================
|
||||||
@@ -861,6 +1216,101 @@ class OptimalService:
|
|||||||
|
|
||||||
return stats
|
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(
|
async def auto_index_on_startup(
|
||||||
self,
|
self,
|
||||||
code_sources: list[str] | None = None,
|
code_sources: list[str] | None = None,
|
||||||
|
|||||||
@@ -12,7 +12,7 @@ from typing import Any, cast
|
|||||||
|
|
||||||
import structlog
|
import structlog
|
||||||
from piragi import AsyncRagi
|
from piragi import AsyncRagi
|
||||||
from piragi.types import Document
|
from piragi.types import Citation, Document
|
||||||
|
|
||||||
from roboco.config import settings
|
from roboco.config import settings
|
||||||
from roboco.models.optimal import IndexType, SearchResult
|
from roboco.models.optimal import IndexType, SearchResult
|
||||||
@@ -32,9 +32,9 @@ class IndexConfig:
|
|||||||
use_hyde: bool = True
|
use_hyde: bool = True
|
||||||
use_hybrid_search: bool = True
|
use_hybrid_search: bool = True
|
||||||
use_cross_encoder: bool = False
|
use_cross_encoder: bool = False
|
||||||
embedding_model: str = "BAAI/bge-base-en-v1.5"
|
embedding_model: str = "embeddinggemma:300m"
|
||||||
llm_model: str = "llama3.2"
|
llm_model: str = "gemma3:4b"
|
||||||
llm_base_url: str = "http://192.168.50.111:11434/v1"
|
llm_base_url: str = "http://roboco-ollama:11434/v1"
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def from_settings(cls, index_type: IndexType) -> "IndexConfig":
|
def from_settings(cls, index_type: IndexType) -> "IndexConfig":
|
||||||
@@ -183,6 +183,52 @@ class BaseIndexPlugin(ABC):
|
|||||||
config["embedding"]["api_key"] = "not-needed"
|
config["embedding"]["api_key"] = "not-needed"
|
||||||
return config
|
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:
|
async def initialize(self) -> None:
|
||||||
"""Initialize the index backend."""
|
"""Initialize the index backend."""
|
||||||
if self._initialized:
|
if self._initialized:
|
||||||
@@ -200,9 +246,11 @@ class BaseIndexPlugin(ABC):
|
|||||||
model=self.config.embedding_model,
|
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
|
# 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()
|
store = self._create_store_with_dimension()
|
||||||
|
|
||||||
# Use config with dummy embedding URL to prevent model loading
|
# Use config with dummy embedding URL to prevent model loading
|
||||||
@@ -541,17 +589,27 @@ class BaseIndexPlugin(ABC):
|
|||||||
import asyncio
|
import asyncio
|
||||||
|
|
||||||
try:
|
try:
|
||||||
# Run retrieve in thread pool to avoid blocking event loop
|
|
||||||
# piragi's retrieve() calls embedder.embed_chunks() synchronously
|
|
||||||
ragi_sync = self.ragi._sync
|
ragi_sync = self.ragi._sync
|
||||||
|
embedder = ragi_sync.embedder
|
||||||
|
|
||||||
def _do_retrieve() -> list:
|
# Use async embedding if available (OllamaEmbedder has aembed_query)
|
||||||
# Embed the query
|
if hasattr(embedder, "aembed_query"):
|
||||||
query_embedding = ragi_sync.embedder.embed_query(query)
|
query_embedding = await embedder.aembed_query(query)
|
||||||
# Search the store
|
else:
|
||||||
return ragi_sync.store.search(query_embedding, top_k=top_k)
|
# 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 = []
|
results = []
|
||||||
for chunk in chunks:
|
for chunk in chunks:
|
||||||
# Apply filters if provided
|
# Apply filters if provided
|
||||||
@@ -579,6 +637,16 @@ class BaseIndexPlugin(ABC):
|
|||||||
)
|
)
|
||||||
return []
|
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(
|
async def ask(
|
||||||
self,
|
self,
|
||||||
query: str,
|
query: str,
|
||||||
@@ -592,57 +660,103 @@ class BaseIndexPlugin(ABC):
|
|||||||
top_k: Number of context chunks to use
|
top_k: Number of context chunks to use
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
Tuple of (answer, citations)
|
Tuple of (answer, citations). Returns ("", []) on failure
|
||||||
|
to allow OptimalService to continue to next index.
|
||||||
"""
|
"""
|
||||||
|
import asyncio
|
||||||
|
|
||||||
import httpx
|
import httpx
|
||||||
|
|
||||||
|
# Per-index timeout to prevent one slow index from blocking everything
|
||||||
|
INDEX_TIMEOUT = 15.0
|
||||||
|
|
||||||
|
search_results: list[SearchResult] = []
|
||||||
|
|
||||||
try:
|
try:
|
||||||
# First get context using our properly async search
|
async with asyncio.timeout(INDEX_TIMEOUT):
|
||||||
search_results = await self.search(query, top_k=top_k)
|
# First get context using our properly async search
|
||||||
|
logger.info(
|
||||||
if not search_results:
|
"ask() starting search",
|
||||||
return "No relevant context found.", []
|
index_type=self.index_type.value,
|
||||||
|
query=query[:50],
|
||||||
# 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,
|
|
||||||
},
|
|
||||||
)
|
)
|
||||||
if resp.is_success:
|
search_results = await self.search(query, top_k=top_k)
|
||||||
data = resp.json()
|
|
||||||
answer_text = data["choices"][0]["message"]["content"]
|
logger.info(
|
||||||
return answer_text, search_results
|
"ask() search completed",
|
||||||
else:
|
index_type=self.index_type.value,
|
||||||
logger.warning(
|
num_results=len(search_results),
|
||||||
"LLM call failed",
|
)
|
||||||
status=resp.status_code,
|
|
||||||
error=resp.text,
|
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:
|
except Exception as e:
|
||||||
logger.warning(
|
logger.warning(
|
||||||
"RAG query failed",
|
"RAG query failed",
|
||||||
index_type=self.index_type.value,
|
index_type=self.index_type.value,
|
||||||
error=str(e),
|
error=str(e),
|
||||||
)
|
)
|
||||||
return "", []
|
# Return whatever search results we have for aggregation
|
||||||
|
return "", search_results
|
||||||
|
|
||||||
async def count(self) -> int:
|
async def count(self) -> int:
|
||||||
"""Get the number of documents in the index."""
|
"""Get the number of documents in the index."""
|
||||||
|
|||||||
@@ -4,8 +4,16 @@ Code Index Plugin
|
|||||||
Handles indexing and searching code files and repositories.
|
Handles indexing and searching code files and repositories.
|
||||||
Uses a simple line-based chunking strategy instead of sentence-based,
|
Uses a simple line-based chunking strategy instead of sentence-based,
|
||||||
since code doesn't have natural sentence boundaries.
|
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 pathlib import Path
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
@@ -18,25 +26,147 @@ from roboco.services.optimal_brain.indexes.base import BaseIndexPlugin, IngestRe
|
|||||||
logger = structlog.get_logger()
|
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(
|
def chunk_code(
|
||||||
content: str,
|
content: str,
|
||||||
source: str,
|
source: str,
|
||||||
chunk_size: int = 1500,
|
chunk_size: int = 1500,
|
||||||
chunk_overlap: int = 200,
|
chunk_overlap: int = 150,
|
||||||
) -> list[Chunk]:
|
) -> 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:
|
Prefers breaking at function/class boundaries, falls back to blank lines.
|
||||||
- Splits by lines (respects code structure)
|
Uses reduced overlap since embedding cache avoids redundant work.
|
||||||
- Uses character-based sizes (not token-based, simpler and faster)
|
|
||||||
- Tries to break at blank lines or function/class boundaries
|
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
content: Source code content
|
content: Source code content
|
||||||
source: Source file path/URI
|
source: Source file path/URI
|
||||||
chunk_size: Target chunk size in characters (~375 tokens)
|
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:
|
Returns:
|
||||||
List of Chunk objects
|
List of Chunk objects
|
||||||
@@ -46,53 +176,59 @@ def chunk_code(
|
|||||||
|
|
||||||
lines = content.split("\n")
|
lines = content.split("\n")
|
||||||
chunks: list[Chunk] = []
|
chunks: list[Chunk] = []
|
||||||
current_chunk_lines: list[str] = []
|
current_lines: list[str] = []
|
||||||
current_size = 0
|
current_size = 0
|
||||||
chunk_index = 0
|
min_chunk = 300 # Minimum chunk size
|
||||||
|
|
||||||
for line in lines:
|
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
|
# Check if we need to break
|
||||||
if current_size + line_size > chunk_size and current_chunk_lines:
|
should_break = current_size + line_size > chunk_size and current_lines
|
||||||
# Save current chunk
|
|
||||||
chunk_text = "\n".join(current_chunk_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(
|
chunks.append(
|
||||||
Chunk(
|
Chunk(
|
||||||
text=chunk_text,
|
text=chunk_text,
|
||||||
source=source,
|
source=source,
|
||||||
chunk_index=chunk_index,
|
chunk_index=len(chunks),
|
||||||
metadata={},
|
metadata={},
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
chunk_index += 1
|
elif chunks:
|
||||||
|
# Merge small trailing content into last chunk
|
||||||
# Calculate overlap: keep last N characters worth of lines
|
last = chunks[-1]
|
||||||
overlap_lines: list[str] = []
|
chunks[-1] = Chunk(
|
||||||
overlap_size = 0
|
text=last.text + "\n" + chunk_text,
|
||||||
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,
|
|
||||||
source=source,
|
source=source,
|
||||||
chunk_index=chunk_index,
|
chunk_index=last.chunk_index,
|
||||||
metadata={},
|
metadata=last.metadata,
|
||||||
)
|
)
|
||||||
)
|
|
||||||
|
|
||||||
return chunks
|
return chunks
|
||||||
|
|
||||||
@@ -121,8 +257,9 @@ CODE_EXTENSIONS = {
|
|||||||
".toml": "toml",
|
".toml": "toml",
|
||||||
}
|
}
|
||||||
|
|
||||||
# Directories to skip during indexing
|
# Directories to skip during indexing (non-essential for understanding codebase)
|
||||||
SKIP_DIRECTORIES = {
|
SKIP_DIRECTORIES = {
|
||||||
|
# VCS and build
|
||||||
".git",
|
".git",
|
||||||
".venv",
|
".venv",
|
||||||
"venv",
|
"venv",
|
||||||
@@ -132,6 +269,7 @@ SKIP_DIRECTORIES = {
|
|||||||
".next",
|
".next",
|
||||||
"dist",
|
"dist",
|
||||||
"build",
|
"build",
|
||||||
|
# Cache directories
|
||||||
".mypy_cache",
|
".mypy_cache",
|
||||||
".pytest_cache",
|
".pytest_cache",
|
||||||
".ruff_cache",
|
".ruff_cache",
|
||||||
@@ -139,8 +277,34 @@ SKIP_DIRECTORIES = {
|
|||||||
".tox",
|
".tox",
|
||||||
"eggs",
|
"eggs",
|
||||||
"*.egg-info",
|
"*.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):
|
class CodeIndexPlugin(BaseIndexPlugin):
|
||||||
"""
|
"""
|
||||||
@@ -150,8 +314,23 @@ class CodeIndexPlugin(BaseIndexPlugin):
|
|||||||
- Source code files (Python, TypeScript, etc.)
|
- Source code files (Python, TypeScript, etc.)
|
||||||
- Repository directories
|
- Repository directories
|
||||||
- Glob patterns for selective indexing
|
- 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
|
@property
|
||||||
def index_type(self) -> IndexType:
|
def index_type(self) -> IndexType:
|
||||||
return IndexType.CODE
|
return IndexType.CODE
|
||||||
@@ -175,84 +354,141 @@ class CodeIndexPlugin(BaseIndexPlugin):
|
|||||||
file_path = kwargs.get("file_path", doc_id or "unknown")
|
file_path = kwargs.get("file_path", doc_id or "unknown")
|
||||||
return f"roboco://code/{file_path}"
|
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(
|
async def index_sources(
|
||||||
self,
|
self,
|
||||||
sources: list[str],
|
sources: list[str],
|
||||||
project: str | None = None,
|
project: str | None = None,
|
||||||
|
max_files: int | None = None,
|
||||||
|
force_reindex: bool = False,
|
||||||
) -> tuple[int, list[dict[str, Any]]]:
|
) -> 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,
|
Uses content hashing to skip unchanged files, dramatically reducing
|
||||||
achieving 10-15x speedup on large codebases.
|
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:
|
Args:
|
||||||
sources: List of file paths, directories, or glob patterns
|
sources: List of file paths, directories, or glob patterns
|
||||||
project: Optional project identifier for filtering
|
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:
|
Returns:
|
||||||
Tuple of (count, indexed_files) where indexed_files contains
|
Tuple of (count, indexed_files) where indexed_files contains
|
||||||
metadata for each file indexed (for database tracking)
|
metadata for each file indexed (for database tracking)
|
||||||
"""
|
"""
|
||||||
# Step 1: Collect all files and their contents
|
# Step 1: Collect all files
|
||||||
files_data: list[dict[str, Any]] = []
|
all_files: list[Path] = []
|
||||||
|
|
||||||
for source in sources:
|
for source in sources:
|
||||||
source_path = Path(source)
|
all_files.extend(self._collect_files_from_source(source))
|
||||||
|
|
||||||
# Expand glob patterns and directories
|
if not all_files:
|
||||||
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:
|
|
||||||
return 0, []
|
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)
|
results = await self._ingest_code_batch(files_data)
|
||||||
count = sum(1 for r in results if r.success)
|
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
|
# Build indexed_files list for database tracking
|
||||||
indexed_files = [
|
indexed_files = [
|
||||||
{
|
{
|
||||||
@@ -265,7 +501,10 @@ class CodeIndexPlugin(BaseIndexPlugin):
|
|||||||
for data in files_data
|
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
|
return count, indexed_files
|
||||||
|
|
||||||
async def _ingest_code_batch(
|
async def _ingest_code_batch(
|
||||||
|
|||||||
@@ -11,8 +11,10 @@ import uuid
|
|||||||
from datetime import UTC, datetime
|
from datetime import UTC, datetime
|
||||||
from typing import Any, cast
|
from typing import Any, cast
|
||||||
|
|
||||||
|
import httpx
|
||||||
import structlog
|
import structlog
|
||||||
|
|
||||||
|
from roboco.config import settings
|
||||||
from roboco.models.optimal import (
|
from roboco.models.optimal import (
|
||||||
IndexType,
|
IndexType,
|
||||||
MentorConversation,
|
MentorConversation,
|
||||||
@@ -295,15 +297,23 @@ class MentorService:
|
|||||||
|
|
||||||
async def _synthesize_answer(
|
async def _synthesize_answer(
|
||||||
self,
|
self,
|
||||||
_question: str,
|
question: str,
|
||||||
sources: list[SearchResult],
|
sources: list[SearchResult],
|
||||||
_conversation_context: str,
|
conversation_context: str,
|
||||||
) -> 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
|
Uses the local Ollama LLM to generate a coherent answer based on
|
||||||
use an LLM to generate a coherent answer.
|
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:
|
if not sources:
|
||||||
return (
|
return (
|
||||||
@@ -311,48 +321,108 @@ class MentorService:
|
|||||||
"Try rephrasing your question or asking about a different topic."
|
"Try rephrasing your question or asking about a different topic."
|
||||||
)
|
)
|
||||||
|
|
||||||
# Build answer from top sources
|
# Build context from sources, grouped by type for clarity
|
||||||
answer_parts = []
|
context_parts = []
|
||||||
|
|
||||||
# Group sources by type
|
|
||||||
by_type: dict[IndexType, list[SearchResult]] = {}
|
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:
|
if source.index_type not in by_type:
|
||||||
by_type[source.index_type] = []
|
by_type[source.index_type] = []
|
||||||
by_type[source.index_type].append(source)
|
by_type[source.index_type].append(source)
|
||||||
|
|
||||||
# Standards first
|
# Format each source type
|
||||||
if IndexType.STANDARDS in by_type:
|
type_labels = {
|
||||||
answer_parts.append("**Standards & Guidelines:**")
|
IndexType.STANDARDS: "Standards & Guidelines",
|
||||||
for s in by_type[IndexType.STANDARDS][:2]:
|
IndexType.DECISIONS: "Past Decisions",
|
||||||
answer_parts.append(f"- {s.content[:200]}...")
|
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
|
for index_type, results in by_type.items():
|
||||||
if IndexType.DECISIONS in by_type:
|
label = type_labels.get(index_type, index_type.value)
|
||||||
answer_parts.append("\n**Past Decisions:**")
|
context_parts.append(f"## {label}")
|
||||||
for s in by_type[IndexType.DECISIONS][:2]:
|
for r in results[:3]: # Max 3 per type
|
||||||
answer_parts.append(f"- {s.content[:200]}...")
|
# Include source reference and content
|
||||||
|
context_parts.append(f"[Source: {r.source}]")
|
||||||
|
context_parts.append(r.content)
|
||||||
|
context_parts.append("") # Blank line separator
|
||||||
|
|
||||||
# Learnings
|
context = "\n".join(context_parts)
|
||||||
if IndexType.LEARNINGS in by_type or IndexType.JOURNALS in by_type:
|
|
||||||
answer_parts.append("\n**Team Learnings:**")
|
# Build the prompt
|
||||||
learnings = by_type.get(IndexType.LEARNINGS, []) + by_type.get(
|
system_prompt = (
|
||||||
IndexType.JOURNALS, []
|
"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
|
user_prompt_parts.append(f"Knowledge base context:\n{context}\n")
|
||||||
if IndexType.CODE in by_type:
|
user_prompt_parts.append(f"Question: {question}")
|
||||||
answer_parts.append("\n**Code References:**")
|
|
||||||
for s in by_type[IndexType.CODE][:2]:
|
|
||||||
answer_parts.append(f"- See: {s.source}")
|
|
||||||
|
|
||||||
if not answer_parts:
|
user_prompt = "\n".join(user_prompt_parts)
|
||||||
first_content = sources[0].content[:500]
|
|
||||||
return f"Based on my search, here's what I found:\n\n{first_content}"
|
|
||||||
|
|
||||||
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(
|
async def get_conversation_history(
|
||||||
self, conversation_id: str
|
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
|
Shared Embedder Singleton
|
||||||
|
|
||||||
Provides a single EmbeddingGenerator instance shared across all index plugins
|
Provides a single embedder instance shared across all index plugins.
|
||||||
to avoid loading the SentenceTransformer model 9 times (~3s each = 27s startup).
|
Supports both Ollama models (embeddinggemma, etc.) and SentenceTransformers (BGE, etc.).
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import asyncio
|
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
|
from roboco.logging import get_logger
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from piragi.embeddings import EmbeddingGenerator
|
from piragi.embeddings import EmbeddingGenerator
|
||||||
|
|
||||||
|
from roboco.services.optimal_brain.ollama_embedder import OllamaEmbedder
|
||||||
|
|
||||||
logger = get_logger(__name__)
|
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:
|
class _SharedEmbedderHolder:
|
||||||
"""Holder class for shared embedder state (avoids global statement)."""
|
"""Holder class for shared embedder state (avoids global statement)."""
|
||||||
|
|
||||||
instance: "EmbeddingGenerator | None" = None
|
instance: Union["EmbeddingGenerator", "OllamaEmbedder", None] = None
|
||||||
lock: asyncio.Lock | None = None
|
lock: asyncio.Lock | None = None
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
@@ -31,21 +74,22 @@ class _SharedEmbedderHolder:
|
|||||||
|
|
||||||
|
|
||||||
async def get_shared_embedder(
|
async def get_shared_embedder(
|
||||||
model: str = "BAAI/bge-base-en-v1.5",
|
model: str | None = None,
|
||||||
device: str | None = None,
|
device: str | None = None,
|
||||||
timeout: float = 60.0,
|
timeout: float = 60.0,
|
||||||
) -> "EmbeddingGenerator":
|
) -> Embedder:
|
||||||
"""Get or create the shared embedder instance.
|
"""Get or create the shared embedder instance.
|
||||||
|
|
||||||
Thread-safe singleton that loads the model only once.
|
Thread-safe singleton that loads the model only once.
|
||||||
|
Automatically selects Ollama or SentenceTransformers based on model name.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
model: Embedding model name (default: BAAI/bge-base-en-v1.5)
|
model: Embedding model name (default: from settings)
|
||||||
device: Device to use (None = auto-detect)
|
device: Device to use for SentenceTransformers (None = auto-detect)
|
||||||
timeout: Max seconds to wait for model loading (default: 60)
|
timeout: Max seconds to wait for model loading (default: 60)
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
Shared EmbeddingGenerator instance
|
Shared embedder instance (OllamaEmbedder or EmbeddingGenerator)
|
||||||
|
|
||||||
Raises:
|
Raises:
|
||||||
TimeoutError: If model loading takes too long
|
TimeoutError: If model loading takes too long
|
||||||
@@ -55,50 +99,97 @@ async def get_shared_embedder(
|
|||||||
return _SharedEmbedderHolder.instance
|
return _SharedEmbedderHolder.instance
|
||||||
|
|
||||||
async with _SharedEmbedderHolder.get_lock():
|
async with _SharedEmbedderHolder.get_lock():
|
||||||
# Double-check after acquiring lock
|
# Double-check after acquiring lock (another coroutine may have created it)
|
||||||
if _SharedEmbedderHolder.instance is not None:
|
if _SharedEmbedderHolder.instance is None:
|
||||||
return _SharedEmbedderHolder.instance
|
model = model or settings.default_embedding_model
|
||||||
|
|
||||||
logger.info(
|
# Use Ollama for Ollama models, SentenceTransformers otherwise
|
||||||
"Creating shared embedder (one-time load)",
|
if _is_ollama_model(model):
|
||||||
model=model,
|
logger.info(
|
||||||
device=device or "auto",
|
"Creating shared Ollama embedder",
|
||||||
)
|
model=model,
|
||||||
|
base_url=settings.ollama_base_url,
|
||||||
# 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
|
|
||||||
|
|
||||||
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
|
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:
|
async def close_shared_embedder() -> None:
|
||||||
"""Release the shared embedder resources."""
|
"""Release the shared embedder resources."""
|
||||||
async with _SharedEmbedderHolder.get_lock():
|
async with _SharedEmbedderHolder.get_lock():
|
||||||
|
|||||||
@@ -207,6 +207,7 @@ class TaskService(BaseService):
|
|||||||
parent_task_id=req.parent_task_id,
|
parent_task_id=req.parent_task_id,
|
||||||
target_date=req.target_date,
|
target_date=req.target_date,
|
||||||
estimated_complexity=req.estimated_complexity,
|
estimated_complexity=req.estimated_complexity,
|
||||||
|
nature=req.nature,
|
||||||
status=req.status if req.status else TaskStatus.PENDING,
|
status=req.status if req.status else TaskStatus.PENDING,
|
||||||
sequence=req.sequence, # Task ordering within siblings
|
sequence=req.sequence, # Task ordering within siblings
|
||||||
dependency_ids=req.dependency_ids, # Task IDs that must complete first
|
dependency_ids=req.dependency_ids, # Task IDs that must complete first
|
||||||
|
|||||||
Reference in New Issue
Block a user