"""Vector database integration using ChromaDB for semantic search."""
import os
from pathlib import Path
from typing import Dict, List, Optional, Any, Tuple
import logging
from datetime import datetime
import chromadb
from chromadb.config import Settings
from chromadb.utils import embedding_functions
import numpy as np
from cachetools import TTLCache
from .models import Expert, ExpertEmbedding
logger = logging.getLogger(__name__)
class VectorDatabaseManager:
"""Manage vector database operations with ChromaDB."""
def __init__(
self,
persist_path: Optional[Path] = None,
embedding_model: str = "all-MiniLM-L6-v2",
cache_ttl: int = 3600
):
# Set up persistence path
self.persist_path = persist_path or Path("./chroma_db")
self.persist_path.mkdir(parents=True, exist_ok=True)
# Initialize ChromaDB client
self.client = chromadb.PersistentClient(
path=str(self.persist_path),
settings=Settings(
anonymized_telemetry=False,
allow_reset=True
)
)
# Set up embedding function
self.embedding_function = embedding_functions.SentenceTransformerEmbeddingFunction(
model_name=embedding_model
)
# Initialize collections
self.collections = {}
self._init_collections()
# Cache for embeddings
self.embedding_cache = TTLCache(maxsize=1000, ttl=cache_ttl)
def _init_collections(self):
"""Initialize vector collections."""
collection_names = [
"expert-descriptions",
"expert-patterns",
"expert-constraints",
"expert-technologies",
"task-history"
]
for name in collection_names:
try:
self.collections[name] = self.client.get_or_create_collection(
name=name,
embedding_function=self.embedding_function,
metadata={"hnsw:space": "cosine"}
)
except Exception as e:
logger.error(f"Failed to create collection {name}: {e}")
async def index_expert(self, expert: Expert) -> bool:
"""Index an expert in vector database."""
try:
expert_id = expert.id
# Index description
self.collections["expert-descriptions"].upsert(
ids=[expert_id],
documents=[expert.description],
metadatas=[{
"expert_name": expert.name,
"version": expert.version,
"domains": ",".join(expert.domains),
"updated_at": datetime.now().isoformat()
}]
)
# Index patterns
if expert.patterns:
pattern_text = "\n".join(expert.patterns)
self.collections["expert-patterns"].upsert(
ids=[f"{expert_id}_patterns"],
documents=[pattern_text],
metadatas=[{"expert_id": expert_id}]
)
# Index constraints
if expert.constraints:
constraint_text = "\n".join(expert.constraints)
self.collections["expert-constraints"].upsert(
ids=[f"{expert_id}_constraints"],
documents=[constraint_text],
metadatas=[{"expert_id": expert_id}]
)
# Index technologies
tech_descriptions = []
for spec in expert.specializations:
tech_desc = f"{spec.technology}: {', '.join(spec.frameworks)} ({spec.expertise_level})"
tech_descriptions.append(tech_desc)
if tech_descriptions:
self.collections["expert-technologies"].upsert(
ids=[f"{expert_id}_tech"],
documents=["\n".join(tech_descriptions)],
metadatas=[{
"expert_id": expert_id,
"technologies": ",".join([s.technology for s in expert.specializations])
}]
)
# Clear cache for this expert
cache_keys_to_remove = [k for k in self.embedding_cache.keys() if expert_id in str(k)]
for key in cache_keys_to_remove:
del self.embedding_cache[key]
return True
except Exception as e:
logger.error(f"Failed to index expert {expert.id}: {e}")
return False
async def search_experts(
self,
query: str,
search_mode: str = "hybrid",
limit: int = 5
) -> List[Tuple[str, float, Dict[str, Any]]]:
"""Search for experts using semantic similarity.
Returns list of (expert_id, score, metadata) tuples.
"""
cache_key = f"search:{query}:{search_mode}:{limit}"
if cache_key in self.embedding_cache:
return self.embedding_cache[cache_key]
results = []
if search_mode in ["description", "hybrid"]:
desc_results = self.collections["expert-descriptions"].query(
query_texts=[query],
n_results=limit * 2 if search_mode == "hybrid" else limit
)
for i, expert_id in enumerate(desc_results['ids'][0]):
score = 1.0 - desc_results['distances'][0][i] # Convert distance to similarity
metadata = desc_results['metadatas'][0][i]
results.append((expert_id, score * 0.4, metadata))
if search_mode in ["patterns", "hybrid"]:
pattern_results = self.collections["expert-patterns"].query(
query_texts=[query],
n_results=limit
)
for i, doc_id in enumerate(pattern_results['ids'][0]):
expert_id = pattern_results['metadatas'][0][i]['expert_id']
score = 1.0 - pattern_results['distances'][0][i]
results.append((expert_id, score * 0.3, {}))
if search_mode in ["constraints", "hybrid"]:
constraint_results = self.collections["expert-constraints"].query(
query_texts=[query],
n_results=limit
)
for i, doc_id in enumerate(constraint_results['ids'][0]):
expert_id = constraint_results['metadatas'][0][i]['expert_id']
score = 1.0 - constraint_results['distances'][0][i]
results.append((expert_id, score * 0.3, {}))
# Combine scores for same expert
expert_scores = {}
expert_metadata = {}
for expert_id, score, metadata in results:
if expert_id not in expert_scores:
expert_scores[expert_id] = 0
expert_metadata[expert_id] = metadata
expert_scores[expert_id] += score
# Sort by combined score
sorted_results = [
(expert_id, score, expert_metadata.get(expert_id, {}))
for expert_id, score in sorted(
expert_scores.items(),
key=lambda x: x[1],
reverse=True
)
][:limit]
self.embedding_cache[cache_key] = sorted_results
return sorted_results
async def find_similar_experts(
self,
expert_id: str,
similarity_type: str = "overall",
limit: int = 3
) -> List[Tuple[str, float]]:
"""Find experts similar to a given expert.
Returns list of (expert_id, similarity_score) tuples.
"""
cache_key = f"similar:{expert_id}:{similarity_type}:{limit}"
if cache_key in self.embedding_cache:
return self.embedding_cache[cache_key]
collection_map = {
"overall": "expert-descriptions",
"patterns": "expert-patterns",
"technology": "expert-technologies"
}
collection_name = collection_map.get(similarity_type, "expert-descriptions")
collection = self.collections[collection_name]
# Get the embedding for the source expert
doc_id = expert_id
if similarity_type != "overall":
doc_id = f"{expert_id}_{similarity_type}"
try:
# For "overall" similarity, query using document text instead of embeddings
if similarity_type == "overall":
# Get the source expert's description text for query
source_doc = collection.get(ids=[expert_id], include=['documents'])
if not source_doc['documents'] or not source_doc['documents'][0]:
logger.warning(f"No document found for expert {expert_id}")
return []
source_text = source_doc['documents'][0]
# Query using text similarity
results = collection.query(
query_texts=[source_text],
n_results=limit + 1, # Get extra to exclude self
)
else:
# Get the embedding for the source expert
source_embedding = self._get_expert_embedding(collection, doc_id)
if source_embedding is None:
logger.warning(f"No embedding found for expert {expert_id}")
return []
# Query for similar items, excluding the source
results = collection.query(
query_embeddings=[source_embedding],
n_results=limit + 1, # Get extra to exclude self
where={"expert_id": {"$ne": expert_id}}
)
similar_experts = []
for i, result_id in enumerate(results['ids'][0]):
if result_id != expert_id: # Exclude self
similarity = 1.0 - results['distances'][0][i]
actual_expert_id = result_id
# Extract expert ID from compound IDs
if "_" in result_id and similarity_type != "overall":
actual_expert_id = results['metadatas'][0][i].get('expert_id', result_id)
similar_experts.append((actual_expert_id, similarity))
if len(similar_experts) >= limit:
break
self.embedding_cache[cache_key] = similar_experts
return similar_experts
except Exception as e:
logger.error(f"Failed to find similar experts: {e}")
return []
def _get_expert_embedding(self, collection: Any, doc_id: str) -> Optional[List[float]]:
"""Get embedding for a document from collection."""
try:
result = collection.get(ids=[doc_id], include=['embeddings'])
# Check if we have embeddings in the result
embeddings = result.get('embeddings', [])
if isinstance(embeddings, list) and len(embeddings) > 0 and embeddings[0] is not None:
embedding = embeddings[0]
# Convert numpy array to list if needed
if hasattr(embedding, 'tolist'):
return embedding.tolist()
return embedding
except Exception as e:
logger.error(f"Error getting embedding for {doc_id}: {e}")
return None
async def add_task_history(
self,
task_id: str,
task_description: str,
selected_expert_id: str,
success: bool,
metadata: Optional[Dict] = None
):
"""Add task to history for collaborative filtering."""
self.collections["task-history"].upsert(
ids=[task_id],
documents=[task_description],
metadatas=[{
"selected_expert": selected_expert_id,
"success": str(success),
"timestamp": datetime.now().isoformat(),
**(metadata or {})
}]
)
async def get_expert_embeddings(self, expert_id: str) -> Optional[ExpertEmbedding]:
"""Get all embeddings for an expert."""
embeddings = {}
# Get description embedding
desc_result = self.collections["expert-descriptions"].get(
ids=[expert_id],
include=['embeddings']
)
if desc_result['embeddings']:
embeddings['description'] = desc_result['embeddings'][0]
# Get pattern embedding
pattern_result = self.collections["expert-patterns"].get(
ids=[f"{expert_id}_patterns"],
include=['embeddings']
)
if pattern_result['embeddings']:
embeddings['patterns'] = pattern_result['embeddings'][0]
# Get constraint embedding
constraint_result = self.collections["expert-constraints"].get(
ids=[f"{expert_id}_constraints"],
include=['embeddings']
)
if constraint_result['embeddings']:
embeddings['constraints'] = constraint_result['embeddings'][0]
# Get technology embedding
tech_result = self.collections["expert-technologies"].get(
ids=[f"{expert_id}_tech"],
include=['embeddings']
)
if tech_result['embeddings']:
embeddings['technologies'] = tech_result['embeddings'][0]
if not embeddings:
return None
# Fill missing embeddings with zeros
embedding_dim = len(next(iter(embeddings.values())))
default_embedding = [0.0] * embedding_dim
return ExpertEmbedding(
expert_id=expert_id,
embeddings={
'description': embeddings.get('description', default_embedding),
'domains': embeddings.get('description', default_embedding), # Reuse description
'technologies': embeddings.get('technologies', default_embedding),
'patterns': embeddings.get('patterns', default_embedding),
'constraints': embeddings.get('constraints', default_embedding)
},
metadata={
'model': 'all-MiniLM-L6-v2',
'timestamp': datetime.now(),
'version': '1.0.0'
}
)
async def reset_database(self):
"""Reset all collections (useful for testing)."""
try:
self.client.reset()
self._init_collections()
self.embedding_cache.clear()
logger.info("Vector database reset successfully")
except Exception as e:
logger.error(f"Failed to reset vector database: {e}")
async def get_collection_stats(self) -> Dict[str, int]:
"""Get statistics about collections."""
stats = {}
for name, collection in self.collections.items():
try:
count = collection.count()
stats[name] = count
except Exception as e:
logger.error(f"Failed to get stats for {name}: {e}")
stats[name] = -1
return stats
