embedding_client.py

"""Client for handling embeddings for vector search"""

import os
import logging
from typing import List, Optional
import numpy as np
from openai import AzureOpenAI

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class EmbeddingClient:
    """Client for generating embeddings using Azure OpenAI"""
    
    def __init__(self, azure_endpoint: str, api_key: str, deployment: str, api_version: str = "2023-05-15"):
        """Initialize the embedding client"""
        self.azure_endpoint = azure_endpoint
        self.api_key = api_key
        self.deployment = deployment
        self.api_version = api_version
        
        # Initialize client
        self.client = None
        if self.azure_endpoint and self.api_key and self.deployment:
            try:
                self.client = AzureOpenAI(
                    api_key=self.api_key,
                    api_version=self.api_version,
                    azure_endpoint=self.azure_endpoint
                )
                logger.info(f"Initialized embedding client with deployment {self.deployment}")
            except Exception as e:
                logger.error(f"Failed to initialize Azure OpenAI client: {e}")
                self.client = None
        else:
            logger.warning("Missing configuration for embedding client")
    
    def get_embedding(self, text: str) -> List[float]:
        """Generate embedding for the given text"""
        if not self.client:
            logger.warning("No embedding client available, falling back to mock embedding")
            return self._get_mock_embedding()
        
        try:
            # Truncate text if needed
            max_chars = 32000  # Approximate limit
            if len(text) > max_chars:
                text = text[:max_chars]
                logger.warning(f"Text truncated to {max_chars} characters")
            
            # Call the embeddings API
            response = self.client.embeddings.create(
                input=text,
                model=self.deployment
            )
            
            # Extract the embedding
            embedding = response.data[0].embedding
            logger.info(f"Successfully generated embedding of dimension {len(embedding)}")
            return embedding
        except Exception as e:
            logger.error(f"Error generating embedding: {e}")
            return self._get_mock_embedding()
    
    def _get_mock_embedding(self) -> List[float]:
        """Generate a mock embedding for fallback"""
        logger.warning("Using mock embedding - this is not suitable for production use")
        # Generate a normalized random vector
        embedding = np.random.normal(size=1536)  # Standard dimension for text-embedding-ada-002
        embedding = embedding / np.linalg.norm(embedding)
        return embedding.tolist()