"""Embedding utilities for the brain system.
Integrates with the local embedding server (all-MiniLM-L6-v2) to provide
semantic search capabilities.
Usage:
from embeddings import get_embedding, search_semantic
emb = get_embedding("your query text")
results = search_semantic("query", limit=5)
"""
import json
import sqlite3
import urllib.request
from pathlib import Path
from typing import List, Optional
EMBEDDING_SERVER = "http://localhost:9001/embed"
BRAIN_DB = Path.home() / ".hermes/brain/index/events.db"
def get_embedding(text: str) -> List[float]:
"""Get embedding for a single text from the local server."""
data = json.dumps({"texts": [text]}).encode()
req = urllib.request.Request(
EMBEDDING_SERVER,
data=data,
headers={"Content-Type": "application/json"},
method="POST"
)
with urllib.request.urlopen(req, timeout=30) as response:
result = json.loads(response.read())
return result["embeddings"][0]
def cosine_similarity(emb1: List[float], emb2: List[float]) -> float:
"""Compute cosine similarity between two embeddings."""
dot_product = sum(a * b for a, b in zip(emb1, emb2))
norm1 = sum(a * a for a in emb1) ** 0.5
norm2 = sum(b * b for b in emb2) ** 0.5
if norm1 == 0 or norm2 == 0:
return 0.0
return dot_product / (norm1 * norm2)
def search_semantic(query: str, limit: int = 5, min_score: float = 0.3) -> List[dict]:
"""Search brain entities using semantic similarity.
Args:
query: Search query text
limit: Maximum number of results
min_score: Minimum similarity score (0.0 to 1.0)
Returns:
List of dicts with 'entity_id', 'category', 'content', 'score'
"""
query_emb = get_embedding(query)
# Connect to brain DB
conn = sqlite3.connect(str(BRAIN_DB))
conn.row_factory = sqlite3.Row
cursor = conn.cursor()
# Get active entities
cursor.execute("""
SELECT entity_id, category, content
FROM entities
WHERE deleted = 0 AND content IS NOT NULL AND content != ''
ORDER BY pinned DESC, access_count DESC
LIMIT 500
""")
entities = cursor.fetchall()
conn.close()
# Compute similarities
results = []
for entity in entities:
try:
# Use first 512 chars for embedding (covers most entity content)
entity_text = entity["category"] + " " + entity["content"][:512]
entity_emb = get_embedding(entity_text)
score = cosine_similarity(query_emb, entity_emb)
if score >= min_score:
results.append({
"entity_id": entity["entity_id"],
"category": entity["category"],
"content": entity["content"][:200], # Preview
"score": round(score, 3)
})
except Exception:
continue
# Sort by score descending and limit
results.sort(key=lambda x: x["score"], reverse=True)
return results[:limit]
def ensure_embeddings_table():
"""Create the entity_embeddings table if it doesn't exist."""
conn = sqlite3.connect(str(BRAIN_DB))
cursor = conn.cursor()
cursor.execute("""
CREATE TABLE IF NOT EXISTS entity_embeddings (
entity_id TEXT PRIMARY KEY,
embedding TEXT NOT NULL,
created_at TEXT NOT NULL DEFAULT (datetime('now')),
FOREIGN KEY (entity_id) REFERENCES entities(entity_id)
)
""")
conn.commit()
conn.close()
def store_entity_embedding(entity_id: str, embedding: List[float]) -> bool:
"""Store embedding for an entity in the database.
Args:
entity_id: Entity ID (string)
embedding: 384-dimensional embedding vector
Returns:
True if stored successfully
"""
ensure_embeddings_table()
conn = sqlite3.connect(str(BRAIN_DB))
cursor = conn.cursor()
try:
cursor.execute("""
INSERT OR REPLACE INTO entity_embeddings (entity_id, embedding)
VALUES (?, ?)
""", (entity_id, json.dumps(embedding)))
conn.commit()
return True
except Exception as e:
print(f"Error storing embedding: {e}")
return False
finally:
conn.close()
def get_entity_embedding(entity_id: str) -> Optional[List[float]]:
"""Retrieve stored embedding for an entity.
Args:
entity_id: Entity ID
Returns:
Embedding list or None if not found
"""
conn = sqlite3.connect(str(BRAIN_DB))
conn.row_factory = sqlite3.Row
cursor = conn.cursor()
cursor.execute("""
SELECT embedding FROM entity_embeddings WHERE entity_id = ?
""", (entity_id,))
row = cursor.fetchone()
conn.close()
if row and row["embedding"]:
try:
return json.loads(row["embedding"])
except json.JSONDecodeError:
return None
return None
def batch_embed_entities(entity_ids: List[str]) -> int:
"""Embed multiple entities in batch, storing results in DB.
Args:
entity_ids: List of entity IDs to embed
Returns:
Number of entities successfully embedded
"""
ensure_embeddings_table()
conn = sqlite3.connect(str(BRAIN_DB))
conn.row_factory = sqlite3.Row
cursor = conn.cursor()
placeholders = ",".join(["?"] * len(entity_ids))
cursor.execute(f"""
SELECT entity_id, category, content
FROM entities
WHERE entity_id IN ({placeholders}) AND deleted = 0
""", entity_ids)
entities = cursor.fetchall()
conn.close()
# Build texts for batch embedding
texts = []
entity_map = {}
for entity in entities:
text = entity["category"] + " " + entity["content"][:512]
texts.append(text)
entity_map[text] = entity["entity_id"]
if not texts:
return 0
# Get batch embeddings
data = json.dumps({"texts": texts}).encode()
req = urllib.request.Request(
EMBEDDING_SERVER,
data=data,
headers={"Content-Type": "application/json"},
method="POST"
)
with urllib.request.urlopen(req, timeout=60) as response:
result = json.loads(response.read())
embeddings = result["embeddings"]
# Store in DB
stored = 0
for text, embedding in zip(texts, embeddings):
if store_entity_embedding(entity_map[text], embedding):
stored += 1
return stored
def search_semantic_cached(query: str, limit: int = 5, min_score: float = 0.3) -> List[dict]:
"""Search using cached entity embeddings for speed.
Falls back to on-the-fly computation for entities without cached embeddings.
Args:
query: Search query text
limit: Maximum number of results
min_score: Minimum similarity score (0.0 to 1.0)
Returns:
List of dicts with 'entity_id', 'category', 'content', 'score'
"""
query_emb = get_embedding(query)
# Check how many cached embeddings we have
conn = sqlite3.connect(str(BRAIN_DB))
conn.row_factory = sqlite3.Row
cursor = conn.cursor()
# Try cached embeddings first
cursor.execute("""
SELECT e.entity_id, e.category, e.content, ee.embedding
FROM entities e
INNER JOIN entity_embeddings ee ON e.entity_id = ee.entity_id
WHERE e.deleted = 0
ORDER BY e.pinned DESC, e.access_count DESC
LIMIT 500
""")
cached_rows = cursor.fetchall()
# If we have cached embeddings, use them
if cached_rows:
results = []
for row in cached_rows:
try:
entity_emb = json.loads(row["embedding"])
score = cosine_similarity(query_emb, entity_emb)
if score >= min_score:
results.append({
"entity_id": row["entity_id"],
"category": row["category"],
"content": row["content"][:200],
"score": round(score, 3)
})
except (json.JSONDecodeError, KeyError):
continue
conn.close()
results.sort(key=lambda x: x["score"], reverse=True)
return results[:limit]
conn.close()
# Fallback: on-the-fly computation
return search_semantic(query, limit, min_score)