Switch LLM provider to z.ai GLM

src/llm_trade_bot/zai_llm.py

@@ -0,0 +1,148 @@
+"""z.ai adapter for the provider-neutral micro-agents interface."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import orjson
+from micro_agents import OpenAIBasedLLM
+from micro_agents.conversation import (
+    Conversation,
+    LlmMessage,
+    Message,
+    ReasoningMessage,
+    ToolCallMessage,
+    ToolCallResponseMessage,
+    UserMessage,
+)
+import micro_agents.tools as tools_package
+
+
+class ZaiLLM(OpenAIBasedLLM):
+    """Use z.ai's OpenAI-compatible Chat Completions endpoint.
+
+    ``OpenAIBasedLLM`` in micro-agents targets OpenAI's Responses API, which
+    z.ai does not expose.  This adapter retains its tool-schema inference while
+    translating the agent conversation to the Chat Completions format.
+    """
+
+    async def ask(
+        self,
+        conversation: Conversation,
+        tools: list[tools_package.Tool],
+        *,
+        instructions: str | None = None,
+        thinking: bool | None = None,
+    ) -> list[Message]:
+        extra_body: dict[str, object] | None = None
+        if thinking is not None:
+            extra_body = {"thinking": {"type": "enabled" if thinking else "disabled"}}
+        response = await self._ai.chat.completions.create(
+            model=self._model,
+            messages=self._serialize_conversation(conversation, instructions),
+            tools=self._serialize_chat_tools(tools),
+            stream=False,
+            extra_body=extra_body,
+        )
+        if not response.choices:
+            raise ValueError("z.ai returned no choices")
+        message = response.choices[0].message
+        tool_calls = message.tool_calls or []
+        out: list[Message] = []
+
+        # GLM returns this field when thinking is enabled.  Keeping it in the
+        # conversation lets the next call include it with the tool invocation.
+        reasoning = getattr(message, "reasoning_content", None)
+        if isinstance(reasoning, str) and reasoning:
+            out.append(ReasoningMessage(reasoning))
+
+        # A Chat Completions response with tool calls is represented by one
+        # assistant message.  The content is normally empty in that case; do
+        # not create a second assistant message before the tool-call message.
+        if not tool_calls and isinstance(message.content, str) and message.content:
+            out.append(LlmMessage(message.content))
+        for call in tool_calls:
+            function = call.function
+            if not isinstance(call.id, str) or not isinstance(function.name, str):
+                raise ValueError("z.ai returned an invalid function call")
+            try:
+                arguments = orjson.loads(function.arguments)
+            except orjson.JSONDecodeError as error:
+                raise ValueError("z.ai returned invalid JSON function arguments") from error
+            if not isinstance(arguments, dict):
+                raise ValueError("z.ai function-call arguments must be an object")
+            out.append(ToolCallMessage(call.id, function.name, arguments))
+        return out
+
+    def _serialize_chat_tools(self, tools: list[tools_package.Tool]) -> list[dict[str, object]]:
+        result: list[dict[str, object]] = []
+        for tool in tools:
+            schema = self._serialize_tool(tool)
+            result.append(
+                {
+                    "type": "function",
+                    "function": {
+                        "name": schema["name"],
+                        "description": schema["description"],
+                        "parameters": schema["parameters"],
+                    },
+                }
+            )
+        return result
+
+    @staticmethod
+    def _serialize_conversation(conversation: Conversation, instructions: str | None) -> list[dict[str, Any]]:
+        messages: list[dict[str, Any]] = []
+        if instructions:
+            messages.append({"role": "system", "content": instructions})
+
+        items = conversation.messages
+        index = 0
+        pending_reasoning: str | None = None
+        while index < len(items):
+            item = items[index]
+            if isinstance(item, ReasoningMessage):
+                pending_reasoning = item.text()
+                index += 1
+                continue
+            if isinstance(item, UserMessage):
+                messages.append({"role": "user", "content": item.message()})
+            elif isinstance(item, LlmMessage):
+                assistant: dict[str, Any] = {"role": "assistant", "content": item.message()}
+                if pending_reasoning:
+                    assistant["reasoning_content"] = pending_reasoning
+                    pending_reasoning = None
+                messages.append(assistant)
+            elif isinstance(item, ToolCallMessage):
+                calls: list[dict[str, object]] = []
+                while index < len(items) and isinstance(items[index], ToolCallMessage):
+                    call = items[index]
+                    calls.append(
+                        {
+                            "id": call.unique_id(),
+                            "type": "function",
+                            "function": {
+                                "name": call.tool_name(),
+                                "arguments": orjson.dumps(call.arguments()).decode("utf-8"),
+                            },
+                        }
+                    )
+                    index += 1
+                assistant = {"role": "assistant", "content": None, "tool_calls": calls}
+                if pending_reasoning:
+                    assistant["reasoning_content"] = pending_reasoning
+                    pending_reasoning = None
+                messages.append(assistant)
+                continue
+            elif isinstance(item, ToolCallResponseMessage):
+                messages.append(
+                    {
+                        "role": "tool",
+                        "tool_call_id": item.unique_id(),
+                        "content": item.text(),
+                    }
+                )
+            else:
+                raise ValueError(f"Unknown type of message {type(item).__name__}")
+            index += 1
+        return messages