pyfect-dev 0.1.0__py3-none-any.whl

pyfect/ai/anthropic/_stream.py

@@ -0,0 +1,377 @@
+"""Streamed message events as encoded stream parts (Effect's
+``makeStreamResponse``)."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, cast
+
+import pyfect as fx
+from pyfect._core import Effect
+from pyfect.ai import _provider as P
+from pyfect.ai import language_model as lm
+from pyfect.ai.ai_error import AiError, ToolParameterValidationError
+from pyfect.ai.anthropic import _schema as W
+from pyfect.ai.anthropic._errors import SECRET_HEADERS
+from pyfect.ai.anthropic._response import (
+    Gen,
+    Part,
+    ServerResults,
+    caller_metadata,
+    citable_documents,
+    citation_source,
+    now_iso,
+    resolve_finish_reason,
+    transform_params,
+    usage_part,
+)
+from pyfect.ai.response import FinishReason
+from pyfect.ai.tool import NameMapper, unsafe_secure_json_parse
+from pyfect.http.response import HttpClientResponse
+from pyfect.stream import Stream
+
+_SERVER_CALLS = frozenset(
+    {
+        "code_execution",
+        "bash_code_execution",
+        "text_editor_code_execution",
+        "web_fetch",
+        "web_search",
+    }
+)
+_SEARCH_CALLS = frozenset({"tool_search_tool_bm25", "tool_search_tool_regex"})
+
+
+@dataclass
+class _ToolBlock:
+    id: str
+    name: str
+    params: P.TextBuffer
+    first_delta: bool
+    provider_name: str | None = None
+    provider_executed: bool | None = None
+    caller: dict[str, Any] | None = None
+
+
+@dataclass
+class _Usage:
+    input_tokens: int = 0
+    output_tokens: int = 0
+    cache_read: int = 0
+    cache_write: int = 0
+
+
+@dataclass
+class _State:
+    options: lm.ProviderOptions
+    response: HttpClientResponse
+    results: ServerResults
+    max_chars: int
+    container: Any = None
+    context_management: Any = None
+    finish_reason: FinishReason = "unknown"
+    stop_sequence: str | None = None
+    raw_usage: dict[str, Any] | None = None
+    usage: _Usage = field(default_factory=_Usage)
+    blocks: dict[int, str | _ToolBlock] = field(default_factory=dict[int, "str | _ToolBlock"])
+    block_type: str | None = None
+
+    @property
+    def names(self) -> NameMapper:
+        return self.results.names
+
+    # message events ---------------------------------------------------------------------
+
+    def message_start(self, event: W.MessageStartEvent) -> Gen[list[Part]]:
+        message = event["message"]
+        usage = message["usage"]
+        self.raw_usage = dict(usage)
+        self.usage.input_tokens = usage["input_tokens"]
+        self.usage.cache_read = usage.get("cache_read_input_tokens") or 0
+        self.usage.cache_write = usage.get("cache_creation_input_tokens") or 0
+        if message.get("container") is not None:
+            self.container = message.get("container")
+        if message["stop_sequence"] is not None:
+            self.stop_sequence = message["stop_sequence"]
+        if message["stop_reason"] is not None:
+            self.finish_reason = resolve_finish_reason(message["stop_reason"])
+        timestamp = yield from now_iso()
+        parts: list[Part] = [
+            {
+                "type": "response-metadata",
+                "id": message["id"],
+                "modelId": message["model"],
+                "timestamp": timestamp,
+                "request": P.request_details(self.response.request, secret_headers=SECRET_HEADERS),
+            }
+        ]
+        for block in message["content"]:
+            if block["type"] == "tool_use":
+                parts.extend((yield from self._prepopulated_tool_use(block)))
+        return parts
+
+    def _prepopulated_tool_use(self, block: W.ContentBlock) -> Gen[list[Part]]:
+        name = self.names.get_custom_name(block["name"])
+        params = yield from transform_params(self.options, name, block.get("input"))
+        call: Part = {"type": "tool-call", "id": block["id"], "name": name, "params": params}
+        caller = caller_metadata(block)
+        if caller is not None:
+            call["metadata"] = {"anthropic": {"caller": caller}}
+        return [
+            {"type": "tool-params-start", "id": block["id"], "name": name},
+            {
+                "type": "tool-params-delta",
+                "id": block["id"],
+                "delta": P.dumps(P.obj(block.get("input"))),
+            },
+            {"type": "tool-params-end", "id": block["id"]},
+            call,
+        ]
+
+    def message_delta(self, event: W.MessageDeltaEvent) -> list[Part]:
+        usage = event["usage"]
+        self.raw_usage = {**(self.raw_usage or {}), **usage}
+        if usage.get("input_tokens") is not None:
+            self.usage.input_tokens = cast("int", usage.get("input_tokens"))
+        self.usage.output_tokens = usage["output_tokens"]
+        if usage.get("cache_read_input_tokens") is not None:
+            self.usage.cache_read = cast("int", usage.get("cache_read_input_tokens"))
+        if usage.get("cache_creation_input_tokens") is not None:
+            self.usage.cache_write = cast("int", usage.get("cache_creation_input_tokens"))
+        delta = event["delta"]
+        if delta.get("container") is not None:
+            self.container = delta.get("container")
+        if event.get("context_management") is not None:
+            self.context_management = event.get("context_management")
+        if delta["stop_reason"] is not None:
+            self.finish_reason = resolve_finish_reason(delta["stop_reason"])
+        if delta["stop_sequence"] is not None:
+            self.stop_sequence = delta["stop_sequence"]
+        return []
+
+    def message_stop(self) -> list[Part]:
+        usage = self.usage
+        return [
+            {
+                "type": "finish",
+                "reason": self.finish_reason,
+                "usage": usage_part(
+                    usage.input_tokens, usage.output_tokens, usage.cache_read, usage.cache_write
+                ),
+                "response": P.response_details(self.response),
+                "metadata": {
+                    "anthropic": {
+                        "container": self.container,
+                        "contextManagement": self.context_management,
+                        "stopSequence": self.stop_sequence,
+                        "usage": self.raw_usage,
+                    }
+                },
+            }
+        ]
+
+    # content blocks ---------------------------------------------------------------------
+
+    def _tool_use_start(self, index: int, block: W.ContentBlock) -> list[Part]:
+        input = P.obj(block.get("input"))
+        initial = P.dumps(input) if input else ""
+        name = self.names.get_custom_name(block["name"])
+        self.blocks[index] = _ToolBlock(
+            block["id"],
+            name,
+            P.TextBuffer(initial, self.max_chars),
+            len(initial) > 0,
+            caller=caller_metadata(block),
+        )
+        return [{"type": "tool-params-start", "id": block["id"], "name": name}]
+
+    def _server_tool_use_start(self, index: int, block: W.ContentBlock) -> list[Part]:
+        provider = block["name"]
+        if provider in _SERVER_CALLS:
+            custom = "code_execution" if provider.endswith("_code_execution") else provider
+        elif provider in _SEARCH_CALLS:
+            self.results.search_calls[block["id"]] = provider
+            custom = provider
+        else:
+            return []
+        name = self.names.get_custom_name(custom)
+        self.blocks[index] = _ToolBlock(
+            block["id"], name, P.TextBuffer("", self.max_chars), True, provider, True
+        )
+        return [
+            {"type": "tool-params-start", "id": block["id"], "name": name, "providerExecuted": True}
+        ]
+
+    def content_block_start(  # noqa: PLR0911
+        self, event: W.ContentBlockStartEvent
+    ) -> Gen[list[Part]]:
+        block = event["content_block"]
+        index = event["index"]
+        kind = block["type"]
+        self.block_type = kind
+        match kind:
+            case "text":
+                self.blocks[index] = "text"
+                return [{"type": "text-start", "id": str(index)}]
+            case "thinking":
+                self.blocks[index] = "reasoning"
+                return [{"type": "reasoning-start", "id": str(index)}]
+            case "redacted_thinking":
+                self.blocks[index] = "reasoning"
+                info = {"type": "redacted_thinking", "redactedData": block.get("data")}
+                return [
+                    {
+                        "type": "reasoning-start",
+                        "id": str(index),
+                        "metadata": {"anthropic": {"info": info}},
+                    }
+                ]
+            case "tool_use":
+                return self._tool_use_start(index, block)
+            case "server_tool_use":
+                return self._server_tool_use_start(index, block)
+            case "mcp_tool_use":
+                return [self.results.mcp_call(block)]
+            case "fallback":
+                timestamp = yield from now_iso()
+                return [
+                    {
+                        "type": "response-metadata",
+                        "modelId": block["to"]["model"],
+                        "timestamp": timestamp,
+                        "metadata": {"anthropic": {"fallback": block}},
+                    }
+                ]
+            case _:
+                parts = yield from self.results.result_parts(block)
+                return parts or []
+
+    def _input_json_delta(self, index: int, delta: str) -> Gen[list[Part]]:
+        block = self.blocks.get(index)
+        if not delta or not isinstance(block, _ToolBlock):
+            return []
+        if block.first_delta and block.provider_name in (
+            "bash_code_execution",
+            "text_editor_code_execution",
+        ):
+            delta = f'{{"type":{P.dumps(block.provider_name)},{delta[1:]}'
+        if not block.params.add(delta):
+            error = P.too_large("AnthropicLanguageModel", "makeStreamResponse", self.max_chars)
+            return (yield from fx.fail(error))
+        block.first_delta = False
+        return [{"type": "tool-params-delta", "id": block.id, "delta": delta}]
+
+    def content_block_delta(self, event: W.ContentBlockDeltaEvent) -> Gen[list[Part]]:
+        delta = event["delta"]
+        index = event["index"]
+        match delta.get("type"):
+            case "text_delta":
+                return [{"type": "text-delta", "id": str(index), "delta": delta.get("text", "")}]
+            case "thinking_delta":
+                return [
+                    {
+                        "type": "reasoning-delta",
+                        "id": str(index),
+                        "delta": delta.get("thinking", ""),
+                    }
+                ]
+            case "signature_delta" if self.block_type == "thinking":
+                info = {"type": "thinking", "signature": delta.get("signature")}
+                return [
+                    {
+                        "type": "reasoning-delta",
+                        "id": str(index),
+                        "delta": "",
+                        "metadata": {"anthropic": {"info": info}},
+                    }
+                ]
+            case "input_json_delta":
+                return (yield from self._input_json_delta(index, delta.get("partial_json", "")))
+            case "citations_delta":
+                source = yield from citation_source(delta["citation"], self.results.documents)
+                return [] if source is None else [source]
+            case _:
+                return []
+
+    def _tool_call(self, block: _ToolBlock) -> Gen[Part]:
+        text = block.params.text() or "{}"
+        if block.provider_name == "code_execution":
+            try:
+                parsed = P.obj(unsafe_secure_json_parse(text))
+                if "code" in parsed and "type" not in parsed:
+                    text = P.dumps({"type": "programmatic-tool-call", **parsed})
+            except ValueError:
+                pass
+        try:
+            params = unsafe_secure_json_parse(text)
+        except ValueError as cause:
+            error = AiError(
+                module="AnthropicLanguageModel",
+                method="makeStreamResponse",
+                reason=ToolParameterValidationError(
+                    tool_name=block.name,
+                    description=f"Failed to securely JSON parse tool parameters: {cause}",
+                ),
+            )
+            return (yield from fx.fail(error))
+        if block.provider_executed is not True:
+            params = yield from transform_params(self.options, block.name, params)
+        part: Part = {"type": "tool-call", "id": block.id, "name": block.name, "params": params}
+        if block.provider_executed is not None:
+            part["providerExecuted"] = block.provider_executed
+        if block.caller is not None:
+            part["metadata"] = {"anthropic": {"caller": block.caller}}
+        return part
+
+    def content_block_stop(self, event: W.ContentBlockStopEvent) -> Gen[list[Part]]:
+        index = event["index"]
+        block = self.blocks.pop(index, None)
+        self.block_type = None
+        if block == "text":
+            return [{"type": "text-end", "id": str(index)}]
+        if block == "reasoning":
+            return [{"type": "reasoning-end", "id": str(index)}]
+        if isinstance(block, _ToolBlock):
+            end: Part = {"type": "tool-params-end", "id": block.id}
+            return [end, (yield from self._tool_call(block))]
+        return []
+
+    def event(self, event: W.MessageStreamEvent) -> Gen[list[Part]]:  # noqa: PLR0911
+        match event["type"]:
+            case "message_start":
+                return (yield from self.message_start(event))
+            case "message_delta":
+                return self.message_delta(event)
+            case "message_stop":
+                return self.message_stop()
+            case "content_block_start":
+                return (yield from self.content_block_start(event))
+            case "content_block_delta":
+                return (yield from self.content_block_delta(event))
+            case "content_block_stop":
+                return (yield from self.content_block_stop(event))
+            case "error":
+                return [
+                    {
+                        "type": "error",
+                        "error": event["error"],
+                        "metadata": {"anthropic": {"requestId": event["request_id"]}},
+                    }
+                ]
+
+
+def make_stream_response(
+    options: lm.ProviderOptions,
+    events: Stream[W.MessageStreamEvent, AiError],
+    response: HttpClientResponse,
+    names: NameMapper,
+    max_chars: int,
+) -> Stream[Part, AiError, Any]:
+    """``max_chars`` bounds each tool call's streamed parameters."""
+    results = ServerResults(names, citable_documents(options.prompt))
+    state = _State(options, response, results, max_chars)
+
+    def convert(event: W.MessageStreamEvent) -> Effect[list[Part], AiError, Any]:
+        return fx.gen(lambda: state.event(event))
+
+    return events.map_effect(convert).flatten_iterable()

