External/roboco
1
0
Fork 0
mirror of https://github.com/rennf93/roboco.git synced 2026-08-03 07:23:24 +02:00
Code Issues Packages Projects Releases Wiki Activity

fix(rag): close audit gaps in the in-house engine

Browse Source 
An adversarial audit of the piragi -> in-house swap surfaced nine confirmed
issues; this fixes all of them.

- Re-ingest now REPLACES a source's chunks instead of appending. Add
  VectorStore.delete_by_source and BaseIndexPlugin.replace_on_reingest (default
  True), called before add_chunks in both ingest paths. Without it every
  startup / periodic / manual reindex appended a fresh copy of each doc's
  chunks, growing the tables unbounded and crowding out distinct results.
  Conversations opt OUT (replace_on_reingest=False): their many messages share
  one source URI, so delete-by-source would wipe history.
- index_* now honor the plugin IngestResult. The explicit record endpoints
  (error / standard / decision / review / learning) raise on failure instead of
  writing a green tracking row for content that never persisted;
  conversation / journal indexing stays best-effort but skips the tracking row
  when the embed fails. index_message / index_entry return IngestResult.
- A deprecated index type (code) now returns 404 instead of a 500 leaked from
  _get_plugin's missing-plugin error: add OptimalService.is_index_registered
  and guard the stats / clear / refresh routes. The panel drops the dead 'Code'
  category, filter, badge, label, and mock data.
- Panel: getContext reads 'results' (matches SearchResponse) instead of a
  non-existent 'context' field; the reindex toast no longer reports phantom
  '0 code files'; the stats 'Updated' label uses the max timestamp across
  indexes rather than indexes[0]; ProactiveContextItem matches the wire shape.
- Drop the always-zero per-document chunk_count from the documents API.
- Remove dead RAG settings (hybrid_search, cross_encoder) the engine never
  consumed, and correct stale piragi / BM25 references in code, README, and
  CLAUDE.md. Delete the unused duplicate roboco/kb embedder package the swap
  shipped.

