ic

packages/azure-gpt/banortegpt/generation/azure_gpt/__init__.py

@@ -0,0 +1,61 @@
+from typing import Protocol, cast
+
+from openai import AsyncAzureOpenAI
+from openai.types.chat import ChatCompletion
+
+
+class AsyncGPT:
+    def __init__(self, azure_endpoint: str, api_key: str, api_version: str) -> None:
+        self.client = AsyncAzureOpenAI(
+            azure_endpoint=azure_endpoint, api_key=api_key, api_version=api_version
+        )
+
+    async def generate(self, messages, model, **kwargs):
+        response = await self.client.chat.completions.create(
+            messages=messages, model=model, **kwargs
+        )
+
+        response = cast(ChatCompletion, response)
+
+        return response.choices[0].message
+
+    async def stream(self, messages, model, **kwargs):
+        response = await self.client.chat.completions.create(
+            messages=messages, model=model, stream=True, **kwargs
+        )
+
+        async for chunk in response:
+            if choices := chunk.choices:
+                yield choices[0].delta
+
+    @staticmethod
+    def build_tool_call(tool_id: str, tool_name: str, tool_buffer: str):
+        tool_call = {
+            "tool_calls": [
+                {
+                    "id": tool_id,
+                    "function": {
+                        "name": tool_name,
+                        "arguments": tool_buffer,
+                    },
+                    "type": "function",
+                }
+            ]
+        }
+
+        return tool_call
+
+    @staticmethod
+    def build_tool_call_id(tool_id: str):
+        return {"tool_call_id": tool_id}
+
+    class Config(Protocol):
+        azure_endpoint: str
+        openai_api_key: str
+        openai_api_version: str
+
+    @classmethod
+    def from_config(cls, config: Config):
+        return cls(
+            config.azure_endpoint, config.openai_api_key, config.openai_api_version
+        )