pyfect/ai/anthropic/_tools.py

@@ -0,0 +1,248 @@
+"""The tools and tool choice of an Anthropic request (Effect's
+``prepareTools``)."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from typing import Any
+
+from pyfect.ai import _provider as P
+from pyfect.ai import _provider_tools as PT
+from pyfect.ai import language_model as lm
+from pyfect.ai import tool as T
+from pyfect.ai._provider_tools import UnsupportedInput
+from pyfect.ai.ai_error import InvalidUserInputError
+from pyfect.ai.anthropic_structured_output import to_codec_anthropic
+
+type Block = dict[str, Any]
+
+
+@dataclass(frozen=True)
+class Capabilities:
+    max_output_tokens: int
+    supports_structured_output: bool
+    supports_mid_conversation_system_messages: bool
+
+
+@dataclass(frozen=True)
+class Prepared:
+    tools: list[Block] | None
+    tool_choice: Block | None
+
+
+MODULE = "AnthropicLanguageModel"
+
+
+def _error(method: str, reason: Any) -> UnsupportedInput:
+    return PT.unsupported(MODULE, method, reason)
+
+
+def json_schema(schema: Any, method: str) -> Block:
+    return PT.json_schema(MODULE, schema, method, transformer=to_codec_anthropic)
+
+
+def tool_json_schema(tool: T.AnyTool) -> Block:
+    return PT.tool_json_schema(MODULE, tool, require_object=True, transformer=to_codec_anthropic)
+
+
+def _computer(args: Mapping[str, Any], name: str, type_: str) -> Block:
+    return {
+        "name": name,
+        "type": type_,
+        "display_height_px": args.get("displayHeightPx"),
+        "display_width_px": args.get("displayWidthPx"),
+        "display_number": args.get("displayNumber"),
+    }
+
+
+def _user_location(args: Mapping[str, Any]) -> Block | None:
+    location = args.get("userLocation")
+    if location is None:
+        return None
+    return {
+        "type": location.get("type"),
+        "region": location.get("region"),
+        "city": location.get("city"),
+        "country": location.get("country"),
+        "timezone": location.get("timezone"),
+    }
+
+
+def _plain(name: str, type_: str) -> Callable[[Mapping[str, Any]], Block]:
+    return lambda _args: {"name": name, "type": type_}
+
+
+_PROVIDER_TOOLS: dict[str, tuple[str | None, Callable[[Mapping[str, Any]], Block]]] = {
+    "anthropic.bash_20241022": ("computer-use-2024-10-22", _plain("bash", "bash_20241022")),
+    "anthropic.bash_20250124": ("computer-use-2025-01-24", _plain("bash", "bash_20250124")),
+    "anthropic.code_execution_20250522": (
+        "code-execution-2025-05-22",
+        _plain("code_execution", "code_execution_20250522"),
+    ),
+    "anthropic.code_execution_20250825": (
+        "code-execution-2025-08-25",
+        _plain("code_execution", "code_execution_20250825"),
+    ),
+    "anthropic.computer_use_20241022": (
+        "computer-use-2024-10-22",
+        lambda args: _computer(args, "computer", "computer_20241022"),
+    ),
+    "anthropic.computer_20250124": (
+        "computer-use-2025-01-24",
+        lambda args: _computer(args, "computer", "computer_20250124"),
+    ),
+    "anthropic.computer_20251124": (
+        "computer-use-2025-11-24",
+        lambda args: {
+            **_computer(args, "computer", "computer_20251124"),
+            "enable_zoom": args.get("enableZoom", False),
+        },
+    ),
+    "anthropic.memory_20250818": (
+        "context-management-2025-06-27",
+        _plain("memory", "memory_20250818"),
+    ),
+    "anthropic.text_editor_20241022": (
+        "computer-use-2024-10-22",
+        _plain("str_replace_editor", "text_editor_20241022"),
+    ),
+    "anthropic.text_editor_20250124": (
+        "computer-use-2025-01-24",
+        _plain("str_replace_editor", "text_editor_20250124"),
+    ),
+    "anthropic.text_editor_20250429": (
+        "computer-use-2025-01-24",
+        _plain("str_replace_based_edit_tool", "text_editor_20250429"),
+    ),
+    "anthropic.text_editor_20250728": (
+        None,
+        lambda args: {
+            "name": "str_replace_based_edit_tool",
+            "type": "text_editor_20250728",
+            "max_characters": args.get("max_characters"),
+        },
+    ),
+    "anthropic.tool_search_tool_bm25_20251119": (
+        "advanced-tool-use-2025-11-20",
+        _plain("tool_search_tool_bm25", "tool_search_tool_bm25_20251119"),
+    ),
+    "anthropic.tool_search_tool_regex_20251119": (
+        None,
+        _plain("tool_search_tool_regex", "tool_search_tool_regex_20251119"),
+    ),
+    "anthropic.web_search_20250305": (
+        None,
+        lambda args: {
+            "name": "web_search",
+            "type": "web_search_20250305",
+            "max_uses": args.get("maxUses"),
+            "allowed_domains": args.get("allowedDomains"),
+            "blocked_domains": args.get("blockedDomains"),
+            "user_location": _user_location(args),
+        },
+    ),
+    "anthropic.web_fetch_20250910": (
+        "web-fetch-2025-09-10",
+        lambda args: {
+            "name": "web_fetch",
+            "type": "web_fetch_20250910",
+            "max_uses": args.get("maxUses"),
+            "allowed_domains": args.get("allowedDomains"),
+            "blocked_domains": args.get("blockedDomains"),
+            "citations": args.get("citations"),
+            "max_content_tokens": args.get("maxContentTokens"),
+        },
+    ),
+}
+
+
+def _provider_tool(
+    tool: T.ProviderDefined[Any, Any, Any, Any, Any, Any], betas: dict[str, None]
+) -> Block:
+    entry = _PROVIDER_TOOLS.get(tool.provider_id)
+    if entry is None:
+        raise _error(
+            "prepareTools",
+            InvalidUserInputError(
+                description=f"Received request to call unknown provider-defined tool '{tool.name}'"
+            ),
+        )
+    beta, build = entry
+    if beta is not None:
+        betas[beta] = None
+    return build(P.obj(tool.args))
+
+
+def _user_tool(
+    tool: T.AnyTool, capabilities: Capabilities, strict_default: bool | None, betas: dict[str, None]
+) -> Block:
+    block: Block = {"name": tool.name, "input_schema": tool_json_schema(tool)}
+    description = T.get_description(tool)
+    if description is not None:
+        block["description"] = description
+    if capabilities.supports_structured_output:
+        strict = T.get_strict_mode(tool)
+        block["strict"] = strict if strict is not None else (strict_default is not False)
+        betas["structured-outputs-2025-11-13"] = None
+    return block
+
+
+def _tool_choice(choice: lm.ToolChoice, tools: list[Block]) -> tuple[list[Block], Block]:
+    if choice == "auto":
+        return tools, {"type": "auto"}
+    if choice == "required":
+        return tools, {"type": "any"}
+    if isinstance(choice, dict) and "tool" in choice:
+        return tools, {"type": "tool", "name": choice["tool"]}
+    assert isinstance(choice, dict)
+    allowed = set(choice.get("one_of", ()))
+    mode = choice.get("mode")
+    return [t for t in tools if t["name"] in allowed], {
+        "type": "any" if mode == "required" else "auto"
+    }
+
+
+def prepare_tools(
+    options: lm.ProviderOptions,
+    capabilities: Capabilities,
+    config: Mapping[str, Any],
+    betas: dict[str, None],
+) -> Prepared:
+    """Raises ``UnsupportedInput`` for a schema or tool Anthropic cannot
+    take."""
+    response_format = options.response_format
+    if isinstance(response_format, lm.JsonFormat) and not capabilities.supports_structured_output:
+        schema = json_schema(response_format.schema, "prepareTools")
+        user_description = PT.schema_description(response_format.schema)
+        prefix = f"{user_description} - " if user_description is not None else ""
+        return Prepared(
+            [
+                {
+                    "name": response_format.object_name,
+                    "description": f"{prefix}You MUST respond with a JSON object.",
+                    "input_schema": schema,
+                }
+            ],
+            {
+                "type": "tool",
+                "name": response_format.object_name,
+                "disable_parallel_tool_use": True,
+            },
+        )
+    if not options.tools or options.tool_choice == "none":
+        return Prepared(None, None)
+    user_tools: list[Block] = []
+    provider_tools: list[Block] = []
+    for tool in options.tools:
+        if isinstance(tool, T.ProviderDefined):
+            provider_tools.append(_provider_tool(tool, betas))
+        else:
+            user_tools.append(
+                _user_tool(tool, capabilities, config.get("strict_json_schema"), betas)
+            )
+    tools, choice = _tool_choice(options.tool_choice, [*user_tools, *provider_tools])
+    disable = config.get("disable_parallel_tool_calls")
+    if disable is not None and choice["type"] != "none":
+        choice["disable_parallel_tool_use"] = disable
+    return Prepared(tools, choice)