mirror of
https://github.com/rennf93/roboco.git
synced 2026-08-03 07:23:24 +02:00
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:
@@ -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
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user