Add Bedrock InvokeModel and InvokeModelWithResponseStream and support multi-turn tool calling across providers

tools/gen-ai/mock-server/src/genai_mock_server/anthropic.py

@@ -82,16 +82,61 @@
 }
 
 
+def _current_turn(messages):
+    """Messages since the last user message that is not a tool result.
+
+    A tool called in an earlier turn must not stop the model from calling it
+    again when the user asks a new question.
+    """
+    last_user = -1
+    for index, message in enumerate(messages):
+        if message.get("role") != "user":
+            continue
+        content = message.get("content")
+        if isinstance(content, list) and any(
+            isinstance(block, dict) and block.get("type") == "tool_result"
+            for block in content
+        ):
+            continue
+        last_user = index
+    return messages[last_user + 1 :]
+
+
 def _has_tool_result(body):
     for message in body.get("messages", []):
         content = message.get("content")
         if isinstance(content, list):
             for block in content:
-                if isinstance(block, dict) and block.get("type") == "tool_result":
+                if (
+                    isinstance(block, dict)
+                    and block.get("type") == "tool_result"
+                ):
                     return True
     return False
 
 
+def _get_anthropic_tool_info(messages):
+    called_names = set()
+    call_ids = []
+    for message in messages:
+        content = message.get("content")
+        if isinstance(content, list):
+            for block in content:
+                if isinstance(block, dict):
+                    if block.get("type") == "tool_use":
+                        call_id = block.get("id")
+                        if call_id:
+                            call_ids.append(call_id)
+                        name = block.get("name")
+                        if name:
+                            called_names.add(name)
+                    elif block.get("type") == "tool_result":
+                        call_id = block.get("tool_use_id")
+                        if call_id:
+                            call_ids.append(call_id)
+    return called_names, call_ids
+
+
 def _sse_event(event_type, data):
     return f"event: {event_type}\ndata: {json.dumps(data)}\n\n"
 
@@ -170,20 +215,41 @@ def messages():
 
     context_management = body.get("context_management") or {}
     edits = context_management.get("edits") or []
-    if any(edit.get("type") == "compact_20260112" for edit in edits if isinstance(edit, dict)):
+    if any(
+        edit.get("type") == "compact_20260112"
+        for edit in edits
+        if isinstance(edit, dict)
+    ):
         resp = copy.deepcopy(MESSAGE_COMPACTION_RESPONSE)
         resp["model"] = body.get("model", resp["model"])
         return resp
 
-    if body.get("tools") and not _has_tool_result(body):
-        resp = copy.deepcopy(MESSAGE_TOOL_USE_RESPONSE)
-        resp["model"] = body.get("model", resp["model"])
-        tool = body.get("tools", [{}])[0]
-        tool_name = tool.get("name")
-        if tool_name:
-            resp["content"][0]["name"] = tool_name
-        resp["content"][0]["input"] = mock_tool_arguments(tool)
-        return resp
+    if body.get("tools"):
+        tools = body.get("tools")
+        first_tool = tools[0] if tools else {}
+        tool_name = first_tool.get("name")
+        messages = body.get("messages", [])
+        _, call_ids = _get_anthropic_tool_info(messages)
+        turn = _current_turn(messages)
+        called_names, _ = _get_anthropic_tool_info(turn)
+        if not _has_tool_result({"messages": turn}):
+            should_call = True
+        elif called_names:
+            should_call = bool(tool_name) and tool_name not in called_names
+        else:
+            # A result with no call to attribute it to: answer rather than loop.
+            should_call = False
+
+        if should_call:
+            call_idx = len(set(call_ids)) + 1
+            call_id = f"toolu_mock_{call_idx:03d}"
+            resp = copy.deepcopy(MESSAGE_TOOL_USE_RESPONSE)
+            resp["model"] = body.get("model", resp["model"])
+            if tool_name:
+                resp["content"][0]["name"] = tool_name
+            resp["content"][0]["id"] = call_id
+            resp["content"][0]["input"] = mock_tool_arguments(first_tool)
+            return resp
 
     resp = copy.deepcopy(MESSAGE_RESPONSE)
     resp["model"] = body.get("model", resp["model"])

tools/gen-ai/mock-server/src/genai_mock_server/bedrock.py

@@ -1,5 +1,6 @@
 """AWS Bedrock-compatible endpoints."""
 
+import base64
 import copy
 import json
 
@@ -75,7 +76,9 @@ def _tool_to_call(body):
     ]
     if not specifications:
         return None
