364 lines
11 KiB
Python
364 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import base64
|
|
import json
|
|
from dataclasses import dataclass
|
|
from typing import Any, Protocol
|
|
|
|
from runtime.planner_config import PlannerConfig
|
|
from runtime.tool_specs import ToolSpec
|
|
|
|
|
|
class ToolCallUnavailable(Exception):
|
|
"""Internal signal for expected tool-calling transport/response failures."""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ToolCallUsage:
|
|
input_tokens: int | None = None
|
|
output_tokens: int | None = None
|
|
total_tokens: int | None = None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ToolCallDecision:
|
|
tool_name: str
|
|
arguments: dict[str, Any]
|
|
usage: ToolCallUsage | None = None
|
|
|
|
|
|
class ToolCallingClient(Protocol):
|
|
def decide(
|
|
self,
|
|
*,
|
|
system_prompt: str,
|
|
user_prompt: str,
|
|
screenshot: bytes | None,
|
|
tools: list[ToolSpec],
|
|
timeout: float,
|
|
) -> ToolCallDecision: ...
|
|
|
|
|
|
class AnthropicToolCallingClient:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
model: str,
|
|
transport: Any | None = None,
|
|
max_tokens: int = 1024,
|
|
api_key: str | None = None,
|
|
) -> None:
|
|
self.model = model
|
|
self._transport = transport
|
|
self.max_tokens = max_tokens
|
|
self._api_key = api_key
|
|
|
|
def decide(
|
|
self,
|
|
*,
|
|
system_prompt: str,
|
|
user_prompt: str,
|
|
screenshot: bytes | None,
|
|
tools: list[ToolSpec],
|
|
timeout: float,
|
|
) -> ToolCallDecision:
|
|
try:
|
|
response = self._create_message(
|
|
system_prompt,
|
|
user_prompt,
|
|
screenshot,
|
|
tools,
|
|
timeout=timeout,
|
|
)
|
|
return _decision_from_anthropic_response(response)
|
|
except ToolCallUnavailable:
|
|
raise
|
|
except Exception as exc:
|
|
raise ToolCallUnavailable(str(exc)) from exc
|
|
|
|
def _create_message(
|
|
self,
|
|
system_prompt: str,
|
|
user_prompt: str,
|
|
screenshot: bytes | None,
|
|
tools: list[ToolSpec],
|
|
*,
|
|
timeout: float,
|
|
) -> Any:
|
|
client = self._client()
|
|
kwargs = {
|
|
"model": self.model,
|
|
"max_tokens": self.max_tokens,
|
|
"timeout": timeout,
|
|
"system": [
|
|
{
|
|
"type": "text",
|
|
"text": system_prompt,
|
|
"cache_control": {"type": "ephemeral"},
|
|
}
|
|
],
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": _anthropic_content(user_prompt, screenshot),
|
|
}
|
|
],
|
|
"tools": [_anthropic_tool(spec) for spec in tools],
|
|
"tool_choice": {"type": "any", "disable_parallel_tool_use": True},
|
|
}
|
|
messages = getattr(client, "messages", None)
|
|
if messages is not None:
|
|
return messages.create(**kwargs)
|
|
return client.create(**kwargs)
|
|
|
|
def _client(self) -> Any:
|
|
if self._transport is not None:
|
|
return self._transport
|
|
|
|
try:
|
|
import anthropic
|
|
except Exception as exc:
|
|
raise ToolCallUnavailable("anthropic SDK is unavailable") from exc
|
|
|
|
self._transport = (
|
|
anthropic.Anthropic(api_key=self._api_key)
|
|
if self._api_key is not None
|
|
else anthropic.Anthropic()
|
|
)
|
|
return self._transport
|
|
|
|
|
|
class OpenAIToolCallingClient:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
model: str,
|
|
transport: Any | None = None,
|
|
max_tokens: int = 1024,
|
|
api_key: str | None = None,
|
|
base_url: str | None = None,
|
|
) -> None:
|
|
self.model = model
|
|
self._transport = transport
|
|
self.max_tokens = max_tokens
|
|
self._api_key = api_key
|
|
self._base_url = base_url
|
|
|
|
def decide(
|
|
self,
|
|
*,
|
|
system_prompt: str,
|
|
user_prompt: str,
|
|
screenshot: bytes | None,
|
|
tools: list[ToolSpec],
|
|
timeout: float,
|
|
) -> ToolCallDecision:
|
|
try:
|
|
response = self._create_completion(
|
|
system_prompt,
|
|
user_prompt,
|
|
screenshot,
|
|
tools,
|
|
timeout=timeout,
|
|
)
|
|
return _decision_from_openai_response(response)
|
|
except ToolCallUnavailable:
|
|
raise
|
|
except Exception as exc:
|
|
raise ToolCallUnavailable(str(exc)) from exc
|
|
|
|
def _create_completion(
|
|
self,
|
|
system_prompt: str,
|
|
user_prompt: str,
|
|
screenshot: bytes | None,
|
|
tools: list[ToolSpec],
|
|
*,
|
|
timeout: float,
|
|
) -> Any:
|
|
client = self._client()
|
|
kwargs = {
|
|
"model": self.model,
|
|
"max_completion_tokens": self.max_tokens,
|
|
"timeout": timeout,
|
|
"messages": [
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": _openai_content(user_prompt, screenshot)},
|
|
],
|
|
"tools": [_openai_tool(spec) for spec in tools],
|
|
"tool_choice": "required",
|
|
"parallel_tool_calls": False,
|
|
}
|
|
chat = getattr(client, "chat", None)
|
|
if chat is not None:
|
|
return chat.completions.create(**kwargs)
|
|
return client.create(**kwargs)
|
|
|
|
def _client(self) -> Any:
|
|
if self._transport is not None:
|
|
return self._transport
|
|
|
|
try:
|
|
from openai import OpenAI
|
|
except Exception as exc:
|
|
raise ToolCallUnavailable("openai SDK is unavailable") from exc
|
|
|
|
kwargs: dict[str, str] = {}
|
|
if self._api_key is not None:
|
|
kwargs["api_key"] = self._api_key
|
|
if self._base_url is not None:
|
|
kwargs["base_url"] = self._base_url
|
|
self._transport = OpenAI(**kwargs)
|
|
return self._transport
|
|
|
|
|
|
def build_client(config: PlannerConfig) -> ToolCallingClient:
|
|
model = config.resolved_model()
|
|
if config.provider == "openai":
|
|
return OpenAIToolCallingClient(model=model)
|
|
return AnthropicToolCallingClient(model=model)
|
|
|
|
|
|
def _anthropic_content(
|
|
user_prompt: str,
|
|
screenshot: bytes | None,
|
|
) -> list[dict[str, Any]]:
|
|
content: list[dict[str, Any]] = []
|
|
if screenshot is not None:
|
|
content.append(
|
|
{
|
|
"type": "image",
|
|
"source": {
|
|
"type": "base64",
|
|
"media_type": "image/png",
|
|
"data": base64.b64encode(screenshot).decode("ascii"),
|
|
},
|
|
}
|
|
)
|
|
content.append({"type": "text", "text": user_prompt})
|
|
return content
|
|
|
|
|
|
def _anthropic_tool(spec: ToolSpec) -> dict[str, Any]:
|
|
return {
|
|
"name": spec.name,
|
|
"description": spec.description,
|
|
"input_schema": spec.parameters,
|
|
}
|
|
|
|
|
|
def _decision_from_anthropic_response(response: Any) -> ToolCallDecision:
|
|
content = _value(response, "content")
|
|
if not isinstance(content, list):
|
|
raise ValueError("anthropic tool-call response missing content list")
|
|
for block in content:
|
|
if _value(block, "type") != "tool_use":
|
|
continue
|
|
name = _value(block, "name")
|
|
arguments = _value(block, "input")
|
|
if isinstance(name, str) and isinstance(arguments, dict):
|
|
return ToolCallDecision(
|
|
tool_name=name,
|
|
arguments=arguments,
|
|
usage=_anthropic_usage(response),
|
|
)
|
|
raise ValueError("anthropic response did not include a tool_use block")
|
|
|
|
|
|
def _openai_content(
|
|
user_prompt: str,
|
|
screenshot: bytes | None,
|
|
) -> str | list[dict[str, Any]]:
|
|
if screenshot is None:
|
|
return user_prompt
|
|
encoded = base64.b64encode(screenshot).decode("ascii")
|
|
return [
|
|
{"type": "text", "text": user_prompt},
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {"url": f"data:image/png;base64,{encoded}"},
|
|
},
|
|
]
|
|
|
|
|
|
def _openai_tool(spec: ToolSpec) -> dict[str, Any]:
|
|
return {
|
|
"type": "function",
|
|
"function": {
|
|
"name": spec.name,
|
|
"description": spec.description,
|
|
"parameters": spec.parameters,
|
|
},
|
|
}
|
|
|
|
|
|
def _decision_from_openai_response(response: Any) -> ToolCallDecision:
|
|
choices = _value(response, "choices")
|
|
if not isinstance(choices, list) or not choices:
|
|
raise ValueError("openai tool-call response missing choices")
|
|
message = _value(choices[0], "message")
|
|
tool_calls = _value(message, "tool_calls")
|
|
if not isinstance(tool_calls, list) or not tool_calls:
|
|
raise ValueError("openai response did not include a tool call")
|
|
function = _value(tool_calls[0], "function")
|
|
name = _value(function, "name")
|
|
if not isinstance(name, str):
|
|
raise ValueError("openai tool call missing a function name")
|
|
arguments = _decode_openai_arguments(_value(function, "arguments"))
|
|
return ToolCallDecision(
|
|
tool_name=name,
|
|
arguments=arguments,
|
|
usage=_openai_usage(response),
|
|
)
|
|
|
|
|
|
def _anthropic_usage(response: Any) -> ToolCallUsage | None:
|
|
usage = _value(response, "usage")
|
|
input_tokens = _integer_value(usage, "input_tokens")
|
|
output_tokens = _integer_value(usage, "output_tokens")
|
|
if input_tokens is None and output_tokens is None:
|
|
return None
|
|
return ToolCallUsage(
|
|
input_tokens=input_tokens,
|
|
output_tokens=output_tokens,
|
|
total_tokens=(input_tokens or 0) + (output_tokens or 0),
|
|
)
|
|
|
|
|
|
def _openai_usage(response: Any) -> ToolCallUsage | None:
|
|
usage = _value(response, "usage")
|
|
input_tokens = _integer_value(usage, "prompt_tokens")
|
|
output_tokens = _integer_value(usage, "completion_tokens")
|
|
total_tokens = _integer_value(usage, "total_tokens")
|
|
if input_tokens is None and output_tokens is None and total_tokens is None:
|
|
return None
|
|
return ToolCallUsage(
|
|
input_tokens=input_tokens,
|
|
output_tokens=output_tokens,
|
|
total_tokens=total_tokens
|
|
if total_tokens is not None
|
|
else (input_tokens or 0) + (output_tokens or 0),
|
|
)
|
|
|
|
|
|
def _decode_openai_arguments(raw_arguments: Any) -> dict[str, Any]:
|
|
if isinstance(raw_arguments, dict):
|
|
return raw_arguments
|
|
if isinstance(raw_arguments, str):
|
|
decoded = json.loads(raw_arguments)
|
|
if isinstance(decoded, dict):
|
|
return decoded
|
|
raise ValueError("openai tool call arguments must decode to a JSON object")
|
|
|
|
|
|
def _value(source: Any, key: str) -> Any:
|
|
if isinstance(source, dict):
|
|
return source.get(key)
|
|
value = getattr(source, key, None)
|
|
return None if callable(value) else value
|
|
|
|
|
|
def _integer_value(source: Any, key: str) -> int | None:
|
|
value = _value(source, key)
|
|
return value if isinstance(value, int) and value >= 0 else None
|