| 23 | logger = logging.getLogger(__name__) |
| 24 | |
| 25 | class EmbeddingProcessor: |
| 26 | def __init__(self, embedding_size: str = "small"): |
| 27 | """Initialize the embedding processor. |
| 28 | |
| 29 | Args: |
| 30 | embedding_size: Size of embedding model to use ("small" or "large") |
| 31 | """ |
| 32 | self.embedding_size = embedding_size.lower() |
| 33 | if self.embedding_size not in ["small", "large"]: |
| 34 | raise ValueError("embedding_size must be 'small' or 'large'") |
| 35 | |
| 36 | # Set model based on size |
| 37 | self.model = f"text-embedding-3-{self.embedding_size}" |
| 38 | logger.info(f"Using embedding model: {self.model}") |
| 39 | |
| 40 | async def get_embedding_async(self, text: str) -> list[float]: |
| 41 | """ |
| 42 | Get embedding for a text using the embedding API. |
| 43 | |
| 44 | Args: |
| 45 | text: Text to embed |
| 46 | |
| 47 | Returns: |
| 48 | List of floats representing the embedding vector |
| 49 | """ |
| 50 | try: |
| 51 | # Use the get_embedding function from core/embedding.py |
| 52 | # Using azure_openai provider with selected model size |
| 53 | embedding = await get_embedding( |
| 54 | text=text, |
| 55 | provider="azure_openai", # Use Azure OpenAI |
| 56 | model=self.model, |
| 57 | timeout=60 # Increase timeout for Azure |
| 58 | ) |
| 59 | |
| 60 | if embedding and isinstance(embedding, list): |
| 61 | return embedding |
| 62 | else: |
| 63 | logger.error(f"Invalid embedding response: {embedding}") |
| 64 | return [] |
| 65 | |
| 66 | except Exception as e: |
| 67 | logger.error(f"Error getting embedding: {e}") |
| 68 | return [] |
| 69 | |
| 70 | def get_embedding(self, text: str) -> list[float]: |
| 71 | """ |
| 72 | Synchronous wrapper for get_embedding_async. |
| 73 | |
| 74 | Args: |
| 75 | text: Text to embed |
| 76 | |
| 77 | Returns: |
| 78 | List of floats representing the embedding vector |
| 79 | """ |
| 80 | return asyncio.run(self.get_embedding_async(text)) |
| 81 | |
| 82 | def create_embedding_text(self, store_data: dict[str, Any]) -> str: |