agentrun-sdk 0.0.4__py3-none-any.whl

agentrun/agent_runtime/api/data.py

@@ -0,0 +1,98 @@
+"""
+This file is auto generated by the code generation script.
+Do not modify this file manually.
+Use the `make codegen` command to regenerate.
+
+当前文件为自动生成的控制 API 客户端代码。请勿手动修改此文件。
+使用 `make codegen` 命令重新生成。
+
+source: agentrun/agent_runtime/api/__data_async_template.py
+"""
+
+from typing import Iterable, Optional, TypedDict
+
+from openai.types.chat import ChatCompletionMessageParam
+from typing_extensions import Required, Unpack
+
+from agentrun.utils.config import Config
+from agentrun.utils.data_api import DataAPI, ResourceType
+
+
+class InvokeArgs(TypedDict):
+    messages: Required[Iterable[ChatCompletionMessageParam]]
+    stream: Required[bool]
+    config: Optional[Config]
+
+
+class AgentRuntimeDataAPI(DataAPI):
+
+    def __init__(
+        self,
+        agent_runtime_name: str,
+        agent_runtime_endpoint_name: str = "Default",
+        config: Optional[Config] = None,
+    ):
+        super().__init__(
+            resource_name=agent_runtime_name,
+            resource_type=ResourceType.Runtime,
+            namespace=f"agent-runtimes/{agent_runtime_name}/endpoints/{agent_runtime_endpoint_name}/invocations",
+            config=config,
+        )
+
+    async def invoke_openai_async(
+        self,
+        **kwargs: Unpack[InvokeArgs],
+    ):
+        messages = kwargs.get("messages", [])
+        stream = kwargs.get("stream", False)
+        config = kwargs.get("config", None)
+
+        cfg = Config.with_configs(self.config, config)
+        api_base = self.with_path("openai/v1")
+        _, headers, _ = self.auth(headers=cfg.get_headers())
+
+        from httpx import AsyncClient
+        from openai import AsyncOpenAI
+
+        client = AsyncOpenAI(
+            api_key="",
+            base_url=api_base,
+            http_client=AsyncClient(headers=headers),
+        )
+        timeout = cfg.get_timeout()
+
+        return client.chat.completions.create(
+            model=self.resource_name,
+            messages=messages,
+            timeout=timeout,
+            stream=stream,
+        )
+
+    def invoke_openai(
+        self,
+        **kwargs: Unpack[InvokeArgs],
+    ):
+        messages = kwargs.get("messages", [])
+        stream = kwargs.get("stream", False)
+        config = kwargs.get("config", None)
+
+        cfg = Config.with_configs(self.config, config)
+        api_base = self.with_path("openai/v1")
+        _, headers, _ = self.auth(headers=cfg.get_headers())
+
+        from httpx import Client
+        from openai import OpenAI
+
+        client = OpenAI(
+            api_key="",
+            base_url=api_base,
+            http_client=Client(headers=headers),
+        )
+        timeout = cfg.get_timeout()
+
+        return client.chat.completions.create(
+            model=self.resource_name,
+            messages=messages,
+            timeout=timeout,
+            stream=stream,
+        )