opt

configs/llms/text_encoder.yml

@@ -1,3 +1,3 @@
-name: OpenaiTextEmbeddingV3
+name: AzureTextEmbeddingV3
 endpoint: ${env| custom_openai_endpoint, https://api.openai.com/v1}
 api_key: ${env| custom_openai_key, openai_api_key}

omagent_core/models/encoders/azure_encoder.py

@@ -0,0 +1,32 @@
+from typing import Any, List
+
+from openai import AsyncAzureOpenAI, AzureOpenAI
+
+from ...utils.registry import registry
+from .base import EncoderBase
+
+
+@registry.register_encoder()
+class AzureTextEmbeddingV3(EncoderBase):
+    model_id: str = "text-embedding-3-large"
+    api_key: str
+    dim: int = 3072
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        protected_namespaces = ()
+        extra = "allow"
+
+    def __init__(self, /, **data: Any) -> None:
+        super().__init__(**data)
+        self.client = AzureOpenAI(base_url=self.endpoint, api_key=self.api_key)
+        self.aclient = AsyncAzureOpenAI(base_url=self.endpoint, api_key=self.api_key)
+
+    def _infer(self, data: List[str], **kwargs) -> List[List[float]]:
+        res = self.client.embeddings.create(input=data, model=self.model_id)
+        return [item.embedding for item in res.data]
+
+    async def _ainfer(self, data: List[str], **kwargs) -> List[List[float]]:
+        res = await self.aclient.embeddings.create(input=data, model=self.model_id)
+        return [item.embedding for item in res.data]

webpage_configs/llms/text_encoder.yml

@@ -1,3 +1,3 @@
-name: OpenaiTextEmbeddingV3
+name: AzureTextEmbeddingV3
 endpoint: ${env| custom_openai_endpoint, https://api.openai.com/v1}
 api_key: ${env| custom_openai_key, openai_api_key}