Skip to content

⚡ Bolt: Cache numpy arrays in embedding search #105

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
daggerstuff wants to merge 2 commits into
base: staging
Choose a base branch
from
Open

⚡ Bolt: Cache numpy arrays in embedding search #105

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
7cd0e2a
Cache numpy arrays in embedding search
google-labs-jules[bot] Mar 30, 2026
a7dc1fd
Fix CodeFactor complex method issue in embedding agent
google-labs-jules[bot] Mar 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 34 additions & 23 deletions api/embedding_agent/service.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -502,6 +502,38 @@ def search_similar(
model_used=self.config.model_name.value,
)

def _compute_similarity(
self, item: Any, query_embedding: List[float], query_np: Any
) -> Optional[float]:
# ⚡ Bolt: Cache normalized numpy arrays safely to avoid regenerating on query
item_np = getattr(item, "_cached_np_embedding", None)

if item_np is None:
item_embedding = getattr(item, "embedding", None)
if item_embedding is None:
return None
if NUMPY_AVAILABLE and query_np is not None:
item_np = np.array(item_embedding)
if self.config.normalize_embeddings:
norm = np.linalg.norm(item_np)
if norm > 0:
item_np = item_np / norm
try:
setattr(item, "_cached_np_embedding", item_np)
except AttributeError:
pass # Skip caching if item is a strict class (e.g., uses slots)
Comment on lines +522 to +524
Copy link
Copy Markdown

@chatgpt-codex-connector chatgpt-codex-connector bot Mar 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Handle cache-write failures beyond AttributeError

In _compute_similarity, the new cache path writes item_np back onto each knowledge item with setattr, but it only catches AttributeError. Strict model objects (for example, classes that reject unknown attributes via __setattr__) commonly raise ValueError or TypeError instead, which will now bubble up and fail the entire similarity search request for that item set. This optimization path should degrade gracefully (skip caching) rather than turning unsupported item types into request-time exceptions.

Useful? React with 👍 / 👎.

Comment on lines +521 to +524
Copy link
Copy Markdown

@coderabbitai coderabbitai bot Mar 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion | 🟠 Major

Replace pass with explicit logging or structured handling.

The pass statement violates coding guidelines which prohibit "stubs, placeholders, pass statements." While the except AttributeError is a valid fallback, silently swallowing it obscures when caching fails. Log at debug level so operators can diagnose performance issues if caching is unexpectedly disabled for certain item types.

♻️ Proposed fix 
                 try:
                     setattr(item, "_cached_np_embedding", item_np)
                 except AttributeError:
-                    pass  # Skip caching if item is a strict class (e.g., uses slots)
+                    logger.debug(
+                        "Cannot cache numpy embedding on item %s: attribute assignment restricted",
+                        type(item).__name__,
+                    )

As per coding guidelines: "Absolutely prohibit stubs, placeholders, pass statements, TODO comments, NotImplementedError, or FIXME comments; every implementation must be production-ready."

🧰 Tools
🪛 Ruff (0.15.7)

[warning] 522-522: Do not call setattr with a constant attribute value. It is not any safer than normal property access.

Replace setattr with assignment

