ic

packages/azure-ada/banortegpt/embedding/azure_ada/nonblocking.py

@@ -0,0 +1,47 @@
+from langfuse.openai import AsyncAzureOpenAI
+from openai.types.embedding import Embedding
+
+from .base import BaseAda
+
+
+class AsyncAda(BaseAda):
+    def __init__(
+        self, model: str | None = None, *, endpoint: str, key: str, version: str
+    ) -> None:
+        super().__init__(model, endpoint=endpoint, key=key, version=version)
+        self.client = AsyncAzureOpenAI(
+            azure_endpoint=endpoint, api_key=key, api_version=version
+        )
+
+    async def embed(
+        self, input: str | list[str], *, model: str | None = None
+    ) -> list[float] | list[list[float]]:
+        if isinstance(input, str):
+            return await self.embed_query(input, model)
+        else:
+            return await self.batch_embed(input, model)
+
+    async def batch_embed(
+        self, texts: list[str], model: str | None = None
+    ) -> list[list[float]]:
+        if model is None:
+            if self.model is None:
+                raise ValueError("No embedding model set")
+            model = self.model
+
+        batches = [texts[i : i + 2048] for i in range(0, len(texts), 2048)]
+        results = [
+            (await self.client.embeddings.create(input=batch, model=model)).data
+            for batch in batches
+        ]
+        flattened_results: list[Embedding] = sum(results, [])
+        return [result.embedding for result in flattened_results]
+
+    async def embed_query(self, text: str, model: str | None = None) -> list[float]:
+        if model is None:
+            if self.model is None:
+                raise ValueError("No embedding model set")
+            model = self.model
+
+        response = await self.client.embeddings.create(input=text, model=model)
+        return response.data[0].embedding