diff --git a/strands-py/strands/_wasm_host.py b/strands-py/strands/_wasm_host.py index 9185ac1bb..4e3f0e35c 100644 --- a/strands-py/strands/_wasm_host.py +++ b/strands-py/strands/_wasm_host.py @@ -82,6 +82,22 @@ def log(self, level: str, message: str, context: typing.Optional[str]) -> None: raise NotImplementedError +class ModelProviderBase(ABC): + @abstractmethod + def invoke( + self, + messages: str, + system_prompt: typing.Optional[str], + tool_specs: typing.Optional[str], + config: str, + ) -> list[str]: + """Invoke the host-side model provider. + + Returns a list of model event JSON strings. + """ + raise NotImplementedError + + def _run_sync(coro: typing.Coroutine) -> typing.Any: """Run an async coroutine from sync context, even if an event loop is running.""" try: @@ -212,6 +228,9 @@ def _build_model_config_variant(cfg: ModelConfigInput) -> Variant: } ) return Variant("gemini", payload) + if provider == "host-model": + payload = _rec(config=cfg.additional_config or "{}") + return Variant("host-model", payload) raise ValueError(f"unknown model provider: {provider}") @@ -476,6 +495,47 @@ def log_fn(store_ctx: typing.Any, entry: typing.Any) -> None: return log_fn +def _make_model_invoke_fn( + provider: ModelProviderBase | None, +) -> typing.Callable[..., typing.Any]: + """Create the callback for the model-provider.invoke WIT import.""" + + async def invoke_fn(store_ctx: typing.Any, args: typing.Any) -> Variant: + if provider is None: + return Variant("err", "no host model provider configured") + messages = getattr(args, "messages") + system_prompt = getattr(args, "system-prompt") + config = getattr(args, "config") + + # Serialize tool specs if present + raw_specs = getattr(args, "tool-specs") + tool_specs_json: str | None = None + if raw_specs is not None: + import json as _json + + tool_specs_json = _json.dumps( + [ + { + "name": getattr(s, "name"), + "description": getattr(s, "description"), + "inputSchema": getattr(s, "input-schema"), + } + for s in raw_specs + ] + ) + + try: + event_jsons = provider.invoke(messages, system_prompt, tool_specs_json, config) + log.debug("model-provider invoke returned %d events", len(event_jsons)) + import json as _json2 + return Variant("ok", _json2.dumps(event_jsons)) + except Exception as exc: + log.error("model-provider invoke failed: %s", exc) + return Variant("err", str(exc)) + + return invoke_fn + + # --------------------------------------------------------------------------- # WasmAgent — drop-in replacement for the former native Agent class # --------------------------------------------------------------------------- @@ -491,6 +551,7 @@ def __init__( tools: list[ToolSpec] | None, tool_dispatcher: ToolDispatcherBase | None, log_handler: LogHandlerBase | None, + model_provider: ModelProviderBase | None = None, use_callback_relay: bool = False, ): engine, component = _get_engine_and_component() @@ -506,6 +567,8 @@ def __init__( tp.add_func("call-tools", _make_call_tools_fn(tool_dispatcher)) with root.add_instance("strands:agent/host-log") as hl: hl.add_func("log", _make_log_fn(log_handler)) + with root.add_instance("strands:agent/model-provider") as mp: + mp.add_func_async("invoke", _make_model_invoke_fn(model_provider)) # --- store --- store = Store(engine) @@ -606,4 +669,4 @@ def get_messages(self) -> str: def set_messages(self, json: str) -> None: """Sync wrapper — safe from any context (inside or outside event loop).""" - _run_sync(self.set_messages_async(json)) + _run_sync(self.set_messages_async(json)) \ No newline at end of file diff --git a/strands-py/strands/agent/__init__.py b/strands-py/strands/agent/__init__.py index cac7de986..e9a035475 100644 --- a/strands-py/strands/agent/__init__.py +++ b/strands-py/strands/agent/__init__.py @@ -279,7 +279,20 @@ def __init__( sp_blocks = json.dumps(system_prompt) sp_str = None - model_config = self._build_model_config(resolve_model(model)) + # Detect host-side model providers (have stream() method — e.g. SageMaker, Mistral, Writer). + # These run on the Python side and are invoked by the WASM guest via model-provider import. + model_provider_callback = None + if model is not None and hasattr(model, "stream") and callable(getattr(model, "stream", None)): + from strands.models.host_adapter import HostModelAdapter + + model_provider_callback = HostModelAdapter(model) + model_config = _ModelConfigInput( + provider="host-model", + additional_config=json.dumps({"provider_type": type(model).__name__}), + ) + else: + model_config = self._build_model_config(resolve_model(model)) + tool_specs = ( [ _ToolSpec( @@ -300,6 +313,7 @@ def __init__( tools=tool_specs, tool_dispatcher=self._dispatcher, log_handler=_LogHandler(), + model_provider=model_provider_callback, use_callback_relay=False, ) @@ -766,4 +780,4 @@ def cleanup(self) -> None: # Re-export for test compatibility from strands.agent.conversation_manager import NullConversationManager # noqa: E402 -__all__ = ["Agent", "AgentResult", "NullConversationManager"] +__all__ = ["Agent", "AgentResult", "NullConversationManager"] \ No newline at end of file diff --git a/strands-py/strands/models/__init__.py b/strands-py/strands/models/__init__.py index b6b7da999..588a75a11 100644 --- a/strands-py/strands/models/__init__.py +++ b/strands-py/strands/models/__init__.py @@ -3,5 +3,6 @@ from strands.models.gemini import GeminiModel from strands.models.model import Model from strands.models.openai import OpenAIModel +from strands.models.sagemaker import SageMakerAIModel -__all__ = ["AnthropicModel", "BedrockModel", "GeminiModel", "Model", "OpenAIModel"] +__all__ = ["AnthropicModel", "BedrockModel", "GeminiModel", "Model", "OpenAIModel", "SageMakerAIModel"] \ No newline at end of file diff --git a/strands-py/strands/models/_validation.py b/strands-py/strands/models/_validation.py new file mode 100644 index 000000000..5919ccb40 --- /dev/null +++ b/strands-py/strands/models/_validation.py @@ -0,0 +1,64 @@ +"""Configuration validation utilities for model providers.""" + +import warnings +from collections.abc import Mapping +from typing import Any + +from typing_extensions import get_type_hints + +from ..types.content import ContentBlock +from ..types.tools import ToolChoice + + +def validate_config_keys(config_dict: Mapping[str, Any], config_class: type) -> None: + """Validate that config keys match the TypedDict fields. + + Args: + config_dict: Dictionary of configuration parameters + config_class: TypedDict class to validate against + """ + valid_keys = set(get_type_hints(config_class).keys()) + provided_keys = set(config_dict.keys()) + invalid_keys = provided_keys - valid_keys + + if invalid_keys: + warnings.warn( + f"Invalid configuration parameters: {sorted(invalid_keys)}." + f"\nValid parameters are: {sorted(valid_keys)}." + f"\n" + f"\nSee https://github.com/strands-agents/sdk-python/issues/815", + stacklevel=4, + ) + + +def warn_on_tool_choice_not_supported(tool_choice: ToolChoice | None) -> None: + """Emits a warning if a tool choice is provided but not supported by the provider. + + Args: + tool_choice: the tool_choice provided to the provider + """ + if tool_choice: + warnings.warn( + "A ToolChoice was provided to this provider but is not supported and will be ignored", + stacklevel=4, + ) + + +def _has_location_source(content: ContentBlock) -> bool: + """Check if a content block contains a location source. + + Providers need to explicitly define an implementation to support content locations. + + Args: + content: Content block to check. + + Returns: + True if the content block contains an location source, False otherwise. + """ + if "image" in content: + return "location" in content["image"].get("source", {}) + if "document" in content: + return "location" in content["document"].get("source", {}) + if "video" in content: + return "location" in content["video"].get("source", {}) + return False \ No newline at end of file diff --git a/strands-py/strands/models/host_adapter.py b/strands-py/strands/models/host_adapter.py new file mode 100644 index 000000000..368a266ba --- /dev/null +++ b/strands-py/strands/models/host_adapter.py @@ -0,0 +1,186 @@ +"""Generic adapter that bridges a Python model provider to the WIT model-provider callback. + +Wraps any Python model that has a `stream()` generator (sync or async) and +collects its StreamEvent output, serializing it for the WIT boundary. +""" + +import asyncio +import inspect +import json +import logging +import threading +from typing import Any + +from strands._wasm_host import ModelProviderBase + +logger = logging.getLogger(__name__) + + +def _run_in_thread(coro: Any) -> Any: + """Run an async coroutine in a new thread with a fresh event loop. + + This is needed because the adapter's invoke() is called from within + wasmtime-py's async execution context, which already owns the event loop. + We cannot use asyncio.run() (loop already running) or await (we're in a + sync callback). Spawning a dedicated thread with its own loop is the + only safe option. + """ + result: list[Any] = [None] + exc: list[BaseException | None] = [None] + + def _target() -> None: + try: + result[0] = asyncio.run(coro) + except BaseException as e: + exc[0] = e + + t = threading.Thread(target=_target) + t.start() + t.join() + if exc[0] is not None: + raise exc[0] + return result[0] + + +class HostModelAdapter(ModelProviderBase): + """Implements the UniFFI ModelProvider foreign trait. + + Wraps a Python model instance (SageMaker, Mistral, Writer, etc.) and + adapts its async stream() output into the JSON event list the WASM + guest expects. + """ + + def __init__(self, model: Any) -> None: + self._model = model + + def invoke( + self, + messages: str, + system_prompt: str | None, + tool_specs: str | None, + config: str, + ) -> list[str]: + """Called when the WASM guest invokes model-provider.invoke(). + + Deserializes the WIT args, calls the Python model's stream(), + collects all events, and returns them as JSON strings. + """ + logger.debug("host model invoke: messages_len=%d", len(messages)) + parsed_messages = json.loads(messages) + + parsed_tool_specs = None + if tool_specs: + raw_specs = json.loads(tool_specs) + parsed_tool_specs = [] + for spec in raw_specs: + input_schema = spec.get("inputSchema", "{}") + if isinstance(input_schema, str): + try: + input_schema = json.loads(input_schema) + except json.JSONDecodeError: + input_schema = {} + parsed_tool_specs.append({ + "name": spec["name"], + "description": spec["description"], + "inputSchema": {"json": input_schema}, + }) + + # Call model.stream() — may return a sync generator or async generator. + gen = self._model.stream( + parsed_messages, tool_specs=parsed_tool_specs, system_prompt=system_prompt, + ) + + events: list[str] = [] + if inspect.isasyncgen(gen): + logger.debug("host model: async generator, running in thread") + collected = _run_in_thread(self._collect_async(gen)) + events = [json.dumps(self._normalize_event(e)) for e in collected] + else: + logger.debug("host model: sync generator") + for event in gen: + events.append(json.dumps(self._normalize_event(event))) + + logger.debug("host model invoke: collected %d events", len(events)) + return events + + @staticmethod + async def _collect_async(gen: Any) -> list[dict[str, Any]]: + """Drain an async generator into a list.""" + events: list[dict[str, Any]] = [] + async for event in gen: + events.append(event) + return events + + # Mapping from Python SDK event keys to TS ModelStreamEvent converters. + # Add new entries here when new event types are introduced. + _EVENT_MAP: dict[str, Any] = {} + + @classmethod + def _normalize_event(cls, event: dict[str, Any]) -> dict[str, Any]: + """Convert a Python StreamEvent dict into the TS ModelStreamEvent JSON format. + + The Python SDK uses keys like messageStart, contentBlockDelta, messageStop. + The TS SDK uses types like modelMessageStartEvent, modelContentBlockDeltaEvent, etc. + Unrecognized events are passed through as-is with a debug log. + """ + for key, converter in cls._EVENT_MAP.items(): + if key in event: + return converter(event[key]) + + logger.debug("unrecognized event keys: %s — passing through", list(event.keys())) + return event + + @staticmethod + def _convert_message_start(payload: dict[str, Any]) -> dict[str, Any]: + return {"type": "modelMessageStartEvent", "role": payload["role"]} + + @staticmethod + def _convert_content_block_start(payload: dict[str, Any]) -> dict[str, Any]: + start = payload.get("start", {}) + result: dict[str, Any] = {"type": "modelContentBlockStartEvent"} + if "toolUse" in start: + result["start"] = { + "type": "toolUseStart", + "name": start["toolUse"]["name"], + "toolUseId": start["toolUse"]["toolUseId"], + } + return result + + @staticmethod + def _convert_content_block_delta(payload: dict[str, Any]) -> dict[str, Any]: + delta = payload["delta"] + if "text" in delta: + return {"type": "modelContentBlockDeltaEvent", "delta": {"type": "textDelta", "text": delta["text"]}} + if "toolUse" in delta: + return {"type": "modelContentBlockDeltaEvent", "delta": {"type": "toolUseInputDelta", "input": delta["toolUse"]["input"]}} + if "reasoningContent" in delta: + return {"type": "modelContentBlockDeltaEvent", "delta": {"type": "reasoningContentDelta", "text": delta["reasoningContent"].get("text", "")}} + return {"type": "modelContentBlockDeltaEvent", "delta": delta} + + @staticmethod + def _convert_content_block_stop(_payload: dict[str, Any]) -> dict[str, Any]: + return {"type": "modelContentBlockStopEvent"} + + @staticmethod + def _convert_message_stop(payload: dict[str, Any]) -> dict[str, Any]: + return {"type": "modelMessageStopEvent", "stopReason": payload["stopReason"]} + + @staticmethod + def _convert_metadata(payload: dict[str, Any]) -> dict[str, Any]: + result: dict[str, Any] = {"type": "modelMetadataEvent"} + if "usage" in payload: + result["usage"] = payload["usage"] + if "metrics" in payload: + result["metrics"] = payload["metrics"] + return result + + +# Register converters — add new event types here. +HostModelAdapter._EVENT_MAP = { + "messageStart": HostModelAdapter._convert_message_start, + "contentBlockStart": HostModelAdapter._convert_content_block_start, + "contentBlockDelta": HostModelAdapter._convert_content_block_delta, + "contentBlockStop": HostModelAdapter._convert_content_block_stop, + "messageStop": HostModelAdapter._convert_message_stop, + "metadata": HostModelAdapter._convert_metadata, +} \ No newline at end of file diff --git a/strands-py/strands/models/sagemaker.py b/strands-py/strands/models/sagemaker.py new file mode 100644 index 000000000..9a24b8d57 --- /dev/null +++ b/strands-py/strands/models/sagemaker.py @@ -0,0 +1,390 @@ +"""Amazon SageMaker model provider (host-side). + +This model runs on the Python/host side and is invoked by the WASM guest +via the model-provider WIT import. +""" + +import json +import logging +import os +from collections.abc import Generator +from dataclasses import dataclass +from typing import Any, TypedDict + +import boto3 +from botocore.config import Config as BotocoreConfig + +from ..types.content import Messages +from ..types.tools import ToolSpec +from ._validation import validate_config_keys + +logger = logging.getLogger(__name__) + + +@dataclass +class UsageMetadata: + """Usage metadata from the model response.""" + + total_tokens: int = 0 + completion_tokens: int = 0 + prompt_tokens: int = 0 + + _KNOWN_KEYS = {"total_tokens", "completion_tokens", "prompt_tokens"} + + def __init__(self, **kwargs: Any) -> None: + self.total_tokens = kwargs.get("total_tokens", 0) + self.completion_tokens = kwargs.get("completion_tokens", 0) + self.prompt_tokens = kwargs.get("prompt_tokens", 0) + unknown = set(kwargs) - self._KNOWN_KEYS + if unknown: + logger.debug("UsageMetadata: ignoring unknown fields: %s", unknown) + + +@dataclass +class FunctionCall: + """Function call returned by the model.""" + + name: str | dict[Any, Any] + arguments: str | dict[Any, Any] + + def __init__(self, **kwargs: Any) -> None: + self.name = kwargs.get("name", "") + self.arguments = kwargs.get("arguments", "") + + +@dataclass +class ToolCall: + """Tool call returned by the model.""" + + id: str + type: str + function: FunctionCall + + def __init__(self, **kwargs: Any) -> None: + self.id = str(kwargs.get("id", "")) + self.type = "function" + self.function = FunctionCall(**kwargs.get("function", {"name": "", "arguments": ""})) + + +class SageMakerAIModel: + """Amazon SageMaker model provider implementation (host-side). + + This model provider runs entirely in Python. The WASM guest delegates + inference to it via the model-provider WIT import. + """ + + class SageMakerAIPayloadSchema(TypedDict, total=False): + """Payload schema for the Amazon SageMaker AI model.""" + + max_tokens: int + stream: bool + temperature: float | None + top_p: float | None + top_k: int | None + stop: list[str] | None + additional_args: dict[str, Any] | None + + class SageMakerAIEndpointConfig(TypedDict, total=False): + """Configuration options for SageMaker endpoints.""" + + endpoint_name: str + region_name: str + inference_component_name: str | None + target_model: str | None + target_variant: str | None + additional_args: dict[str, Any] | None + + def __init__( + self, + endpoint_config: SageMakerAIEndpointConfig, + payload_config: SageMakerAIPayloadSchema, + boto_session: boto3.Session | None = None, + boto_client_config: BotocoreConfig | None = None, + ) -> None: + validate_config_keys(endpoint_config, self.SageMakerAIEndpointConfig) + validate_config_keys(payload_config, self.SageMakerAIPayloadSchema) + payload_config.setdefault("stream", True) + self.endpoint_config = self.SageMakerAIEndpointConfig(**endpoint_config) + self.payload_config = self.SageMakerAIPayloadSchema(**payload_config) + + region = self.endpoint_config.get("region_name") or os.getenv("AWS_REGION") or "us-west-2" + session = boto_session or boto3.Session(region_name=str(region)) + + if boto_client_config: + existing_ua = getattr(boto_client_config, "user_agent_extra", None) + new_ua = f"{existing_ua} strands-agents" if existing_ua else "strands-agents" + client_config = boto_client_config.merge(BotocoreConfig(user_agent_extra=new_ua)) + else: + client_config = BotocoreConfig(user_agent_extra="strands-agents") + + self.client = session.client(service_name="sagemaker-runtime", config=client_config) + + def get_config(self) -> "SageMakerAIModel.SageMakerAIEndpointConfig": + return self.endpoint_config + + # ── Message formatting (inlined from upstream OpenAI base) ── + + @staticmethod + def _format_content_block(content: dict[str, Any]) -> dict[str, Any]: + """Format a content block for OpenAI-compatible chat format.""" + if "text" in content: + return {"type": "text", "text": content["text"]} + if "image" in content: + import base64 + + img = content["image"] + b64 = base64.b64encode(img["source"]["bytes"]).decode("utf-8") + fmt = img.get("format", "jpeg") + return {"type": "image_url", "image_url": {"detail": "auto", "url": f"data:image/{fmt};base64,{b64}"}} + return {"type": "text", "text": str(content)} + + @staticmethod + def _format_tool_call(tool_use: dict[str, Any]) -> dict[str, Any]: + return { + "function": {"arguments": json.dumps(tool_use["input"]), "name": tool_use["name"]}, + "id": tool_use["toolUseId"], + "type": "function", + } + + @staticmethod + def _format_tool_message(tool_result: dict[str, Any]) -> dict[str, Any]: + parts: list[str] = [] + for c in tool_result.get("content", []): + if "json" in c: + parts.append(json.dumps(c["json"])) + elif "text" in c: + parts.append(c["text"]) + else: + parts.append(str(c)) + return {"role": "tool", "tool_call_id": tool_result["toolUseId"], "content": " ".join(parts)} + + @classmethod + def format_request_messages( + cls, messages: Messages, system_prompt: str | None = None, + ) -> list[dict[str, Any]]: + formatted: list[dict[str, Any]] = [] + if system_prompt: + formatted.append({"role": "system", "content": system_prompt}) + + for message in messages: + contents = message.get("content", []) + text_parts: list[dict[str, Any]] = [] + tool_calls: list[dict[str, Any]] = [] + tool_msgs: list[dict[str, Any]] = [] + + for c in contents: + if "toolUse" in c: + tool_calls.append(cls._format_tool_call(c["toolUse"])) + elif "toolResult" in c: + tool_msgs.append(cls._format_tool_message(c["toolResult"])) + elif "reasoningContent" not in c: + text_parts.append(cls._format_content_block(c)) + + msg: dict[str, Any] = {"role": message["role"]} + if text_parts: + msg["content"] = text_parts + if tool_calls: + msg["tool_calls"] = tool_calls + if text_parts or tool_calls: + formatted.append(msg) + formatted.extend(tool_msgs) + + return formatted + + # ── Chunk formatting (inlined from upstream OpenAI base) ── + + @staticmethod + def format_chunk(event: dict[str, Any]) -> dict[str, Any]: + match event["chunk_type"]: + case "message_start": + return {"messageStart": {"role": "assistant"}} + case "content_start": + if event.get("data_type") == "tool": + tc = event["data"] + return {"contentBlockStart": {"start": {"toolUse": {"name": tc.function.name, "toolUseId": tc.id}}}} + return {"contentBlockStart": {"start": {}}} + case "content_delta": + if event.get("data_type") == "tool": + tc = event["data"] + return {"contentBlockDelta": {"delta": {"toolUse": {"input": tc.function.arguments or ""}}}} + if event.get("data_type") == "reasoning_content": + return {"contentBlockDelta": {"delta": {"reasoningContent": {"text": event["data"]}}}} + return {"contentBlockDelta": {"delta": {"text": event["data"]}}} + case "content_stop": + return {"contentBlockStop": {}} + case "message_stop": + match event["data"]: + case "tool_calls": + return {"messageStop": {"stopReason": "tool_use"}} + case "length": + return {"messageStop": {"stopReason": "max_tokens"}} + case _: + return {"messageStop": {"stopReason": "end_turn"}} + case "metadata": + usage = event["data"] + return { + "metadata": { + "usage": { + "inputTokens": usage.prompt_tokens, + "outputTokens": usage.completion_tokens, + "totalTokens": usage.total_tokens, + }, + "metrics": {"latencyMs": 0}, + }, + } + case _: + raise RuntimeError(f"chunk_type=<{event['chunk_type']}> | unknown type") + + # ── Request formatting ── + + def format_request( + self, + messages: Messages, + tool_specs: list[ToolSpec] | None = None, + system_prompt: str | None = None, + ) -> dict[str, Any]: + formatted_messages = self.format_request_messages(messages, system_prompt) + + payload: dict[str, Any] = { + "messages": formatted_messages, + "tools": [ + { + "type": "function", + "function": { + "name": ts["name"], + "description": ts["description"], + "parameters": ts.get("inputSchema", {}).get("json", ts.get("inputSchema", {})), + }, + } + for ts in (tool_specs or []) + ], + **{k: v for k, v in self.payload_config.items() if k not in ["additional_args"]}, + } + + extra = self.payload_config.get("additional_args") + if extra: + payload.update(extra) + + if not payload["tools"]: + payload.pop("tools") + payload.pop("tool_choice", None) + else: + payload["tool_choice"] = "auto" + + # Assistant messages: if tool_calls present, drop content + for msg in payload["messages"]: + if msg.get("role") == "assistant" and msg.get("tool_calls"): + msg.pop("content", None) + + request: dict[str, Any] = { + "EndpointName": self.endpoint_config["endpoint_name"], + "Body": json.dumps(payload), + "ContentType": "application/json", + "Accept": "application/json", + } + + inf_comp = self.endpoint_config.get("inference_component_name") + if inf_comp: + request["InferenceComponentName"] = inf_comp + target_model = self.endpoint_config.get("target_model") + if target_model: + request["TargetModel"] = target_model + target_variant = self.endpoint_config.get("target_variant") + if target_variant: + request["TargetVariant"] = target_variant + ep_extra = self.endpoint_config.get("additional_args") + if ep_extra: + request.update(ep_extra) + + return request + + # ── Streaming ── + + def stream( + self, + messages: Messages, + tool_specs: list[ToolSpec] | None = None, + system_prompt: str | None = None, + **kwargs: Any, + ) -> Any: + """Stream conversation with the SageMaker model (synchronous generator).""" + request = self.format_request(messages, tool_specs, system_prompt) + + if self.payload_config.get("stream", True): + response = self.client.invoke_endpoint_with_response_stream(**request) + yield self.format_chunk({"chunk_type": "message_start"}) + + finish_reason = "" + partial_content = "" + tool_calls: dict[int, list[Any]] = {} + text_started = False + + for event in response["Body"]: + chunk = event["PayloadPart"]["Bytes"].decode("utf-8") + partial_content += chunk[6:] if chunk.startswith("data: ") else chunk + try: + content = json.loads(partial_content) + partial_content = "" + choice = content["choices"][0] + + if choice["delta"].get("content"): + if not text_started: + yield self.format_chunk({"chunk_type": "content_start", "data_type": "text"}) + text_started = True + yield self.format_chunk({"chunk_type": "content_delta", "data_type": "text", "data": choice["delta"]["content"]}) + + for tc in choice["delta"].get("tool_calls", []): + tool_calls.setdefault(tc["index"], []).append(tc) + + if choice["finish_reason"] is not None: + finish_reason = choice["finish_reason"] + break + except json.JSONDecodeError: + continue + + if text_started: + yield self.format_chunk({"chunk_type": "content_stop", "data_type": "text"}) + + for tool_deltas in tool_calls.values(): + if not tool_deltas[0]["function"].get("name"): + raise Exception("The model did not provide a tool name.") + yield self.format_chunk({"chunk_type": "content_start", "data_type": "tool", "data": ToolCall(**tool_deltas[0])}) + for td in tool_deltas: + yield self.format_chunk({"chunk_type": "content_delta", "data_type": "tool", "data": ToolCall(**td)}) + yield self.format_chunk({"chunk_type": "content_stop", "data_type": "tool"}) + + if not text_started and not tool_calls: + yield self.format_chunk({"chunk_type": "content_start", "data_type": "text"}) + yield self.format_chunk({"chunk_type": "content_stop", "data_type": "text"}) + + yield self.format_chunk({"chunk_type": "message_stop", "data": finish_reason}) + else: + # Non-streaming path + response = self.client.invoke_endpoint(**request) + body = json.loads(response["Body"].read().decode("utf-8")) + message = body["choices"][0]["message"] + stop_reason = body["choices"][0]["finish_reason"] + + yield self.format_chunk({"chunk_type": "message_start"}) + + if message.get("content", ""): + yield self.format_chunk({"chunk_type": "content_start", "data_type": "text"}) + yield self.format_chunk({"chunk_type": "content_delta", "data_type": "text", "data": message["content"]}) + yield self.format_chunk({"chunk_type": "content_stop", "data_type": "text"}) + + if message.get("tool_calls") or stop_reason == "tool_calls": + tcs = message["tool_calls"] + if not isinstance(tcs, list): + tcs = [tcs] + for tc in tcs: + if not isinstance(tc["function"]["arguments"], str): + tc["function"]["arguments"] = json.dumps(tc["function"]["arguments"]) + yield self.format_chunk({"chunk_type": "content_start", "data_type": "tool", "data": ToolCall(**tc)}) + yield self.format_chunk({"chunk_type": "content_delta", "data_type": "tool", "data": ToolCall(**tc)}) + yield self.format_chunk({"chunk_type": "content_stop", "data_type": "tool"}) + stop_reason = "tool_calls" + + yield self.format_chunk({"chunk_type": "message_stop", "data": stop_reason}) + + if body.get("usage"): + yield self.format_chunk({"chunk_type": "metadata", "data": UsageMetadata(**body["usage"])}) \ No newline at end of file diff --git a/strands-py/strands/types/tools.py b/strands-py/strands/types/tools.py index 804f70fc1..47c39b944 100644 --- a/strands-py/strands/types/tools.py +++ b/strands-py/strands/types/tools.py @@ -30,6 +30,7 @@ def interrupt(self, name: str, reason: str = "") -> str: # Type aliases matching the existing SDK. +ToolChoice = dict[str, Any] ToolResult = dict[str, Any] ToolSpec = dict[str, Any] ToolUse = dict[str, Any] diff --git a/strands-ts/package.json b/strands-ts/package.json index 60a2a9ad0..a8d87139c 100644 --- a/strands-ts/package.json +++ b/strands-ts/package.json @@ -34,6 +34,10 @@ "types": "./dist/src/models/vercel.d.ts", "default": "./dist/src/models/vercel.js" }, + "./models/host-model": { + "types": "./dist/src/models/host-model.d.ts", + "default": "./dist/src/models/host-model.js" + }, "./multiagent": { "types": "./dist/src/multiagent/index.d.ts", "default": "./dist/src/multiagent/index.js" diff --git a/strands-ts/src/models/host-model.ts b/strands-ts/src/models/host-model.ts new file mode 100644 index 000000000..b803fc6da --- /dev/null +++ b/strands-ts/src/models/host-model.ts @@ -0,0 +1,108 @@ +/** + * Host-side model provider proxy. + * + * Instead of making HTTP calls directly, this Model delegates inference + * to the WASM host via the `model-provider` WIT import. The host + * (Python) runs the actual model client and returns serialized + * ModelStreamEvent JSON blobs. + */ + +import { type BaseModelConfig, Model, type StreamOptions } from './model.js' +import type { ModelStreamEvent } from './streaming.js' +import type { Message } from '../types/messages.js' +import { logger } from '../logging/logger.js' + +export interface HostModelConfig extends BaseModelConfig { + /** Opaque provider config JSON passed through to the host. */ + hostConfig: string +} + +export class HostModel extends Model { + private _config: HostModelConfig + private _invoke: (args: { + messages: string + systemPrompt?: string + toolSpecs?: Array<{ name: string; description: string; inputSchema: string }> + config: string + }) => Array<{ data: string }> + + constructor( + config: HostModelConfig, + invoke: (args: { + messages: string + systemPrompt?: string + toolSpecs?: Array<{ name: string; description: string; inputSchema: string }> + config: string + }) => Array<{ data: string }>, + ) { + super() + this._config = config + this._invoke = invoke + } + + updateConfig(modelConfig: HostModelConfig): void { + this._config = { ...this._config, ...modelConfig } + } + + getConfig(): HostModelConfig { + return this._config + } + + async *stream(messages: Message[], options?: StreamOptions): AsyncIterable { + // Serialize messages to JSON for the WIT boundary. + const messagesJson = JSON.stringify( + messages.map((m) => ({ role: m.role, content: m.content })), + ) + + // Serialize system prompt. + let systemPrompt: string | undefined + if (options?.systemPrompt !== undefined) { + systemPrompt = + typeof options.systemPrompt === 'string' + ? options.systemPrompt + : JSON.stringify(options.systemPrompt) + } + + // Serialize tool specs. + const toolSpecs = options?.toolSpecs?.map((spec) => ({ + name: spec.name, + description: spec.description, + inputSchema: JSON.stringify(spec.inputSchema), + })) + + logger.debug('HostModel: invoking host model provider') + + let events: Array<{ data: string }> + try { + const args: { + messages: string + systemPrompt?: string + toolSpecs?: Array<{ name: string; description: string; inputSchema: string }> + config: string + } = { + messages: messagesJson, + config: this._config.hostConfig, + } + if (systemPrompt !== undefined) { + args.systemPrompt = systemPrompt + } + if (toolSpecs !== undefined) { + args.toolSpecs = toolSpecs + } + events = this._invoke(args) + } catch (err: any) { + logger.error('HostModel: host invoke failed', err) + throw new Error(`Host model provider error: ${err?.message ?? err}`) + } + + // Deserialize each event and yield as ModelStreamEvent. + for (const event of events) { + try { + const parsed = JSON.parse(event.data) as ModelStreamEvent + yield parsed + } catch (err) { + logger.warn('HostModel: failed to parse event', { data: event.data }) + } + } + } +} \ No newline at end of file diff --git a/strands-wasm/entry.ts b/strands-wasm/entry.ts index 236c6fed0..ac398a7e5 100644 --- a/strands-wasm/entry.ts +++ b/strands-wasm/entry.ts @@ -26,11 +26,14 @@ import type { import { callTool } from 'strands:agent/tool-provider'; import { log as hostLog } from 'strands:agent/host-log'; -import { Agent, FunctionTool, SessionManager, FileStorage, S3Storage } from '@strands-agents/sdk'; -import { AnthropicModel } from '@strands-agents/sdk/anthropic'; -import { BedrockModel } from '@strands-agents/sdk/bedrock'; -import { OpenAIModel } from '@strands-agents/sdk/openai'; -import { GeminiModel } from '@strands-agents/sdk/gemini'; +import { invoke as hostModelInvoke } from 'strands:agent/model-provider'; +import { Agent, FunctionTool, SessionManager, FileStorage } from '@strands-agents/sdk'; +import { S3Storage } from '@strands-agents/sdk/session/s3-storage'; +import { AnthropicModel } from '@strands-agents/sdk/models/anthropic'; +import { BedrockModel } from '@strands-agents/sdk/models/bedrock'; +import { OpenAIModel } from '@strands-agents/sdk/models/openai'; +import { GoogleModel } from '@strands-agents/sdk/models/google'; +import { HostModel } from '@strands-agents/sdk/models/host-model'; import type { StopReason, AgentStreamEvent, Model, BaseModelConfig } from '@strands-agents/sdk'; // All log calls go through `hostLog` (the WIT import). The host can @@ -164,6 +167,27 @@ function modelParamsConfig(params?: ModelParams): Record { }; } +/** Unwrap the result from hostModelInvoke into {data}[] for HostModel. */ +function unwrapHostModelResult(args: { + messages: string; + systemPrompt?: string; + toolSpecs?: Array<{ name: string; description: string; inputSchema: string }>; + config: string; +}): Array<{ data: string }> { + const result = hostModelInvoke(args); + let jsonStr: string; + if (typeof result === 'object' && result !== null && 'tag' in result) { + if ((result as any).tag === 'err') { + throw new Error((result as any).val); + } + jsonStr = (result as any).val; + } else { + jsonStr = result as any; + } + const eventStrings: string[] = JSON.parse(jsonStr); + return eventStrings.map((s: string) => ({ data: s })); +} + function createModel(config?: ModelConfig, params?: ModelParams): Model { const base = modelParamsConfig(params); @@ -215,13 +239,20 @@ function createModel(config?: ModelConfig, params?: ModelParams): Model, } + /// Opaque config for a host-side model provider. + /// The guest does not interpret this — it passes it back to the host + /// via the model-provider import. + record host-model-config { + /// Provider-specific configuration as a JSON blob. + config: string, + } + variant model-config { anthropic(anthropic-config), bedrock(bedrock-config), openai(openai-config), gemini(gemini-config), + host-model(host-model-config), + } + + /// Arguments for a host-side model invocation. + record model-invoke-args { + /// Conversation messages as a JSON array. + messages: string, + /// System prompt (optional). + system-prompt: option, + /// Tool specifications available to the model. + tool-specs: option>, + /// Provider-specific config JSON (passed through from host-model-config). + config: string, + } + + /// A single model stream event serialized as JSON. + /// Currently unused in the invoke return type due to a wasmtime-py + /// limitation — see the comment on model-provider.invoke below. + /// Retained for future use when the return type can be upgraded to + /// result, string>. + record model-event { + data: string, } record model-params { @@ -212,6 +242,25 @@ interface host-log { log: func(entry: log-entry); } +/// Host-side model provider. +/// +/// When the agent is configured with a `host-model` model-config, the guest +/// delegates model inference to the host via this import instead of +/// making HTTP calls directly. +interface model-provider { + use types.{model-invoke-args, model-event}; + + /// Invoke the host-side model and return the full sequence of events as a JSON array. + /// + /// Stopgap: the ideal signature is result, string>, + /// but wasmtime-py's async component model cannot correctly return + /// list from import callbacks during async execution (triggers + /// "cannot enter component instance" trap). Using result + /// with a JSON-encoded event array as a workaround until the wasmtime-py + /// fork merges async component model fixes upstream. + invoke: func(args: model-invoke-args) -> result; +} + interface api { use types.{agent-config, stream-event, stream-args, respond-args, set-messages-args}; @@ -235,9 +284,11 @@ interface api { world agent { import tool-provider; import host-log; + import model-provider; export api; } world agent-types { export types; } +