149 lines
6.0 KiB
Python
149 lines
6.0 KiB
Python
"""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
|