# MCP Vectorizer Module """ Embedding generation and similarity computation for MCP integration. Features: - Text embedding generation - Batch embedding - Similarity computation - Embedding caching """ from typing import List, Dict, Optional, Tuple, Any import numpy as np from functools import lru_cache # Embedding cache _embedding_cache: Dict[str, List[float]] = {} _cache_max_size = 1000 def generate_embeddings(text: str, use_cache: bool = True) -> Dict: """ Generate embeddings for text. Args: text: Input text to embed use_cache: Whether to use cached embeddings Returns: Embedding vector and metadata """ global _embedding_cache if not text or not text.strip(): return { "success": False, "error": "Empty text provided", "embedding": None } # Check cache cache_key = text[:500] # Use first 500 chars as key if use_cache and cache_key in _embedding_cache: return { "success": True, "embedding": _embedding_cache[cache_key], "cached": True, "dimensions": len(_embedding_cache[cache_key]) } try: from core.llm import embed_text embedding = embed_text(text) if embedding is None: return { "success": False, "error": "Embedding generation failed", "embedding": None } # Convert to list if numpy array if isinstance(embedding, np.ndarray): embedding = embedding.tolist() # Cache the result if use_cache: if len(_embedding_cache) >= _cache_max_size: # Remove oldest entries oldest = list(_embedding_cache.keys())[:100] for k in oldest: del _embedding_cache[k] _embedding_cache[cache_key] = embedding return { "success": True, "embedding": embedding, "cached": False, "dimensions": len(embedding) } except Exception as e: return { "success": False, "error": str(e), "embedding": None } def batch_embed( texts: List[str], batch_size: int = 10, use_cache: bool = True ) -> Dict: """ Generate embeddings for multiple texts. Args: texts: List of texts to embed batch_size: Number of texts per batch use_cache: Whether to use cached embeddings Returns: List of embeddings with metadata """ if not texts: return { "success": False, "error": "No texts provided", "embeddings": [] } embeddings = [] errors = [] cached_count = 0 for i, text in enumerate(texts): result = generate_embeddings(text, use_cache=use_cache) if result["success"]: embeddings.append({ "index": i, "text_preview": text[:50], "embedding": result["embedding"], "cached": result.get("cached", False) }) if result.get("cached"): cached_count += 1 else: embeddings.append({ "index": i, "text_preview": text[:50], "embedding": None, "error": result.get("error") }) errors.append(f"Text {i}: {result.get('error')}") return { "success": len(errors) == 0, "embeddings": embeddings, "total": len(texts), "successful": len(texts) - len(errors), "cached_hits": cached_count, "errors": errors if errors else None } def compute_similarity( embedding1: List[float], embedding2: List[float], method: str = "cosine" ) -> Dict: """ Compute similarity between two embeddings. Args: embedding1: First embedding vector embedding2: Second embedding vector method: "cosine", "euclidean", "dot" Returns: Similarity score and metadata """ try: vec1 = np.array(embedding1) vec2 = np.array(embedding2) if len(vec1) != len(vec2): return { "success": False, "error": f"Dimension mismatch: {len(vec1)} vs {len(vec2)}", "similarity": None } if method == "cosine": # Cosine similarity dot = np.dot(vec1, vec2) norm1 = np.linalg.norm(vec1) norm2 = np.linalg.norm(vec2) if norm1 == 0 or norm2 == 0: similarity = 0.0 else: similarity = dot / (norm1 * norm2) elif method == "euclidean": # Euclidean distance (converted to similarity) distance = np.linalg.norm(vec1 - vec2) similarity = 1.0 / (1.0 + distance) elif method == "dot": # Dot product similarity = float(np.dot(vec1, vec2)) else: return { "success": False, "error": f"Unknown method: {method}", "similarity": None } return { "success": True, "similarity": float(similarity), "method": method, "dimensions": len(vec1) } except Exception as e: return { "success": False, "error": str(e), "similarity": None } def find_most_similar( query_embedding: List[float], candidate_embeddings: List[Dict], top_k: int = 5, threshold: float = 0.0 ) -> Dict: """ Find most similar embeddings from candidates. Args: query_embedding: Query vector candidate_embeddings: List of {"id": str, "embedding": List[float]} top_k: Number of top results threshold: Minimum similarity threshold Returns: Top k most similar candidates """ try: query_vec = np.array(query_embedding) similarities = [] for candidate in candidate_embeddings: cand_vec = np.array(candidate["embedding"]) # Cosine similarity dot = np.dot(query_vec, cand_vec) norm_q = np.linalg.norm(query_vec) norm_c = np.linalg.norm(cand_vec) if norm_q > 0 and norm_c > 0: sim = dot / (norm_q * norm_c) else: sim = 0.0 if sim >= threshold: similarities.append({ "id": candidate.get("id", "unknown"), "similarity": float(sim), "metadata": candidate.get("metadata", {}) }) # Sort by similarity descending similarities.sort(key=lambda x: x["similarity"], reverse=True) return { "success": True, "results": similarities[:top_k], "total_candidates": len(candidate_embeddings), "above_threshold": len(similarities) } except Exception as e: return { "success": False, "error": str(e), "results": [] } def compute_centroid(embeddings: List[List[float]]) -> Dict: """ Compute centroid of multiple embeddings. Args: embeddings: List of embedding vectors Returns: Centroid vector """ try: if not embeddings: return { "success": False, "error": "No embeddings provided", "centroid": None } vectors = np.array(embeddings) centroid = np.mean(vectors, axis=0) return { "success": True, "centroid": centroid.tolist(), "num_vectors": len(embeddings), "dimensions": len(centroid) } except Exception as e: return { "success": False, "error": str(e), "centroid": None } def clear_embedding_cache() -> Dict: """Clear the embedding cache""" global _embedding_cache count = len(_embedding_cache) _embedding_cache = {} return { "success": True, "cleared_entries": count } def get_cache_stats() -> Dict: """Get embedding cache statistics""" global _embedding_cache return { "entries": len(_embedding_cache), "max_size": _cache_max_size, "utilization": len(_embedding_cache) / _cache_max_size }