-    chosen = ((tool_config.get("toolChoice") or {}).get("tool") or {}).get("name")
+    chosen = ((tool_config.get("toolChoice") or {}).get("tool") or {}).get(
+        "name"
+    )
     for specification in specifications:
         if specification["name"] == chosen:
             return specification
@@ -96,14 +99,23 @@ def _stream_converse():
     events = []
     events.append(("messageStart", {"role": "assistant"}))
     for word in ["This ", "is ", "a ", "mock ", "streamed ", "response."]:
-        events.append(("contentBlockDelta", {"delta": {"text": word}, "contentBlockIndex": 0}))
+        events.append(
+            (
+                "contentBlockDelta",
+                {"delta": {"text": word}, "contentBlockIndex": 0},
+            )
+        )
     events.append(("contentBlockStop", {"contentBlockIndex": 0}))
     events.append(("messageStop", {"stopReason": "end_turn"}))
     events.append(
         (
             "metadata",
             {
-                "usage": {"inputTokens": 25, "outputTokens": 6, "totalTokens": 31},
+                "usage": {
+                    "inputTokens": 25,
+                    "outputTokens": 6,
+                    "totalTokens": 31,
+                },
                 "metrics": {"latencyMs": 100},
             },
         )
@@ -125,14 +137,91 @@ def bedrock_converse(model_id):
 
 @bp.route("/model/<path:model_id>/converse-stream", methods=["POST"])
 def bedrock_converse_stream(model_id):
-    return Response(_stream_converse(), mimetype="application/vnd.amazon.eventstream")
+    return Response(
+        _stream_converse(), mimetype="application/vnd.amazon.eventstream"
+    )
 
 
 @bp.route("/model/<path:model_id>/invoke", methods=["POST"])
 def bedrock_invoke(model_id):
-    """Handle Bedrock InvokeModel — used for Titan Embeddings."""
+    """Handle Bedrock InvokeModel."""
+    if "embed" in model_id:
+        resp = {
+            "embedding": [0.001] * 256,
+            "inputTextTokenCount": 8,
+        }
+        headers = {
+            "x-amzn-bedrock-input-token-count": "8",
+            "x-amzn-bedrock-content-type": "application/json",
+        }
+        return Response(
+            json.dumps(resp), mimetype="application/json", headers=headers
+        )
+
     resp = {
-        "embedding": [0.001] * 256,
-        "inputTextTokenCount": 8,
+        "inputTextTokenCount": 5,
+        "results": [
+            {
+                "tokenCount": 10,
+                "outputText": "This is a response from the mock server.",
+                "completionReason": "FINISH",
+            }
+        ],
     }
-    return Response(json.dumps(resp), mimetype="application/json")
+    headers = {
+        "x-amzn-bedrock-input-token-count": "5",
+        "x-amzn-bedrock-output-token-count": "10",
+        "x-amzn-bedrock-content-type": "application/json",
+    }
+    return Response(
+        json.dumps(resp), mimetype="application/json", headers=headers
+    )
+
+
+def _stream_invoke():
+    """Yield Bedrock InvokeModelWithResponseStream event-stream chunks in binary format."""
+    # Titan streams `totalOutputTextTokenCount`, not the `tokenCount` its
+    # non-streaming response carries, and only fills it on the final chunk.
+    chunks = [
+        {
+            "outputText": "This is ",
+            "index": 0,
+            "totalOutputTextTokenCount": None,
+            "completionReason": None,
+            "inputTextTokenCount": 5,
+        },
+        {
+            "outputText": "a test",
+            "index": 0,
+            "totalOutputTextTokenCount": 10,
+            "completionReason": "FINISH",
+            "inputTextTokenCount": 5,
+            "amazon-bedrock-invocationMetrics": {
+                "inputTokenCount": 5,
+                "outputTokenCount": 10,
+                "firstByteLatency": 100,
+                "invocationLatency": 200,
+            },
+        },
+    ]
+    for chunk in chunks:
+        raw_bytes = json.dumps(chunk).encode("utf-8")
+        payload = json.dumps(
+            {"bytes": base64.b64encode(raw_bytes).decode("ascii")}
+        ).encode("utf-8")
+        yield encode_aws_event_stream_message("chunk", payload)
+
+
+@bp.route(
+    "/model/<path:model_id>/invoke-with-response-stream", methods=["POST"]
+)
+def bedrock_invoke_stream(model_id):
+    """Handle Bedrock InvokeModelWithResponseStream."""
+    headers = {
+        "x-amzn-bedrock-content-type": "application/json",
+    }
+    return Response(
+        _stream_invoke(),
+        mimetype="application/vnd.amazon.eventstream",
+        headers=headers,
+    )