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
|
# These are composed at runtime by compose_prompt() when spawning agents
|
||||||
COPY agents /app/agents
|
COPY agents /app/agents
|
||||||
COPY docker /app/docker
|
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 pyproject.toml uv.lock alembic.ini README.md /app/
|
||||||
COPY alembic /app/alembic
|
COPY alembic /app/alembic
|
||||||
|
|
||||||
|
|||||||
@@ -120,9 +120,9 @@ class OptimalService:
|
|||||||
"""
|
"""
|
||||||
# Find the docs directory relative to the project root
|
# Find the docs directory relative to the project root
|
||||||
possible_docs_roots = [
|
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() / "docs", # Current working directory
|
||||||
Path.cwd() / "roboco" / "docs", # From project root
|
|
||||||
]
|
]
|
||||||
|
|
||||||
docs_root = None
|
docs_root = None
|
||||||
@@ -198,6 +198,11 @@ class OptimalService:
|
|||||||
await plugin.close()
|
await plugin.close()
|
||||||
self._plugins.clear()
|
self._plugins.clear()
|
||||||
self._initialized = False
|
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")
|
logger.info("OptimalService closed")
|
||||||
|
|
||||||
def _get_plugin(self, index_type: IndexType) -> BaseIndexPlugin:
|
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.
|
and implements specialized chunking, metadata handling, and search strategies.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import contextlib
|
|
||||||
from abc import ABC, abstractmethod
|
from abc import ABC, abstractmethod
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
from typing import Any, cast
|
from typing import Any, cast
|
||||||
@@ -170,6 +169,13 @@ class BaseIndexPlugin(ABC):
|
|||||||
persist_dir=self.config.persist_dir,
|
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(
|
self._ragi = AsyncRagi(
|
||||||
[],
|
[],
|
||||||
persist_dir=self.config.persist_dir,
|
persist_dir=self.config.persist_dir,
|
||||||
@@ -177,6 +183,10 @@ class BaseIndexPlugin(ABC):
|
|||||||
store=self.config.store_url,
|
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
|
self._initialized = True
|
||||||
logger.info(f"{self.index_type.value} index plugin initialized")
|
logger.info(f"{self.index_type.value} index plugin initialized")
|
||||||
|
|
||||||
@@ -288,9 +298,17 @@ class BaseIndexPlugin(ABC):
|
|||||||
index_type=self.index_type.value,
|
index_type=self.index_type.value,
|
||||||
attempt=attempt + 1,
|
attempt=attempt + 1,
|
||||||
)
|
)
|
||||||
|
# Force connection reset - rollback alone isn't enough
|
||||||
if hasattr(ragi_sync.store, "_conn") and ragi_sync.store._conn:
|
if hasattr(ragi_sync.store, "_conn") and ragi_sync.store._conn:
|
||||||
with contextlib.suppress(Exception):
|
try:
|
||||||
ragi_sync.store._conn.rollback()
|
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
|
continue
|
||||||
raise
|
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