Comprehensive comparison for Embeddings technology in AI applications

import cohere
import numpy as np
from typing import List, Dict, Optional
import os
from dataclasses import dataclass
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

@dataclass
class SearchResult:
    text: str
    score: float
    metadata: Dict

class SemanticSearchEngine:
    """Production-grade semantic search using Cohere Embed API"""
    
    def __init__(self, api_key: Optional[str] = None):
        self.api_key = api_key or os.getenv('COHERE_API_KEY')
        if not self.api_key:
            raise ValueError("Cohere API key must be provided or set in COHERE_API_KEY env variable")
        
        self.client = cohere.Client(self.api_key)
        self.model = 'embed-english-v3.0'
        self.input_type_search = 'search_document'
        self.input_type_query = 'search_query'
        
    def embed_documents(self, texts: List[str]) -> np.ndarray:
        """Embed a batch of documents for indexing"""
        try:
            if not texts:
                raise ValueError("texts list cannot be empty")
            
            response = self.client.embed(
                texts=texts,
                model=self.model,
                input_type=self.input_type_search,
                truncate='END'
            )
            
            embeddings = np.array(response.embeddings)
            logger.info(f"Successfully embedded {len(texts)} documents")
            return embeddings
            
        except cohere.CohereError as e:
            logger.error(f"Cohere API error during document embedding: {e}")
            raise
        except Exception as e:
            logger.error(f"Unexpected error during document embedding: {e}")
            raise
    
    def embed_query(self, query: str) -> np.ndarray:
        """Embed a search query"""
        try:
            if not query or not query.strip():
                raise ValueError("Query cannot be empty")
            
            response = self.client.embed(
                texts=[query],
                model=self.model,
                input_type=self.input_type_query,
                truncate='END'
            )
            
            embedding = np.array(response.embeddings[0])
            logger.info(f"Successfully embedded query")
            return embedding
            
        except cohere.CohereError as e:
            logger.error(f"Cohere API error during query embedding: {e}")
            raise
        except Exception as e:
            logger.error(f"Unexpected error during query embedding: {e}")
            raise
    
    def cosine_similarity(self, vec1: np.ndarray, vec2: np.ndarray) -> float:
        """Calculate cosine similarity between two vectors"""
        return np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))
    
    def search(self, query: str, documents: List[Dict], top_k: int = 5) -> List[SearchResult]:
        """Perform semantic search on documents"""
        try:
            if not documents:
                logger.warning("No documents provided for search")
                return []
            
            doc_texts = [doc['text'] for doc in documents]
            doc_embeddings = self.embed_documents(doc_texts)
            query_embedding = self.embed_query(query)
            
            similarities = []
            for idx, doc_emb in enumerate(doc_embeddings):
                score = self.cosine_similarity(query_embedding, doc_emb)
                similarities.append((idx, score))
            
            similarities.sort(key=lambda x: x[1], reverse=True)
            top_results = similarities[:top_k]
            
            results = [
                SearchResult(
                    text=documents[idx]['text'],
                    score=float(score),
                    metadata=documents[idx].get('metadata', {})
                )
                for idx, score in top_results
            ]
            
            logger.info(f"Search completed. Found {len(results)} results")
            return results
            
        except Exception as e:
            logger.error(f"Error during search: {e}")
            raise

if __name__ == "__main__":
    search_engine = SemanticSearchEngine()
    
    documents = [
        {"text": "Python is a high-level programming language", "metadata": {"id": 1}},
        {"text": "Machine learning models require training data", "metadata": {"id": 2}},
        {"text": "Natural language processing uses embeddings", "metadata": {"id": 3}}
    ]
    
    results = search_engine.search("What is NLP?", documents, top_k=2)
    
    for result in results:
        print(f"Score: {result.score:.4f} - {result.text}")