Adds tests for replace-on-reingest (incl. the conversations carve-out) and the
deprecated-index 404.
This commit is contained in:
Renn F
2026-06-15 04:55:16 +02:00
parent e53eb5b7ee
commit 6422f77bb9
24 changed files with 251 additions and 1134 deletions
Download Patch File Download Diff File Expand all files Collapse all files
roboco/agent_sdk/transcript_usage.py
+1 -1
View File
@@ -3,7 +3,7 @@
A dependency-light helper (only ``json`` + ``pathlib``) so callers that need
durable token counts — notably the orchestrator's session-finalization path —
can read them without importing the agent SDK server, which pulls in the
FastAPI / RAG (piragi / openai) stack.
FastAPI / RAG (in-house pgvector / openai) stack.
"""
from __future__ import annotations
roboco/api/routes/optimal.py
+15 -1
View File
@@ -495,6 +495,11 @@ async def get_single_index_stats(
) from e
service = await get_optimal_service()
if not service.is_index_registered(idx_type):
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail=f"Index '{index_type}' is not available",
)
stats = await service.get_index_stats(idx_type)
return SingleIndexStatsResponse(
@@ -544,6 +549,11 @@ async def clear_index(
) from e
service = await get_optimal_service()
if not service.is_index_registered(idx_type):
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail=f"Index '{index_type}' is not available",
)
await service.clear_index(idx_type)
return ClearIndexResponse(status="cleared", index_type=index_type)
@@ -583,7 +593,6 @@ async def list_documents(
metadata={
"title": d["title"],
"preview": d["preview"],
"chunk_count": d["chunk_count"],
**d["extra_data"],
},
)
@@ -618,6 +627,11 @@ async def refresh_index(
) from e
service = await get_optimal_service()
if not service.is_index_registered(idx_type):
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail=f"Index '{request.index_type}' is not available",
)
# Empty sources = "refresh everything currently registered in this
# index". The service knows how to enumerate them so the UI's
roboco/config.py
-6
View File
@@ -136,12 +136,6 @@ class Settings(BaseSettings):
description="Use HyDE (hypothetical document embeddings). "
"Makes one LLM call per query for better semantic matching.",
)
rag_use_hybrid_search: bool = Field(
default=True, description="Use BM25 + vector hybrid search"
)
rag_use_cross_encoder: bool = Field(
default=True, description="Use neural reranking (slower but more accurate)"
)
rag_auto_update_enabled: bool = Field(default=True)
rag_auto_update_interval: int = Field(
default=300, ge=60, description="Seconds between auto-updates"
roboco/kb/__init__.py
-35
View File
@@ -1,35 +0,0 @@
"""
roboco.kb — Knowledge Base primitives (piragi-free).
Public API:
Chunk — lightweight text+embedding container
OllamaEmbedder — HTTP embedder via Ollama native /api/embed
get_shared_embedder — async singleton factory
close_shared_embedder — release the shared singleton
"""
from roboco.kb.ollama_embedder import (
Chunk,
EmbeddingCache,
OllamaConnectionError,
OllamaEmbedder,
OllamaEmbedderError,
OllamaModelError,
)
from roboco.kb.shared_embedder import (
Embedder,
close_shared_embedder,
get_shared_embedder,
)
__all__ = [
"Chunk",
"Embedder",
"EmbeddingCache",
"OllamaConnectionError",
"OllamaEmbedder",
"OllamaEmbedderError",
"OllamaModelError",
"close_shared_embedder",
"get_shared_embedder",
]
roboco/kb/ollama_embedder.py
-858
View File
@@ -1,858 +0,0 @@
"""
Ollama Embedder
Provides embedding generation using Ollama's native API.
Drop-in replacement for the embedding layer 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
"""
from __future__ import annotations
import asyncio
import hashlib
import time
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from collections.abc import Callable
import httpx
from roboco.config import settings
from roboco.logging import get_logger
from roboco.services.exceptions import (
HTTP_TOO_MANY_REQUESTS,
RateLimitError,
parse_retry_after_header,
)
logger = get_logger(__name__)
# Retry configuration — ConnectError / Timeout (existing, unchanged)
MAX_RETRIES = 3
RETRY_DELAY_BASE = 0.5 # seconds, exponential backoff
# Retry configuration — HTTP 429 / RateLimitError (new outer loop)
RATE_LIMIT_MAX_RETRIES = 5
# Parallel processing configuration
MAX_CONCURRENT_BATCHES = 4 # Number of batches to process in parallel
DEFAULT_BATCH_SIZE = 32 # Default batch size
# Keep the embedding model resident in Ollama. It runs on CPU and Ollama's
# default 5-min idle unload means a `say` after an idle window pays a cold 2.4 GB
# reload before embedding; under contention with glm-5:cloud that overran the
# embed retry window and dropped the background conversation ingest. -1 = never
# unload (sent as `keep_alive` on every /api/embed request).
EMBED_KEEP_ALIVE = -1
@dataclass
class Chunk:
"""Lightweight document chunk carrying text and its embedding vector.
Replaces the piragi.types.Chunk dependency so this module compiles
without piragi installed.
"""
text: str
embedding: list[float] = field(default_factory=list)
metadata: dict[str, Any] = field(default_factory=dict)
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 = "qwen3-embedding:0.6b",
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 sync client (async clients created per-operation)
self._sync_client: httpx.Client | None = None
# Semaphore for limiting concurrent requests. Track the loop it was
# created on so we can rebuild when the loop rotates.
self._semaphore: asyncio.Semaphore | None = None
self._semaphore_loop: asyncio.AbstractEventLoop | None = None
def _embed_payload(self, input_data: str | list[str]) -> dict[str, object]:
"""Build the /api/embed JSON body, pinning keep_alive so the model stays
resident in Ollama (see EMBED_KEEP_ALIVE)."""
return {
"model": self.model,
"input": input_data,
"keep_alive": EMBED_KEEP_ALIVE,
}
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
def _create_async_client(self) -> httpx.AsyncClient:
"""Create a fresh async HTTP client.
Always creates a new client to avoid 'Event loop is closed' errors
that occur when a cached client is bound to a different/closed loop.
"""
timeout = httpx.Timeout(
connect=10.0,
read=self.timeout,
write=30.0,
pool=10.0,
)
return httpx.AsyncClient(
timeout=timeout,
limits=httpx.Limits(
max_connections=self.max_concurrent * 2,
max_keepalive_connections=self.max_concurrent,
),
)
def _get_semaphore(self) -> asyncio.Semaphore:
"""Get or create semaphore for limiting concurrent requests.
asyncio.Semaphore binds to the event loop it was created in. If the
orchestrator's loop rotates (e.g. lifespan restart, test teardown),
a cached semaphore raises "bound to a different event loop". Detect
loop rotation by comparing the current running loop to the one we
recorded at creation time, and rebuild if they differ.
"""
current_loop = asyncio.get_running_loop()
if self._semaphore is None or self._semaphore_loop is not current_loop:
self._semaphore = asyncio.Semaphore(self.max_concurrent)
self._semaphore_loop = current_loop
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."""
# Async clients are now created per-request, nothing to close
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
@staticmethod
def _rl_backoff(retry_after: float | None, rl_attempt: int) -> float:
"""Backoff seconds for a 429: honor Retry-After, else exponential."""
return retry_after if retry_after is not None else float(2**rl_attempt)
@staticmethod
def _map_embed_error(e: Exception, base_url: str) -> Exception:
"""Map a raw request exception to the appropriate Ollama embedder error."""
if isinstance(e, httpx.ConnectError):
return OllamaConnectionError(f"Cannot connect to Ollama at {base_url}: {e}")
if isinstance(e, httpx.TimeoutException):
return OllamaConnectionError(f"Ollama request timed out: {e}")
return OllamaEmbedderError(f"Unexpected error: {e}")
@staticmethod
def _log_429(rl_attempt: int, backoff: float) -> None:
"""Log an Ollama 429 rate-limit retry."""
logger.warning(
"Ollama rate limited (429), retrying",
provider="ollama",
attempt=rl_attempt + 1,
max_retries=RATE_LIMIT_MAX_RETRIES,
backoff_duration=backoff,
)
@staticmethod
def _sleep_connect_retry(
attempt: int, last_error: Exception | None, label: str, **extra: Any
) -> None:
"""Exponential backoff between ConnectError/Timeout retries.
No sleep on the final attempt — the caller raises ``last_error`` then.
"""
if attempt < MAX_RETRIES - 1:
delay = RETRY_DELAY_BASE * (2**attempt)
logger.warning(
label, attempt=attempt + 1, delay=delay, error=str(last_error), **extra
)
time.sleep(delay)
@staticmethod
async def _asleep_connect_retry(
attempt: int, last_error: Exception | None, label: str, **extra: Any
) -> None:
"""Async counterpart of :meth:`_sleep_connect_retry`."""
if attempt < MAX_RETRIES - 1:
delay = RETRY_DELAY_BASE * (2**attempt)
logger.warning(
label, attempt=attempt + 1, delay=delay, error=str(last_error), **extra
)
await asyncio.sleep(delay)
def embed_query(
self,
query: str,
task_instruction: str | None = None,
) -> list[float]:
"""Generate embedding for a single query.
Retry behaviour (two independent concerns, non-overlapping):
- ConnectError / TimeoutException: up to MAX_RETRIES=3 attempts with
0.5/1/2 s exponential backoff (existing behaviour, unchanged).
- HTTP 429 (rate limit): outer loop up to RATE_LIMIT_MAX_RETRIES=5,
respecting Retry-After header. A 429 response does NOT trigger the
ConnectError path.
"""
_ = task_instruction
# Check cache first
cached = self._cache.get(query)
if cached is not None:
return cached
client = self._get_sync_client()
last_rl_retry_after: float | None = None
for rl_attempt in range(RATE_LIMIT_MAX_RETRIES):
got_429 = False
last_error: Exception | None = None
# --- inner loop: ConnectError / Timeout (unchanged) ---
for attempt in range(MAX_RETRIES):
try:
response = client.post(
f"{self.base_url}/api/embed",
json=self._embed_payload(query),
)
# 429 check — must NOT enter the ConnectError path
if response.status_code == HTTP_TOO_MANY_REQUESTS:
last_rl_retry_after = parse_retry_after_header(response)
self._log_429(
rl_attempt,
self._rl_backoff(last_rl_retry_after, rl_attempt),
)
got_429 = True
break # break inner loop; outer loop will sleep + retry
embeddings = self._handle_embed_response(response, input_count=1)
result = embeddings[0]
self._cache.put(query, result)
return result
except (OllamaModelError, OllamaEmbedderError):
raise
except Exception as e:
last_error = self._map_embed_error(e, self.base_url)
self._sleep_connect_retry(
attempt, last_error, "Ollama embed_query retry"
)
# --- end inner loop ---
if not got_429:
# ConnectError / Timeout exhausted — same behaviour as before
raise last_error or OllamaEmbedderError("Max retries exceeded")
# 429: sleep and try again (outer loop)
if rl_attempt < RATE_LIMIT_MAX_RETRIES - 1:
time.sleep(self._rl_backoff(last_rl_retry_after, rl_attempt))
raise RateLimitError(provider="ollama", retry_after=last_rl_retry_after)
def _embed_batch_sync(
self, client: httpx.Client, batch: list[str], batch_index: int
) -> list[list[float]]:
"""Embed a single batch synchronously.
Same two-concern retry composition as :meth:`embed_query`:
inner ConnectError/Timeout loop (unchanged) + outer 429 loop.
"""
last_rl_retry_after: float | None = None
for rl_attempt in range(RATE_LIMIT_MAX_RETRIES):
got_429 = False
last_error: Exception | None = None
for attempt in range(MAX_RETRIES):
try:
response = client.post(
f"{self.base_url}/api/embed",
json=self._embed_payload(batch),
)
if response.status_code == HTTP_TOO_MANY_REQUESTS:
last_rl_retry_after = parse_retry_after_header(response)
self._log_429(
rl_attempt,
self._rl_backoff(last_rl_retry_after, rl_attempt),
)
got_429 = True
break
return self._handle_embed_response(response, input_count=len(batch))
except (OllamaModelError, OllamaEmbedderError):
raise
except Exception as e:
last_error = self._map_embed_error(e, self.base_url)
self._sleep_connect_retry(
attempt,
last_error,
"Ollama embed_documents retry",
batch_index=batch_index,
)
if not got_429:
raise last_error or OllamaEmbedderError("Max retries exceeded")
if rl_attempt < RATE_LIMIT_MAX_RETRIES - 1:
time.sleep(self._rl_backoff(last_rl_retry_after, rl_attempt))
raise RateLimitError(provider="ollama", retry_after=last_rl_retry_after)
def _partition_cached_documents(
self, documents: list[str]
) -> tuple[list[list[float] | None], list[int], list[str]]:
"""Split documents into (pre-filled slots, uncached indices, uncached texts)."""
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)
return result_embeddings, uncached_indices, uncached_docs
@staticmethod
def _merge_embeddings(
result_embeddings: list[list[float] | None],
uncached_indices: list[int],
new_embeddings: list[list[float]],
) -> list[list[float]]:
"""Merge freshly-computed embeddings into preallocated result list."""
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_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 []
(
result_embeddings,
uncached_indices,
uncached_docs,
) = self._partition_cached_documents(documents)
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]
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)
for doc, emb in zip(batch, embeddings, strict=True):
self._cache.put(doc, emb)
return self._merge_embeddings(
result_embeddings, uncached_indices, new_embeddings
)
async def _embed_batch_async(
self,
client: httpx.AsyncClient,
batch: list[str],
batch_index: int,
) -> list[list[float]]:
"""Embed a single batch with semaphore-limited concurrency.
Same two-concern retry composition as :meth:`embed_query`:
inner ConnectError/Timeout loop (unchanged) + outer 429 loop (async).
"""
semaphore = self._get_semaphore()
async with semaphore:
last_rl_retry_after: float | None = None
for rl_attempt in range(RATE_LIMIT_MAX_RETRIES):
got_429 = False
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=self._embed_payload(batch),
)
if response.status_code == HTTP_TOO_MANY_REQUESTS:
last_rl_retry_after = parse_retry_after_header(response)
self._log_429(
rl_attempt,
self._rl_backoff(last_rl_retry_after, rl_attempt),
)
got_429 = True
break
return self._handle_embed_response(
response, input_count=len(batch)
)
except (OllamaModelError, OllamaEmbedderError):
raise
except Exception as e:
last_error = self._map_embed_error(e, self.base_url)
await self._asleep_connect_retry(
attempt,
last_error,
"Parallel embed batch retry",
batch_index=batch_index,
)
if not got_429:
raise last_error or OllamaEmbedderError("Max retries exceeded")
if rl_attempt < RATE_LIMIT_MAX_RETRIES - 1:
await asyncio.sleep(
self._rl_backoff(last_rl_retry_after, rl_attempt)
)
raise RateLimitError(provider="ollama", retry_after=last_rl_retry_after)
async def _run_parallel_batches(
self, batches: list[list[str]]
) -> list[list[list[float]]]:
"""Execute all batches concurrently via a shared async client."""
async with self._create_async_client() as client:
tasks = [
self._embed_batch_async(client, batch, i)
for i, batch in enumerate(batches)
]
return await asyncio.gather(*tasks)
def _collect_and_cache(
self,
batches: list[list[str]],
batch_results: list[list[list[float]]],
) -> list[list[float]]:
"""Flatten batch results and populate the embedding cache."""
new_embeddings: list[list[float]] = []
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)
return new_embeddings
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 []
(
result_embeddings,
uncached_indices,
uncached_docs,
) = self._partition_cached_documents(documents)
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]
batches: list[list[str]] = [
uncached_docs[i : i + batch_size]
for i in range(0, len(uncached_docs), 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()
batch_results = await self._run_parallel_batches(batches)
new_embeddings = self._collect_and_cache(batches, batch_results)
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}",
)
return self._merge_embeddings(
result_embeddings, uncached_indices, new_embeddings
)
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.
Same two-concern retry composition as :meth:`embed_query`:
inner ConnectError/Timeout loop (unchanged) + outer 429 loop (async).
"""
# Check cache first
cached = self._cache.get(query)
if cached is not None:
return cached
last_rl_retry_after: float | None = None
for rl_attempt in range(RATE_LIMIT_MAX_RETRIES):
got_429 = False
last_error: Exception | None = None
for attempt in range(MAX_RETRIES):
# Create fresh client each attempt to avoid event loop issues
async with self._create_async_client() as client:
try:
response = await client.post(
f"{self.base_url}/api/embed",
json=self._embed_payload(query),
)
if response.status_code == HTTP_TOO_MANY_REQUESTS:
last_rl_retry_after = parse_retry_after_header(response)
self._log_429(
rl_attempt,
self._rl_backoff(last_rl_retry_after, rl_attempt),
)
got_429 = True
break
embeddings = self._handle_embed_response(
response, input_count=1
)
result = embeddings[0]
self._cache.put(query, result)
return result
except (OllamaModelError, OllamaEmbedderError):
raise
except Exception as e:
last_error = self._map_embed_error(e, self.base_url)
# A 429 already broke the inner loop above; otherwise back off.
await self._asleep_connect_retry(
attempt, last_error, "Ollama aembed_query retry"
)
if not got_429:
raise last_error or OllamaEmbedderError("Max retries exceeded")
if rl_attempt < RATE_LIMIT_MAX_RETRIES - 1:
await asyncio.sleep(self._rl_backoff(last_rl_retry_after, rl_attempt))
raise RateLimitError(provider="ollama", retry_after=last_rl_retry_after)
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)
roboco/kb/shared_embedder.py
-159
View File
@@ -1,159 +0,0 @@
"""
Shared Embedder Singleton
Provides a single embedder instance shared across all index plugins.
Supports Ollama models (qwen3-embedding, ...) via OllamaEmbedder.
"""
from __future__ import annotations
import asyncio
from typing import TYPE_CHECKING, Protocol
from roboco.config import settings
from roboco.kb.ollama_embedder import Chunk, OllamaEmbedder
from roboco.logging import get_logger
if TYPE_CHECKING:
from collections.abc import Callable
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 = {
"qwen3-embedding",
"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(":", maxsplit=1)[0].lower()
return model_base in OLLAMA_EMBEDDING_MODELS
class _SharedEmbedderHolder:
"""Holder class for shared embedder state (avoids global statement)."""
instance: OllamaEmbedder | 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 | None = None,
_device: str | None = None,
_timeout: float = 60.0,
) -> Embedder:
"""Get or create the shared embedder instance.
Thread-safe singleton that loads the model only once.
Uses OllamaEmbedder for all embedding models (no piragi dependency).
Args:
model: Embedding model name (default: from settings)
_device: Unused — kept for API compatibility with the old signature
_timeout: Unused — Ollama handles load timing server-side
Returns:
Shared embedder instance (OllamaEmbedder)
Raises:
RuntimeError: If model loading fails
"""
if _SharedEmbedderHolder.instance is not None:
return _SharedEmbedderHolder.instance
async with _SharedEmbedderHolder.get_lock():
# Double-check after acquiring lock (another coroutine may have created it)
if _SharedEmbedderHolder.instance is None:
resolved_model = model or settings.default_embedding_model
logger.info(
"Creating shared Ollama embedder",
model=resolved_model,
base_url=settings.ollama_base_url,
)
_SharedEmbedderHolder.instance = OllamaEmbedder(
model=resolved_model,
base_url=settings.ollama_base_url,
)
# Validate embedder implements required protocol methods.
# Using explicit check (not assert) so this survives `python -O`.
if _SharedEmbedderHolder.instance is None:
raise RuntimeError(
"Shared embedder construction succeeded but instance is None"
)
_validate_embedder_protocol(_SharedEmbedderHolder.instance, resolved_model)
logger.info("Shared embedder created successfully", model=resolved_model)
if _SharedEmbedderHolder.instance is None:
raise RuntimeError("Shared embedder not initialized")
return _SharedEmbedderHolder.instance
def _validate_embedder_protocol(embedder: Embedder, model: str) -> None:
"""
Validate that embedder implements the required protocol methods.
Checks at creation time rather than failing during first use.
Args:
embedder: The embedder instance to validate
model: Model name for error messages
Raises:
RuntimeError: If embedder is missing required methods
"""
required_methods = ["embed_query", "embed_documents", "embed_chunks"]
missing = [m for m in required_methods if not callable(getattr(embedder, m, None))]
if missing:
raise RuntimeError(
f"Embedder for model '{model}' is missing required methods: {missing}. "
f"Embedder type: {type(embedder).__name__}"
)
async def close_shared_embedder() -> None:
"""Release the shared embedder resources."""
async with _SharedEmbedderHolder.get_lock():
if _SharedEmbedderHolder.instance is not None:
logger.info("Closing shared embedder")
_SharedEmbedderHolder.instance = None
roboco/mcp/__init__.py
+1 -1
View File
@@ -15,7 +15,7 @@ Servers:
Do NOT eagerly re-export server factories here. Each server is launched
as its own subprocess via ``python -m roboco.mcp.<name>``, and importing
the ``roboco.mcp`` package first forces every sibling module to load —
most notably ``optimal_server``, which pulls in piragi/ollama and adds
most notably ``optimal_server``, which pulls in the pgvector/ollama stack and adds
~6s to startup. Claude Code times out slow MCP servers during
init, which manifests as "roboco-flow/do tools never register". Keep
this file empty-of-imports; callers import the specific module they
roboco/services/optimal.py
+57 -11
View File
@@ -1,7 +1,8 @@
"""
Optimal API Service (Refactored with Plugin Architecture)
Knowledge base, RAG queries, and prompt optimization using piragi.
Knowledge base, RAG queries, and prompt optimization over an in-house
PostgreSQL/pgvector engine.
This service provides semantic search across documentation,
conversations, journal entries, errors, standards, decisions, reviews, and learnings.
@@ -15,7 +16,7 @@ import asyncio
import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any
from typing import TYPE_CHECKING, Any
import structlog
@@ -49,6 +50,9 @@ from roboco.services.optimal_brain.indexes.reviews import (
RecordReviewParams as ReviewParams,
)
if TYPE_CHECKING:
from roboco.services.optimal_brain.indexes.base import IngestResult
logger = structlog.get_logger()
# Max chars per citation content - increased for better synthesis quality
@@ -147,8 +151,8 @@ class OptimalService:
"""
Service for knowledge base operations and RAG queries.
Uses a plugin-based architecture with piragi and PostgreSQL/pgvector
for vector storage. Manages multiple indexes for different content types:
Uses a plugin-based architecture over an in-house PostgreSQL/pgvector
vector store. Manages multiple indexes for different content types:
Existing:
- Code: Repositories, functions, classes
@@ -594,6 +598,15 @@ class OptimalService:
)
return plugin
def is_index_registered(self, index_type: IndexType) -> bool:
"""Whether a live plugin is registered for *index_type*.
Distinguishes a valid-but-deprecated index (e.g. ``code``) from an
active one, so callers can return 404 instead of leaking a 500 from
:meth:`_get_plugin`.
"""
return self._initialized and index_type in self._plugins
# =========================================================================
# INDEXING OPERATIONS (Existing - Backwards Compatible)
# =========================================================================
@@ -741,9 +754,9 @@ class OptimalService:
"""Index a conversation message."""
plugin = self._get_plugin(IndexType.CONVERSATIONS)
if isinstance(plugin, ConversationsIndexPlugin):
await plugin.index_message(params)
result = await plugin.index_message(params)
else:
await plugin.ingest(
result = await plugin.ingest(
content=params.content,
channel_id=params.channel_id,
session_id=params.session_id,
@@ -760,6 +773,16 @@ class OptimalService:
"build doc-source with placeholder. Caller must pass a "
"flushed session UUID."
)
# Best-effort: an actual failure (e.g. embedder down) must not record
# the message as indexed. A successful-but-empty result (short message
# filtered to zero chunks) still tracks, matching prior behavior.
if not result.success:
logger.warning(
"Conversation indexing failed; skipping tracking row",
session_id=str(params.session_id),
error=result.error,
)
return
source = f"roboco://conversations/{params.session_id}"
await self._track_indexed_document(
IndexType.CONVERSATIONS,
@@ -777,9 +800,9 @@ class OptimalService:
"""Index a journal entry."""
plugin = self._get_plugin(IndexType.JOURNALS)
if isinstance(plugin, JournalsIndexPlugin):
await plugin.index_entry(params)
result = await plugin.index_entry(params)
else:
await plugin.ingest(
result = await plugin.ingest(
content=params.content,
entry_id=params.entry_id,
agent_id=params.agent_id,
@@ -799,6 +822,14 @@ class OptimalService:
"events without a journal entry must use a different "
"indexing path."
)
# Best-effort: don't record a failed embed as an indexed entry.
if not result.success:
logger.warning(
"Journal indexing failed; skipping tracking row",
entry_id=str(params.entry_id),
error=result.error,
)
return
source = f"roboco://journals/{params.entry_id}"
await self._track_indexed_document(
IndexType.JOURNALS,
@@ -820,8 +851,11 @@ class OptimalService:
async def index_error(self, params: IndexErrorParams) -> None:
"""Index an error pattern with solution."""
plugin = self._get_plugin(IndexType.ERRORS)
result: IngestResult | None = None
if isinstance(plugin, ErrorsIndexPlugin):
await plugin.record_error(params)
result = await plugin.record_error(params)
if result is not None and not result.success:
raise RuntimeError(f"Failed to index error pattern: {result.error}")
# Track in database
import hashlib
@@ -845,8 +879,11 @@ class OptimalService:
async def index_standard(self, params: IndexStandardParams) -> None:
"""Index a coding/security/workflow standard."""
plugin = self._get_plugin(IndexType.STANDARDS)
result: IngestResult | None = None
if isinstance(plugin, StandardsIndexPlugin):
await plugin.index_standard(params)
result = await plugin.index_standard(params)
if result is not None and not result.success:
raise RuntimeError(f"Failed to index standard: {result.error}")
# Track in database
source = f"roboco://standards/{params.domain or 'general'}"
@@ -866,8 +903,11 @@ class OptimalService:
async def index_decision(self, params: IndexDecisionParams) -> None:
"""Index an architectural/design decision."""
plugin = self._get_plugin(IndexType.DECISIONS)
result: IngestResult | None = None
if isinstance(plugin, DecisionsIndexPlugin):
await plugin.record_decision(params)
result = await plugin.record_decision(params)
if result is not None and not result.success:
raise RuntimeError(f"Failed to index decision: {result.error}")
# Track in database
import hashlib
@@ -895,6 +935,7 @@ class OptimalService:
"""Record a code review for future reference."""
plugin = self._get_plugin(IndexType.REVIEWS)
doc_id = ""
result: IngestResult | None = None
if isinstance(plugin, ReviewsIndexPlugin):
review_params = ReviewParams(
comment=params.summary,
@@ -906,6 +947,8 @@ class OptimalService:
)
result = await plugin.record_review(review_params)
doc_id = result.doc_id
if result is not None and not result.success:
raise RuntimeError(f"Failed to record review: {result.error}")
# Track in database. Refuse to fall back to 'unknown' — file_path
# is typed `str` (required) and an empty value means the caller is
@@ -935,9 +978,12 @@ class OptimalService:
"""Record a learning for cross-agent knowledge sharing."""
plugin = self._get_plugin(IndexType.LEARNINGS)
doc_id = ""
result: IngestResult | None = None
if isinstance(plugin, LearningsIndexPlugin):
result = await plugin.record_learning(params)
doc_id = result.doc_id
if result is not None and not result.success:
raise RuntimeError(f"Failed to record learning: {result.error}")
# Track in database
import hashlib
roboco/services/optimal_brain/indexes/base.py
+13 -4
View File
@@ -68,8 +68,6 @@ class IndexConfig:
chunk_size: int = 512
chunk_overlap: int = 50
use_hyde: bool = True
use_hybrid_search: bool = True
use_cross_encoder: bool = False
embedding_model: str = "qwen3-embedding:0.6b"
llm_model: str = "glm-5:cloud"
llm_base_url: str = "http://roboco-ollama:11434/v1"
@@ -91,8 +89,6 @@ class IndexConfig:
chunk_size=chunk_size,
chunk_overlap=settings.rag_chunk_overlap,
use_hyde=settings.rag_use_hyde,
use_hybrid_search=settings.rag_use_hybrid_search,
use_cross_encoder=settings.rag_use_cross_encoder,
embedding_model=settings.default_embedding_model,
llm_model=settings.local_llm_model,
llm_base_url=settings.local_llm_base_url,
@@ -137,6 +133,13 @@ class BaseIndexPlugin(ABC):
- validate_content() for content validation
"""
# When True (default), re-ingesting a source first deletes that source's
# existing chunks so a reindex *replaces* rather than appends — preventing
# unbounded duplicate-chunk growth on startup/periodic/manual reindex.
# Plugins whose multiple records share one source URI (conversations) set
# this False to preserve append semantics.
replace_on_reingest: bool = True
def __init__(self, config: IndexConfig | None = None) -> None:
"""Initialize the plugin with optional config override."""
self._config = config
@@ -463,6 +466,8 @@ class BaseIndexPlugin(ABC):
embedder.embed_chunks, chunks
)
if self.replace_on_reingest:
await store.delete_by_source(doc.source)
await store.add_chunks(chunks_with_embeddings)
return len(chunks)
@@ -564,6 +569,10 @@ class BaseIndexPlugin(ABC):
chunks_with_embeddings = await asyncio.to_thread(
embedder.embed_chunks, all_chunks
)
if self.replace_on_reingest:
for idx, (doc, _, _) in enumerate(docs_to_process):
if chunk_counts.get(idx, 0) > 0:
await store.delete_by_source(doc.source)
await store.add_chunks(chunks_with_embeddings)
@staticmethod
roboco/services/optimal_brain/indexes/conversations.py
+15 -3
View File
@@ -8,7 +8,11 @@ from typing import Any
from uuid import UUID
from roboco.models.optimal import IndexConversationParams, IndexType
from roboco.services.optimal_brain.indexes.base import BaseIndexPlugin, build_doc_source
from roboco.services.optimal_brain.indexes.base import (
BaseIndexPlugin,
IngestResult,
build_doc_source,
)
class ConversationsIndexPlugin(BaseIndexPlugin):
@@ -21,6 +25,11 @@ class ConversationsIndexPlugin(BaseIndexPlugin):
- Channel conversations
"""
# Many messages share one source URI per session+agent, so deleting by
# source on re-ingest would wipe earlier messages. Keep append semantics
# for this index (the only one whose source is not 1:1 with a record).
replace_on_reingest = False
@property
def index_type(self) -> IndexType:
return IndexType.CONVERSATIONS
@@ -55,14 +64,17 @@ class ConversationsIndexPlugin(BaseIndexPlugin):
combined = f"{raw_session}-{agent_id}"
return build_doc_source(kind="conversations", id_=combined)
async def index_message(self, params: IndexConversationParams) -> None:
async def index_message(self, params: IndexConversationParams) -> IngestResult:
"""
Index a conversation message.
Args:
params: IndexConversationParams containing message details
Returns:
IngestResult so the caller can tell whether the message persisted.
"""
await self.ingest(
return await self.ingest(
content=params.content,
doc_id=f"{params.session_id}-{params.agent_id}"[:50],
channel_id=params.channel_id,
roboco/services/optimal_brain/indexes/journals.py
+10 -3
View File
@@ -8,7 +8,11 @@ from typing import Any
from uuid import UUID
from roboco.models.optimal import IndexJournalEntryParams, IndexType
from roboco.services.optimal_brain.indexes.base import BaseIndexPlugin, build_doc_source
from roboco.services.optimal_brain.indexes.base import (
BaseIndexPlugin,
IngestResult,
build_doc_source,
)
class JournalsIndexPlugin(BaseIndexPlugin):
@@ -49,14 +53,17 @@ class JournalsIndexPlugin(BaseIndexPlugin):
entry_id = str(raw) if raw is not None else None
return build_doc_source(kind="journals", id_=entry_id)
async def index_entry(self, params: IndexJournalEntryParams) -> None:
async def index_entry(self, params: IndexJournalEntryParams) -> IngestResult:
"""
Index a journal entry.
Args:
params: IndexJournalEntryParams containing entry details
Returns:
IngestResult so the caller can tell whether the entry persisted.
"""
await self.ingest(
return await self.ingest(
content=params.content,
doc_id=str(params.entry_id)[:50],
entry_id=params.entry_id,
roboco/services/optimal_brain/text_chunker.py
+2 -2
View File
@@ -1,7 +1,7 @@
"""
Text Chunker — in-house character-based sliding-window chunker.
This module replaces the piragi ``Chunker`` and provides the data-transfer
This module is the in-house ``Chunker`` and provides the data-transfer
types (``Chunk``, ``Document``, ``Citation``) that the rest of the optimal
brain stack uses. **No external tokenizers** are required: chunking is done
by character count only, which is a good approximation for our embedding
@@ -25,7 +25,7 @@ from dataclasses import dataclass, field
from typing import Any
# ---------------------------------------------------------------------------
# Transfer types — replicate the piragi.types surface used by the stack
# Transfer types — the data-transfer surface used across the RAG stack
# ---------------------------------------------------------------------------
roboco/services/optimal_brain/vector_store.py
+14
View File
@@ -196,6 +196,20 @@ class VectorStore:
records,
)
async def delete_by_source(self, source: str) -> None:
"""Delete every chunk row for *source* (idempotent — no-op if absent).
Makes re-ingestion of a source *replace* its chunks instead of
appending a duplicate set on each reindex. ``source`` is bound as a
query parameter; only the validated table identifier is interpolated.
"""
pool = self._require_pool()
async with pool.acquire() as conn:
await conn.execute(
self._q("DELETE FROM {table} WHERE source = $1"),
source,
)
# ------------------------------------------------------------------
# Read
# ------------------------------------------------------------------