ruby2210/rag-chatbot

1"""2Embedding utilities for text processing in the RAG Chatbot application.3Handles text embedding generation and processing.4"""5import cohere6from typing import List, Optional7from .config import settings8import numpy as np9 10 11class EmbeddingUtils:12    """13    Utility class for generating and processing text embeddings.14    """15    def __init__(self):16        self.client = cohere.Client(api_key=settings.COHERE_API_KEY)17 18    def create_embeddings(self, texts: List[str]) -> Optional[List[List[float]]]:19        """20        Create embeddings for a list of texts using Cohere's embedding API.21        """22        try:23            # Use Cohere's embed API24            response = self.client.embed(25                texts=texts,26                model="embed-english-v3.0",  # Using Cohere's standard embedding model27                input_type="search_document"  # Specify the input type for better embeddings28            )29 30            # Extract embeddings from the response31            embeddings = [embedding for embedding in response.embeddings]32 33            return embeddings34        except Exception as e:35            print(f"Error creating embeddings: {e}")36            return None37 38    def cosine_similarity(self, vec1: List[float], vec2: List[float]) -> float:39        """40        Calculate cosine similarity between two embedding vectors.41        """42        try:43            # Convert to numpy arrays for calculation44            v1 = np.array(vec1)45            v2 = np.array(v2)46 47            # Calculate cosine similarity48            dot_product = np.dot(v1, v2)49            norm_v1 = np.linalg.norm(v1)50            norm_v2 = np.linalg.norm(v2)51 52            if norm_v1 == 0 or norm_v2 == 0:53                return 0.054 55            return float(dot_product / (norm_v1 * norm_v2))56        except Exception as e:57            print(f"Error calculating cosine similarity: {e}")58            return 0.059 60    def chunk_text(self, text: str, chunk_size: int = 1000, overlap: int = 100) -> List[str]:61        """62        Split text into overlapping chunks for embedding.63        """64        if len(text) <= chunk_size:65            return [text]66 67        chunks = []68        start = 069 70        while start < len(text):71            end = start + chunk_size72 73            # If this is not the last chunk, try to break at sentence boundary74            if end < len(text):75                # Look for sentence endings near the end76                chunk = text[start:end]77                last_sentence_end = max(78                    chunk.rfind('. '),79                    chunk.rfind('?'),80                    chunk.rfind('!'),81                    chunk.rfind('\n')82                )83 84                if last_sentence_end > chunk_size // 2:  # Only if sentence end is reasonably far in85                    end = start + last_sentence_end + 186 87            chunk_text = text[start:end]88            chunks.append(chunk_text)89 90            # Move start forward, with overlap91            start = end - overlap if end < len(text) else end92 93        return chunks94 95 96# Create a singleton instance97embedding_utils = EmbeddingUtils()