# llm_summarizer/embeddings.py

import os
from openai import AzureOpenAI


class AzureEmbeddingClient:

    def __init__(
        self,
        azure_endpoint=None,
        api_key=None,
        api_version="2024-02-15-preview",
        deployment_name="text-embedding-3-large"
    ):
        self.azure_endpoint = azure_endpoint or os.getenv("AZURE_OPENAI_ENDPOINT")
        self.api_key = api_key or os.getenv("AZURE_OPENAI_API_KEY")
        self.deployment_name = deployment_name

        if not self.azure_endpoint or not self.api_key:
            raise ValueError("Azure endpoint and API key must be provided")

        self.client = AzureOpenAI(
            api_key=self.api_key,
            api_version=api_version,
            azure_endpoint=self.azure_endpoint
        )

    def embed(self, texts):
        response = self.client.embeddings.create(
            model=self.deployment_name,
            input=texts
        )
        return [r.embedding for r in response.data]