From f4bf7ed7afa602cf649fb7454fccbd6d5b61f1b2 Mon Sep 17 00:00:00 2001 From: Nightingalelyy Date: Mon, 17 Aug 2026 00:32:39 +0800 Subject: [PATCH] fix(examples): repair Arize, AutoGen, and AWS Bedrock tracing --- .../arize/02_datasets_projects_spaces.py | 8 ++ .../03_experiments_prompts_evaluators.py | 2 +- python/tracing/arize/04_admin_operations.py | 21 ++--- python/tracing/arize/_shared.py | 2 + python/tracing/arize/run_all.py | 7 +- python/tracing/autogen/01_assistant_run.py | 37 +++++--- python/tracing/autogen/02_tool_use.py | 41 ++++++--- python/tracing/autogen/03_round_robin_team.py | 39 +++++--- python/tracing/autogen/run_all.py | 28 ++++++ python/tracing/aws-bedrock/01_invoke_model.py | 37 +++++--- python/tracing/aws-bedrock/02_converse.py | 35 +++++--- .../tracing/aws-bedrock/03_converse_stream.py | 49 +++++++---- .../tracing/aws-bedrock/04_converse_tool.py | 82 +++++++++++++++++ .../tracing/aws-bedrock/05_converse_error.py | 62 +++++++++++++ python/tracing/aws-bedrock/_shared.py | 88 ++++++++++++++++++- python/tracing/aws-bedrock/run_all.py | 30 +++++++ 16 files changed, 472 insertions(+), 96 deletions(-) create mode 100644 python/tracing/autogen/run_all.py create mode 100644 python/tracing/aws-bedrock/04_converse_tool.py create mode 100644 python/tracing/aws-bedrock/05_converse_error.py create mode 100644 python/tracing/aws-bedrock/run_all.py diff --git a/python/tracing/arize/02_datasets_projects_spaces.py b/python/tracing/arize/02_datasets_projects_spaces.py index f5e2ea9..25c64cb 100644 --- a/python/tracing/arize/02_datasets_projects_spaces.py +++ b/python/tracing/arize/02_datasets_projects_spaces.py @@ -43,6 +43,14 @@ def run_dataset_project_space_operations() -> str: print_result("datasets.list", client.datasets.list(space="space-offline")) print_result("datasets.create", client.datasets.create(name="offline-dataset", space="space-offline", examples=examples)) print_result("datasets.get", client.datasets.get(dataset="offline-dataset", space="space-offline")) + try: + client.datasets.get( + dataset="missing-dataset", + space="space-offline", + _respan_force_error=True, + ) + except RuntimeError as error: + print(f"datasets.get expected error: {error}") print_result("datasets.update", client.datasets.update(dataset="offline-dataset", name="renamed-dataset", space="space-offline")) print_result("datasets.list_examples", client.datasets.list_examples(dataset="offline-dataset", space="space-offline")) print_result("datasets.append_examples", client.datasets.append_examples(dataset="offline-dataset", space="space-offline", examples=examples)) diff --git a/python/tracing/arize/03_experiments_prompts_evaluators.py b/python/tracing/arize/03_experiments_prompts_evaluators.py index b647de7..6fdc7dc 100644 --- a/python/tracing/arize/03_experiments_prompts_evaluators.py +++ b/python/tracing/arize/03_experiments_prompts_evaluators.py @@ -48,7 +48,7 @@ def _template_config() -> models.TemplateConfig: def _code_config() -> models.CustomCodeConfig: return models.CustomCodeConfig( - type="custom", + type="CUSTOM", name="offline-code", code="def evaluate(row): return 1", variables=[], diff --git a/python/tracing/arize/04_admin_operations.py b/python/tracing/arize/04_admin_operations.py index 14a8be2..5d59ad5 100644 --- a/python/tracing/arize/04_admin_operations.py +++ b/python/tracing/arize/04_admin_operations.py @@ -7,6 +7,7 @@ from __future__ import annotations from arize._generated.api_client import models +from arize.api_keys.types import OrgBinding, SpaceBinding from _shared import ( create_arize_client, @@ -25,7 +26,7 @@ def _run_configuration() -> models.TemplateEvaluationRunConfig: return models.TemplateEvaluationRunConfig( - experiment_type="template_evaluation", + experiment_type="TEMPLATE_EVALUATION", ai_integration_id="ai-integration-offline", model_name="gpt-4o-mini", template="score {{output}}", @@ -50,19 +51,11 @@ def run_admin_operations() -> str: example_name=EXAMPLE_NAME, ): print_result("ai_integrations.list", client.ai_integrations.list(space="space-offline")) - print_result("ai_integrations.create", client.ai_integrations.create(name="offline-integration", provider=models.AiIntegrationProvider.OPENAI)) + print_result("ai_integrations.create", client.ai_integrations.create(name="offline-integration", provider=models.AiIntegrationProvider.OPEN_AI)) print_result("ai_integrations.get", client.ai_integrations.get(integration="offline-integration", space="space-offline")) print_result("ai_integrations.update", client.ai_integrations.update(integration="offline-integration", space="space-offline", name="renamed-integration")) print_result("ai_integrations.delete", client.ai_integrations.delete(integration="offline-integration", space="space-offline")) print_result("annotation_configs.list", client.annotation_configs.list(space="space-offline")) - print_result( - "annotation_configs.create", - client.annotation_configs.create( - name="quality", - config_type=models.AnnotationConfigType.FREEFORM, - space="space-offline", - ), - ) print_result("annotation_configs.get", client.annotation_configs.get(annotation_config="quality", space="space-offline")) print_result("annotation_configs.delete", client.annotation_configs.delete(annotation_config="quality", space="space-offline")) print_result("annotation_queues.list", client.annotation_queues.list(space="space-offline")) @@ -176,8 +169,12 @@ def run_admin_operations() -> str: "api_keys.create_service_key", client.api_keys.create_service_key( name="offline-service-key", - space="space-offline", - space_role=models.ApiKeySpaceRole.MEMBER, + orgs=[ + OrgBinding( + org_id="org-offline", + spaces=[SpaceBinding(space="space-offline")], + ) + ], ), ) print_result("api_keys.refresh", client.api_keys.refresh(api_key_id="key-offline")) diff --git a/python/tracing/arize/_shared.py b/python/tracing/arize/_shared.py index 371f35c..03291b6 100644 --- a/python/tracing/arize/_shared.py +++ b/python/tracing/arize/_shared.py @@ -60,6 +60,8 @@ def _offline_result(resource: str, method_name: str, kwargs: dict[str, Any]) -> def _make_offline_method(resource: str, method_name: str) -> Callable[..., Any]: def offline_method(self: Any, *args: Any, **kwargs: Any) -> Any: + if kwargs.pop("_respan_force_error", False): + raise RuntimeError(f"offline {resource}.{method_name} failure") return _offline_result(resource, method_name, kwargs) offline_method.__name__ = method_name diff --git a/python/tracing/arize/run_all.py b/python/tracing/arize/run_all.py index 7bdacd9..33b72a4 100644 --- a/python/tracing/arize/run_all.py +++ b/python/tracing/arize/run_all.py @@ -14,9 +14,14 @@ def run() -> None: here = Path(__file__).resolve().parent + failures: list[str] = [] for example in EXAMPLES: print(f"\n### running {example}", flush=True) - subprocess.run([sys.executable, str(here / example)], check=True) + result = subprocess.run([sys.executable, str(here / example)], check=False) + if result.returncode: + failures.append(f"{example} (exit {result.returncode})") + if failures: + raise SystemExit(f"Arize example failures: {', '.join(failures)}") if __name__ == "__main__": diff --git a/python/tracing/autogen/01_assistant_run.py b/python/tracing/autogen/01_assistant_run.py index 10b90aa..52b78e9 100644 --- a/python/tracing/autogen/01_assistant_run.py +++ b/python/tracing/autogen/01_assistant_run.py @@ -29,7 +29,7 @@ @workflow(name=SCRIPT_NAME) -async def run_assistant_agent() -> None: +async def run_assistant_agent() -> str: model_client = OpenAIChatCompletionClient( model=RESPAN_MODEL, api_key=RESPAN_API_KEY, @@ -43,26 +43,39 @@ async def run_assistant_agent() -> None: ) try: - with propagate_attributes( - customer_identifier="autogen-example-user", - thread_identifier="autogen-assistant-thread", - metadata={"script": SCRIPT_NAME}, - ): - result = await agent.run( - task="In one sentence, explain why tracing helps multi-agent apps." - ) - print(result.messages[-1].content) + result = await agent.run( + task="In one sentence, explain why tracing helps multi-agent apps." + ) + return str(result.messages[-1].content) finally: await model_client.close() async def main() -> None: + run_id = os.getenv("RESPAN_EXAMPLE_RUN_ID", f"autogen-{Path(__file__).stem}") respan = Respan( api_key=RESPAN_API_KEY, base_url=RESPAN_BASE_URL, instrumentations=[AutoGenInstrumentor()], - metadata={"example": "autogen-assistant-run", "script": SCRIPT_NAME}, + metadata={ + "example": "autogen-assistant-run", + "script": SCRIPT_NAME, + "run_id": run_id, + }, ) - await run_assistant_agent() + try: + with propagate_attributes( + customer_identifier="autogen-example-user", + thread_identifier="autogen-assistant-thread", + group_identifier=SCRIPT_NAME, + custom_identifier=run_id, + metadata={"script": SCRIPT_NAME, "run_id": run_id}, + ): + print(await run_assistant_agent()) + finally: + respan.shutdown() + print(f"RESPAN_EXAMPLE_RUN_ID={run_id}") + + if __name__ == "__main__": asyncio.run(main()) diff --git a/python/tracing/autogen/02_tool_use.py b/python/tracing/autogen/02_tool_use.py index 493d90b..fd581db 100644 --- a/python/tracing/autogen/02_tool_use.py +++ b/python/tracing/autogen/02_tool_use.py @@ -29,7 +29,7 @@ @workflow(name=SCRIPT_NAME) -async def run_tool_agent() -> None: +async def run_tool_agent() -> str: async def estimate_latency(service: str, requests_per_minute: int) -> str: """Estimate API latency for a service under load.""" baseline_ms = 120 @@ -54,29 +54,42 @@ async def estimate_latency(service: str, requests_per_minute: int) -> str: ) try: - with propagate_attributes( - customer_identifier="autogen-example-user", - thread_identifier="autogen-tool-thread", - metadata={"script": SCRIPT_NAME}, - ): - result = await agent.run( - task=( - "Estimate the p95 latency for the tracing-api service at " - "240 requests per minute." - ) + result = await agent.run( + task=( + "Estimate the p95 latency for the tracing-api service at " + "240 requests per minute." ) - print(result.messages[-1].content) + ) + return str(result.messages[-1].content) finally: await model_client.close() async def main() -> None: + run_id = os.getenv("RESPAN_EXAMPLE_RUN_ID", f"autogen-{Path(__file__).stem}") respan = Respan( api_key=RESPAN_API_KEY, base_url=RESPAN_BASE_URL, instrumentations=[AutoGenInstrumentor()], - metadata={"example": "autogen-tool-use", "script": SCRIPT_NAME}, + metadata={ + "example": "autogen-tool-use", + "script": SCRIPT_NAME, + "run_id": run_id, + }, ) - await run_tool_agent() + try: + with propagate_attributes( + customer_identifier="autogen-example-user", + thread_identifier="autogen-tool-thread", + group_identifier=SCRIPT_NAME, + custom_identifier=run_id, + metadata={"script": SCRIPT_NAME, "run_id": run_id}, + ): + print(await run_tool_agent()) + finally: + respan.shutdown() + print(f"RESPAN_EXAMPLE_RUN_ID={run_id}") + + if __name__ == "__main__": asyncio.run(main()) diff --git a/python/tracing/autogen/03_round_robin_team.py b/python/tracing/autogen/03_round_robin_team.py index 2e2d424..6cc0016 100644 --- a/python/tracing/autogen/03_round_robin_team.py +++ b/python/tracing/autogen/03_round_robin_team.py @@ -31,7 +31,7 @@ @workflow(name=SCRIPT_NAME) -async def run_round_robin_team() -> None: +async def run_round_robin_team() -> str: model_client = OpenAIChatCompletionClient( model=RESPAN_MODEL, api_key=RESPAN_API_KEY, @@ -59,29 +59,44 @@ async def run_round_robin_team() -> None: ) try: - with propagate_attributes( - customer_identifier="autogen-example-user", - thread_identifier="autogen-team-thread", - metadata={"script": SCRIPT_NAME}, - ): - result = await team.run( - task="Plan a small release checklist for a new Python tracing plugin." - ) + result = await team.run( + task="Plan a small release checklist for a new Python tracing plugin." + ) + lines: list[str] = [] for message in result.messages: content = getattr(message, "content", None) if content: - print(f"{message.source}: {content}") + lines.append(f"{message.source}: {content}") + return "\n".join(lines) finally: await model_client.close() async def main() -> None: + run_id = os.getenv("RESPAN_EXAMPLE_RUN_ID", f"autogen-{Path(__file__).stem}") respan = Respan( api_key=RESPAN_API_KEY, base_url=RESPAN_BASE_URL, instrumentations=[AutoGenInstrumentor()], - metadata={"example": "autogen-round-robin-team", "script": SCRIPT_NAME}, + metadata={ + "example": "autogen-round-robin-team", + "script": SCRIPT_NAME, + "run_id": run_id, + }, ) - await run_round_robin_team() + try: + with propagate_attributes( + customer_identifier="autogen-example-user", + thread_identifier="autogen-team-thread", + group_identifier=SCRIPT_NAME, + custom_identifier=run_id, + metadata={"script": SCRIPT_NAME, "run_id": run_id}, + ): + print(await run_round_robin_team()) + finally: + respan.shutdown() + print(f"RESPAN_EXAMPLE_RUN_ID={run_id}") + + if __name__ == "__main__": asyncio.run(main()) diff --git a/python/tracing/autogen/run_all.py b/python/tracing/autogen/run_all.py new file mode 100644 index 0000000..c3a86d4 --- /dev/null +++ b/python/tracing/autogen/run_all.py @@ -0,0 +1,28 @@ +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + + +EXAMPLES = [ + "01_assistant_run.py", + "02_tool_use.py", + "03_round_robin_team.py", +] + + +def run() -> None: + here = Path(__file__).resolve().parent + failures: list[str] = [] + for example in EXAMPLES: + print(f"\n### running {example}", flush=True) + result = subprocess.run([sys.executable, str(here / example)], check=False) + if result.returncode: + failures.append(f"{example} (exit {result.returncode})") + if failures: + raise SystemExit(f"AutoGen example failures: {', '.join(failures)}") + + +if __name__ == "__main__": + run() diff --git a/python/tracing/aws-bedrock/01_invoke_model.py b/python/tracing/aws-bedrock/01_invoke_model.py index 42b555a..c5e6a69 100644 --- a/python/tracing/aws-bedrock/01_invoke_model.py +++ b/python/tracing/aws-bedrock/01_invoke_model.py @@ -11,10 +11,12 @@ deactivate_stubber, get_model_id, maybe_stub_invoke_model, + new_run_id, ) WORKFLOW_NAME = "aws_bedrock_invoke_model" +EXAMPLE_NAME = "01_invoke_model" @workflow(name=WORKFLOW_NAME) @@ -24,22 +26,33 @@ def run_invoke_model() -> str: client = create_bedrock_client() stubber = maybe_stub_invoke_model(client, model_id=model_id, body=body) try: - with propagate_attributes( - trace_group_identifier=WORKFLOW_NAME, - metadata={"example": "invoke_model"}, - ): - response = client.invoke_model( - modelId=model_id, - body=body, - contentType="application/json", - accept="application/json", - ) + response = client.invoke_model( + modelId=model_id, + body=body, + contentType="application/json", + accept="application/json", + ) payload = json.loads(response["body"].read()) return payload["content"][0]["text"] finally: deactivate_stubber(stubber) +def main() -> str: + run_id = new_run_id(EXAMPLE_NAME) + respan = create_respan(example_name=EXAMPLE_NAME, run_id=run_id) + try: + with propagate_attributes( + group_identifier=WORKFLOW_NAME, + custom_identifier=run_id, + metadata={"example": "invoke_model", "run_id": run_id}, + ): + print(run_invoke_model()) + finally: + respan.shutdown() + print(f"RESPAN_EXAMPLE_RUN_ID={run_id}") + return run_id + + if __name__ == "__main__": - respan = create_respan() - print(run_invoke_model()) + main() diff --git a/python/tracing/aws-bedrock/02_converse.py b/python/tracing/aws-bedrock/02_converse.py index 0a26c65..007f4e1 100644 --- a/python/tracing/aws-bedrock/02_converse.py +++ b/python/tracing/aws-bedrock/02_converse.py @@ -8,10 +8,12 @@ deactivate_stubber, get_model_id, maybe_stub_converse, + new_run_id, ) WORKFLOW_NAME = "aws_bedrock_converse" +EXAMPLE_NAME = "02_converse" @workflow(name=WORKFLOW_NAME) @@ -26,20 +28,31 @@ def run_converse() -> str: client = create_bedrock_client() stubber = maybe_stub_converse(client, model_id=model_id, messages=messages) try: - with propagate_attributes( - trace_group_identifier=WORKFLOW_NAME, - metadata={"example": "converse"}, - ): - response = client.converse( - modelId=model_id, - messages=messages, - inferenceConfig={"maxTokens": 96}, - ) + response = client.converse( + modelId=model_id, + messages=messages, + inferenceConfig={"maxTokens": 96}, + ) return response["output"]["message"]["content"][0]["text"] finally: deactivate_stubber(stubber) +def main() -> str: + run_id = new_run_id(EXAMPLE_NAME) + respan = create_respan(example_name=EXAMPLE_NAME, run_id=run_id) + try: + with propagate_attributes( + group_identifier=WORKFLOW_NAME, + custom_identifier=run_id, + metadata={"example": "converse", "run_id": run_id}, + ): + print(run_converse()) + finally: + respan.shutdown() + print(f"RESPAN_EXAMPLE_RUN_ID={run_id}") + return run_id + + if __name__ == "__main__": - respan = create_respan() - print(run_converse()) + main() diff --git a/python/tracing/aws-bedrock/03_converse_stream.py b/python/tracing/aws-bedrock/03_converse_stream.py index ecffb1c..af38c1b 100644 --- a/python/tracing/aws-bedrock/03_converse_stream.py +++ b/python/tracing/aws-bedrock/03_converse_stream.py @@ -8,10 +8,12 @@ deactivate_stubber, get_model_id, maybe_stub_converse_stream, + new_run_id, ) WORKFLOW_NAME = "aws_bedrock_converse_stream" +EXAMPLE_NAME = "03_converse_stream" @workflow(name=WORKFLOW_NAME) @@ -25,28 +27,39 @@ def run_converse_stream() -> str: ] client = create_bedrock_client() stubber = maybe_stub_converse_stream(client, model_id=model_id, messages=messages) + try: + response = client.converse_stream( + modelId=model_id, + messages=messages, + inferenceConfig={"maxTokens": 96}, + ) + + parts: list[str] = [] + for event in response["stream"]: + delta = event.get("contentBlockDelta", {}).get("delta", {}) + text = delta.get("text") + if text: + parts.append(text) + return "".join(parts) + finally: + deactivate_stubber(stubber) + + +def main() -> str: + run_id = new_run_id(EXAMPLE_NAME) + respan = create_respan(example_name=EXAMPLE_NAME, run_id=run_id) try: with propagate_attributes( - trace_group_identifier=WORKFLOW_NAME, - metadata={"example": "converse_stream"}, + group_identifier=WORKFLOW_NAME, + custom_identifier=run_id, + metadata={"example": "converse_stream", "run_id": run_id}, ): - response = client.converse_stream( - modelId=model_id, - messages=messages, - inferenceConfig={"maxTokens": 96}, - ) - - parts: list[str] = [] - for event in response["stream"]: - delta = event.get("contentBlockDelta", {}).get("delta", {}) - text = delta.get("text") - if text: - parts.append(text) - return "".join(parts) + print(run_converse_stream()) finally: - deactivate_stubber(stubber) + respan.shutdown() + print(f"RESPAN_EXAMPLE_RUN_ID={run_id}") + return run_id if __name__ == "__main__": - respan = create_respan() - print(run_converse_stream()) + main() diff --git a/python/tracing/aws-bedrock/04_converse_tool.py b/python/tracing/aws-bedrock/04_converse_tool.py new file mode 100644 index 0000000..0e8289f --- /dev/null +++ b/python/tracing/aws-bedrock/04_converse_tool.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +from respan import propagate_attributes, workflow + +from _shared import ( + create_bedrock_client, + create_respan, + deactivate_stubber, + get_model_id, + maybe_stub_converse_tool, + new_run_id, +) + + +WORKFLOW_NAME = "aws_bedrock_converse_tool" +EXAMPLE_NAME = "04_converse_tool" + + +@workflow(name=WORKFLOW_NAME) +def run_converse_tool() -> str: + model_id = get_model_id() + messages = [ + { + "role": "user", + "content": [{"text": "What is the weather in Tokyo?"}], + } + ] + tool_config = { + "tools": [ + { + "toolSpec": { + "name": "get_weather", + "description": "Get the current weather for a city.", + "inputSchema": { + "json": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + } + }, + } + } + ] + } + client = create_bedrock_client() + stubber = maybe_stub_converse_tool( + client, + model_id=model_id, + messages=messages, + tool_config=tool_config, + ) + try: + response = client.converse( + modelId=model_id, + messages=messages, + toolConfig=tool_config, + inferenceConfig={"maxTokens": 96}, + ) + tool_use = response["output"]["message"]["content"][0]["toolUse"] + return f"{tool_use['name']}({tool_use['input']})" + finally: + deactivate_stubber(stubber) + + +def main() -> str: + run_id = new_run_id(EXAMPLE_NAME) + respan = create_respan(example_name=EXAMPLE_NAME, run_id=run_id) + try: + with propagate_attributes( + group_identifier=WORKFLOW_NAME, + custom_identifier=run_id, + metadata={"example": "converse_tool", "run_id": run_id}, + ): + print(run_converse_tool()) + finally: + respan.shutdown() + print(f"RESPAN_EXAMPLE_RUN_ID={run_id}") + return run_id + + +if __name__ == "__main__": + main() diff --git a/python/tracing/aws-bedrock/05_converse_error.py b/python/tracing/aws-bedrock/05_converse_error.py new file mode 100644 index 0000000..8b155ce --- /dev/null +++ b/python/tracing/aws-bedrock/05_converse_error.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +from botocore.exceptions import ClientError +from respan import propagate_attributes, workflow + +from _shared import ( + create_bedrock_client, + create_respan, + deactivate_stubber, + get_model_id, + maybe_stub_converse_error, + new_run_id, +) + + +WORKFLOW_NAME = "aws_bedrock_converse_error" +EXAMPLE_NAME = "05_converse_error" + + +@workflow(name=WORKFLOW_NAME) +def run_converse_error() -> str: + model_id = get_model_id() + messages = [ + { + "role": "user", + "content": [{"text": "Trigger the deterministic missing-model route."}], + } + ] + client = create_bedrock_client() + stubber = maybe_stub_converse_error( + client, + model_id=model_id, + messages=messages, + ) + try: + client.converse(modelId=model_id, messages=messages) + except ClientError as error: + status_code = error.response["ResponseMetadata"]["HTTPStatusCode"] + return f"expected Bedrock error {status_code}" + finally: + deactivate_stubber(stubber) + raise AssertionError("the deterministic Bedrock error route unexpectedly succeeded") + + +def main() -> str: + run_id = new_run_id(EXAMPLE_NAME) + respan = create_respan(example_name=EXAMPLE_NAME, run_id=run_id) + try: + with propagate_attributes( + group_identifier=WORKFLOW_NAME, + custom_identifier=run_id, + metadata={"example": "converse_error", "run_id": run_id}, + ): + print(run_converse_error()) + finally: + respan.shutdown() + print(f"RESPAN_EXAMPLE_RUN_ID={run_id}") + return run_id + + +if __name__ == "__main__": + main() diff --git a/python/tracing/aws-bedrock/_shared.py b/python/tracing/aws-bedrock/_shared.py index 4dec785..dee99bf 100644 --- a/python/tracing/aws-bedrock/_shared.py +++ b/python/tracing/aws-bedrock/_shared.py @@ -3,6 +3,7 @@ import io import json import os +import uuid from pathlib import Path from typing import Any @@ -37,11 +38,19 @@ def should_use_stubs() -> bool: return not (os.getenv("AWS_ACCESS_KEY_ID") or os.getenv("AWS_PROFILE")) -def create_respan() -> Respan: +def new_run_id(example_name: str) -> str: + return os.getenv("RESPAN_EXAMPLE_RUN_ID") or f"{example_name}-{uuid.uuid4().hex[:10]}" + + +def create_respan(*, example_name: str, run_id: str) -> Respan: return Respan( app_name="aws-bedrock-examples", instrumentations=[AWSBedrockInstrumentor()], - metadata={"example_set": "aws-bedrock"}, + metadata={ + "example_set": "aws-bedrock", + "example_name": example_name, + "run_id": run_id, + }, environment=os.getenv("RESPAN_ENVIRONMENT", "examples"), ) @@ -161,7 +170,15 @@ def maybe_stub_converse_stream( "contentBlockDelta": { "contentBlockIndex": 0, "delta": {"text": "Stubbed stream."}, - } + }, + "metadata": { + "usage": { + "inputTokens": 12, + "outputTokens": 13, + "totalTokens": 25, + }, + "metrics": {"latencyMs": 7}, + }, }, "ResponseMetadata": {"HTTPStatusCode": 200}, }, @@ -171,6 +188,71 @@ def maybe_stub_converse_stream( return stubber +def maybe_stub_converse_tool( + client, + *, + model_id: str, + messages: list[dict[str, Any]], + tool_config: dict[str, Any], +) -> Stubber | None: + if not should_use_stubs(): + return None + + stubber = Stubber(client) + stubber.add_response( + "converse", + { + "output": { + "message": { + "role": "assistant", + "content": [ + { + "toolUse": { + "toolUseId": "tooluse-weather-1", + "name": "get_weather", + "input": {"city": "Tokyo"}, + } + } + ], + } + }, + "stopReason": "tool_use", + "usage": {"inputTokens": 14, "outputTokens": 6, "totalTokens": 20}, + "metrics": {"latencyMs": 8}, + "ResponseMetadata": {"HTTPStatusCode": 200}, + }, + { + "modelId": model_id, + "messages": messages, + "toolConfig": tool_config, + "inferenceConfig": {"maxTokens": 96}, + }, + ) + stubber.activate() + return stubber + + +def maybe_stub_converse_error( + client, + *, + model_id: str, + messages: list[dict[str, Any]], +) -> Stubber | None: + if not should_use_stubs(): + return None + + stubber = Stubber(client) + stubber.add_client_error( + "converse", + service_error_code="ResourceNotFoundException", + service_message="The requested model was not found.", + http_status_code=404, + expected_params={"modelId": model_id, "messages": messages}, + ) + stubber.activate() + return stubber + + def deactivate_stubber(stubber: Stubber | None) -> None: if stubber is not None: stubber.deactivate() diff --git a/python/tracing/aws-bedrock/run_all.py b/python/tracing/aws-bedrock/run_all.py new file mode 100644 index 0000000..4594393 --- /dev/null +++ b/python/tracing/aws-bedrock/run_all.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + + +EXAMPLES = [ + "01_invoke_model.py", + "02_converse.py", + "03_converse_stream.py", + "04_converse_tool.py", + "05_converse_error.py", +] + + +def run() -> None: + here = Path(__file__).resolve().parent + failures: list[str] = [] + for example in EXAMPLES: + print(f"\n### running {example}", flush=True) + result = subprocess.run([sys.executable, str(here / example)], check=False) + if result.returncode: + failures.append(f"{example} (exit {result.returncode})") + if failures: + raise SystemExit(f"AWS Bedrock example failures: {', '.join(failures)}") + + +if __name__ == "__main__": + run()