File size: 2,941 Bytes
fe52ef9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | """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()
|