Fixes Applied

1. Transaction Rollback Bug

  File: roboco/services/optimal_brain/indexes/base.py
  - Changed: On transaction error, now closes and reconnects instead of just rollback
  - Before: conn.rollback() (didn't fully reset connection state)
  - After: conn.close() + store._init_schema() (forces fresh connection)

  2. Docs Directory Path

  Files:
  - roboco/services/optimal.py - Added /app/docs as first path to check
  - docker/orchestrator.Dockerfile - Added COPY docs /app/docs

  3. Shared SentenceTransformer

  Files:
  - Created roboco/services/optimal_brain/shared_embedder.py - Singleton holder
  - Modified roboco/services/optimal_brain/indexes/base.py - Uses shared embedder
  - Modified roboco/services/optimal.py - Closes shared embedder on shutdown
This commit is contained in:
Renn F
2025-12-27 22:05:07 +01:00
parent 1116e597d0
commit ecc049684d
4 changed files with 113 additions and 5 deletions
+2
View File
@@ -41,6 +41,8 @@ COPY roboco /app/roboco
# These are composed at runtime by compose_prompt() when spawning agents
COPY agents /app/agents
COPY docker /app/docker
# docs/ contains standards and workflows for RAG auto-indexing
COPY docs /app/docs
COPY pyproject.toml uv.lock alembic.ini README.md /app/
COPY alembic /app/alembic
+7 -2
View File
@@ -120,9 +120,9 @@ class OptimalService:
"""
# Find the docs directory relative to the project root
possible_docs_roots = [
Path(__file__).parent.parent.parent / "docs", # roboco/docs
Path("/app/docs"), # Docker absolute path
Path(__file__).parent.parent.parent / "docs", # roboco/docs (local)
Path.cwd() / "docs", # Current working directory
Path.cwd() / "roboco" / "docs", # From project root
]
docs_root = None
@@ -198,6 +198,11 @@ class OptimalService:
await plugin.close()
self._plugins.clear()
self._initialized = False
# Close shared embedder
from roboco.services.optimal_brain.shared_embedder import close_shared_embedder
await close_shared_embedder()
logger.info("OptimalService closed")
def _get_plugin(self, index_type: IndexType) -> BaseIndexPlugin:
+21 -3
View File
@@ -6,7 +6,6 @@ Each plugin handles a specific content type (code, docs, errors, standards, etc.
and implements specialized chunking, metadata handling, and search strategies.
"""
import contextlib
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, cast
@@ -170,6 +169,13 @@ class BaseIndexPlugin(ABC):
persist_dir=self.config.persist_dir,
)
# Get shared embedder FIRST (one-time load for all plugins)
from roboco.services.optimal_brain.shared_embedder import get_shared_embedder
shared_embedder = await get_shared_embedder(
model=self.config.embedding_model,
)
self._ragi = AsyncRagi(
[],
persist_dir=self.config.persist_dir,
@@ -177,6 +183,10 @@ class BaseIndexPlugin(ABC):
store=self.config.store_url,
)
# Replace AsyncRagi's embedder with shared instance
# This avoids loading the model 9 times (saves ~24s startup)
self._ragi._sync.embedder = shared_embedder
self._initialized = True
logger.info(f"{self.index_type.value} index plugin initialized")
@@ -288,9 +298,17 @@ class BaseIndexPlugin(ABC):
index_type=self.index_type.value,
attempt=attempt + 1,
)
# Force connection reset - rollback alone isn't enough
if hasattr(ragi_sync.store, "_conn") and ragi_sync.store._conn:
with contextlib.suppress(Exception):
ragi_sync.store._conn.rollback()
try:
ragi_sync.store._conn.close()
# Force reconnection by reinitializing schema
ragi_sync.store._init_schema()
except Exception as reset_err:
logger.warning(
"Failed to reset connection",
error=str(reset_err),
)
continue
raise
@@ -0,0 +1,83 @@
"""
Shared Embedder Singleton
Provides a single EmbeddingGenerator instance shared across all index plugins
to avoid loading the SentenceTransformer model 9 times (~3s each = 27s startup).
"""
import asyncio
from typing import TYPE_CHECKING
from roboco.logging import get_logger
if TYPE_CHECKING:
from piragi.embeddings import EmbeddingGenerator
logger = get_logger(__name__)
class _SharedEmbedderHolder:
"""Holder class for shared embedder state (avoids global statement)."""
instance: "EmbeddingGenerator | None" = None
lock: asyncio.Lock | None = None
@classmethod
def get_lock(cls) -> asyncio.Lock:
"""Get or create the lock."""
if cls.lock is None:
cls.lock = asyncio.Lock()
return cls.lock
async def get_shared_embedder(
model: str = "all-MiniLM-L6-v2",
device: str | None = None,
) -> "EmbeddingGenerator":
"""Get or create the shared embedder instance.
Thread-safe singleton that loads the model only once.
Args:
model: Embedding model name (default: all-MiniLM-L6-v2)
device: Device to use (None = auto-detect)
Returns:
Shared EmbeddingGenerator instance
"""
if _SharedEmbedderHolder.instance is not None:
return _SharedEmbedderHolder.instance
async with _SharedEmbedderHolder.get_lock():
# Double-check after acquiring lock
if _SharedEmbedderHolder.instance is not None:
return _SharedEmbedderHolder.instance
logger.info(
"Creating shared embedder (one-time load)",
model=model,
device=device or "auto",
)
# Import here to avoid circular imports and defer heavy import
from piragi.embeddings import EmbeddingGenerator
# Run model loading in thread to not block event loop
def _create_embedder() -> "EmbeddingGenerator":
return EmbeddingGenerator(
model=model,
device=device,
batch_size=32,
)
_SharedEmbedderHolder.instance = await asyncio.to_thread(_create_embedder)
logger.info("Shared embedder created successfully")
return _SharedEmbedderHolder.instance
async def close_shared_embedder() -> None:
"""Release the shared embedder resources."""
async with _SharedEmbedderHolder.get_lock():
if _SharedEmbedderHolder.instance is not None:
logger.info("Closing shared embedder")
_SharedEmbedderHolder.instance = None