(B010)

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@api/embedding_agent/service.py` around lines 521 - 524, Replace the silent
except block that swallows AttributeError when running setattr(item,
"_cached_np_embedding", item_np) with an explicit debug log so caching failures
are observable; in the except AttributeError handler call a logger.debug (e.g.,
logger.debug or get module logger via logging.getLogger(__name__)) and include
the item repr/type and the exception message to show why caching was skipped,
rather than using pass. Ensure the logger reference you use is consistent with
the module (or import logging if needed) and keep the handler limited to
debug-level output.


if NUMPY_AVAILABLE and query_np is not None and item_np is not None:
return float(np.dot(query_np, item_np))

item_embedding = getattr(item, "embedding", None)
if item_embedding is None:
return None
dot_product = sum(a * b for a, b in zip(query_embedding, item_embedding))
norm_q = sum(x**2 for x in query_embedding) ** 0.5
norm_i = sum(x**2 for x in item_embedding) ** 0.5
return dot_product / (norm_q * norm_i) if norm_q * norm_i > 0 else 0.0
Comment on lines +505 to +535
Copy link
Copy Markdown

@coderabbitai coderabbitai bot Mar 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Stale cache risk when normalize_embeddings config changes.

The cached _cached_np_embedding is computed based on the value of self.config.normalize_embeddings at cache time. Since EmbeddingAgentConfig is mutable (it's a Pydantic BaseModel without frozen=True), if normalize_embeddings changes after items are cached, subsequent similarity computations will use incorrectly normalized/unnormalized embeddings.

Consider one of:

  1. Include the normalization state in the cache (e.g., _cached_np_embedding_normalized: bool)
  2. Make the config immutable (frozen=True)
  3. Invalidate item caches when config changes
🛡️ Proposed fix: track normalization state with the cache 
-        item_np = getattr(item, "_cached_np_embedding", None)
+        cached = getattr(item, "_cached_np_embedding", None)
+        cached_normalized = getattr(item, "_cached_np_normalized", None)
+        
+        # Invalidate cache if normalization setting changed
+        if cached is not None and cached_normalized != self.config.normalize_embeddings:
+            cached = None
+        item_np = cached

         if item_np is None:
             item_embedding = getattr(item, "embedding", None)
             if item_embedding is None:
                 return None
             if NUMPY_AVAILABLE and query_np is not None:
                 item_np = np.array(item_embedding)
                 if self.config.normalize_embeddings:
                     norm = np.linalg.norm(item_np)
                     if norm > 0:
                         item_np = item_np / norm
                 try:
                     setattr(item, "_cached_np_embedding", item_np)
+                    setattr(item, "_cached_np_normalized", self.config.normalize_embeddings)
                 except AttributeError:
-                    pass  # Skip caching if item is a strict class (e.g., uses slots)
+                    # Item class restricts attribute assignment (e.g., __slots__)
+                    logger.debug("Cannot cache embedding on item: attribute assignment not allowed")
🧰 Tools
🪛 Ruff (0.15.7)

[warning] 522-522: Do not call setattr with a constant attribute value. It is not any safer than normal property access.

Replace setattr with assignment

(B010)

[warning] 532-532: zip() without an explicit strict= parameter

Add explicit value for parameter strict=

(B905)

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@api/embedding_agent/service.py` around lines 505 - 535, The cached numpy
embedding in _compute_similarity uses self.config.normalize_embeddings at cache
time, so if that flag changes later cached embeddings become stale; update the
caching to also record the normalization state (e.g., store a companion
attribute like _cached_np_embedding_normalized alongside _cached_np_embedding)
and when reading the cache in _compute_similarity verify that getattr(item,
"_cached_np_embedding_normalized", None) matches
self.config.normalize_embeddings and, if not, recompute (and overwrite) the
cached embedding using the current self.config.normalize_embeddings; reference
_compute_similarity, _cached_np_embedding, and self.config.normalize_embeddings
when implementing this change.


def _search_knowledge_items(
self,
query_embedding: List[float],
Expand All @@ -526,36 +558,15 @@ def _search_knowledge_items(
query_np = query_np / norm

for item in self._knowledge_items:
# Filter by knowledge type if specified
if knowledge_types:
item_type = getattr(item, "knowledge_type", "general")
if item_type not in [kt.value for kt in knowledge_types]:
continue

# Get item embedding
item_embedding = getattr(item, "embedding", None)
if item_embedding is None:
similarity = self._compute_similarity(item, query_embedding, query_np)
if similarity is None:
continue

# Calculate similarity
if NUMPY_AVAILABLE and query_np is not None:
item_np = np.array(item_embedding)
if self.config.normalize_embeddings:
norm = np.linalg.norm(item_np)
if norm > 0:
item_np = item_np / norm
similarity = float(np.dot(query_np, item_np))
else:
# Simple cosine similarity
dot_product = sum(
a * b for a, b in zip(query_embedding, item_embedding)
)
norm_q = sum(x**2 for x in query_embedding) ** 0.5
norm_i = sum(x**2 for x in item_embedding) ** 0.5
similarity = (
dot_product / (norm_q * norm_i) if norm_q * norm_i > 0 else 0.0
)

# Apply threshold
if similarity >= min_similarity:
match = SimilarityMatch(
Expand Down