Skip to content

fix(examples-py): validate SageMaker, Semantic Kernel, and Smolagents OTel 2.x spans #73

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Nightingalelyy wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(examples-py): validate SageMaker, Semantic Kernel, and Smolagents OTel 2.x spans #73

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 13 additions & 9 deletions python/tracing/sagemaker/01_invoke_endpoint_text.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,5 @@
from __future__ import annotations

from respan import workflow

from _shared import (
custom_attributes,
endpoint_name,
Expand All @@ -17,15 +15,17 @@
stubbed_response,
workflow_name,
)
from respan import workflow

EXAMPLE_NAME = "invoke-endpoint-text"


@workflow(name=workflow_name(EXAMPLE_NAME))
def _invoke_text_workflow(client) -> dict:
def _invoke_text_workflow(prompt: str) -> dict:
client = make_client()
request_body = json_bytes(
{
"inputs": "Reply with one concise sentence about SageMaker observability.",
"inputs": prompt,
"parameters": {"max_new_tokens": 32, "temperature": 0.1},
}
)
Expand All @@ -48,21 +48,25 @@ def _invoke_text_workflow(client) -> dict:
"ContentType": "application/json",
}

with stubbed_response(client, "invoke_endpoint", response, params):
result = client.invoke_endpoint(**params)
return {"response": read_json_body(result)}
try:
with stubbed_response(client, "invoke_endpoint", response, params):
result = client.invoke_endpoint(**params)
return {"response": read_json_body(result)}
finally:
client.close()


def run_invoke_endpoint_text() -> None:
respan = make_respan(EXAMPLE_NAME)
client = make_client()
custom_identifier = make_custom_identifier(EXAMPLE_NAME)
result: dict = {}

try:
with example_attributes(EXAMPLE_NAME, custom_identifier):
print_run_header(EXAMPLE_NAME, custom_identifier)
result = _invoke_text_workflow(client)
result = _invoke_text_workflow(
"Reply with one concise sentence about SageMaker observability."
)
finally:
respan.shutdown()

Expand Down
130 changes: 67 additions & 63 deletions python/tracing/sagemaker/02_invoke_endpoint_chat_tools.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,6 @@

import json

from respan import tool, workflow

from _shared import (
custom_attributes,
endpoint_name,
Expand All @@ -19,6 +17,7 @@
stubbed_response,
workflow_name,
)
from respan import tool, workflow

EXAMPLE_NAME = "invoke-endpoint-chat-tools"

Expand Down Expand Up @@ -49,13 +48,13 @@ def _extract_tool_call(response_payload: dict) -> dict:
first_choice = choices[0]
message = first_choice.get("message") if isinstance(first_choice, dict) else None
if not isinstance(message, dict):
raise RuntimeError("SageMaker tool example expected an assistant message.")
raise TypeError("SageMaker tool example expected an assistant message.")
tool_calls = message.get("tool_calls")
if not isinstance(tool_calls, list) or not tool_calls:
raise RuntimeError("SageMaker tool example expected an assistant tool call.")
tool_call = tool_calls[0]
if not isinstance(tool_call, dict):
raise RuntimeError("SageMaker tool example received an invalid tool call.")
raise TypeError("SageMaker tool example received an invalid tool call.")
return tool_call


Expand All @@ -70,10 +69,11 @@ def _tool_arguments(tool_call: dict) -> dict:


@workflow(name=workflow_name(EXAMPLE_NAME))
def _invoke_chat_tools_workflow(client) -> dict:
def _invoke_chat_tools_workflow(city: str) -> dict:
client = make_client()
user_message = {
"role": "user",
"content": "What is the weather in Tokyo? Use the tool.",
"content": f"What is the weather in {city}? Use the tool.",
}
first_body = {
"messages": [user_message],
Expand All @@ -100,7 +100,7 @@ def _invoke_chat_tools_workflow(client) -> dict:
"type": "function",
"function": {
"name": "get_weather",
"arguments": "{\"city\": \"Tokyo\"}",
"arguments": '{"city": "Tokyo"}',
},
}
],
Expand All @@ -117,73 +117,77 @@ def _invoke_chat_tools_workflow(client) -> dict:
"ContentType": "application/json",
}

with stubbed_response(client, "invoke_endpoint", first_response, first_params):
first_result = read_json_body(client.invoke_endpoint(**first_params))

tool_call = _extract_tool_call(first_result)
tool_result = get_weather(**_tool_arguments(tool_call))

assistant_message = first_result["choices"][0]["message"]
second_body = {
"messages": [
user_message,
assistant_message,
{
"role": "tool",
"tool_call_id": tool_call.get("id"),
"content": tool_result,
},
],
"tools": [TOOL_SCHEMA],
}
second_params = {
"EndpointName": endpoint_name(),
"Body": json_bytes(second_body),
"ContentType": "application/json",
"Accept": "application/json",
"CustomAttributes": custom_attributes(),
}
second_response = {
"Body": streaming_body(
{
"choices": [
{
"message": {
"role": "assistant",
"content": "Tokyo is sunny and 22 C.",
}
}
],
"usage": {
"prompt_tokens": 32,
"completion_tokens": 7,
"total_tokens": 39,
try:
with stubbed_response(client, "invoke_endpoint", first_response, first_params):
first_result = read_json_body(client.invoke_endpoint(**first_params))

tool_call = _extract_tool_call(first_result)
tool_result = get_weather(**_tool_arguments(tool_call))

assistant_message = first_result["choices"][0]["message"]
second_body = {
"messages": [
user_message,
assistant_message,
{
"role": "tool",
"tool_call_id": tool_call.get("id"),
"content": tool_result,
},
}
),
"ContentType": "application/json",
}

with stubbed_response(client, "invoke_endpoint", second_response, second_params):
final_result = read_json_body(client.invoke_endpoint(**second_params))

return {
"tool_call": tool_call,
"tool_result": tool_result,
"final_response": final_result,
}
],
"tools": [TOOL_SCHEMA],
}
second_params = {
"EndpointName": endpoint_name(),
"Body": json_bytes(second_body),
"ContentType": "application/json",
"Accept": "application/json",
"CustomAttributes": custom_attributes(),
}
second_response = {
"Body": streaming_body(
{
"choices": [
{
"message": {
"role": "assistant",
"content": f"{city} is sunny and 22 C.",
}
}
],
"usage": {
"prompt_tokens": 32,
"completion_tokens": 7,
"total_tokens": 39,
},
}
),
"ContentType": "application/json",
}

with stubbed_response(
client, "invoke_endpoint", second_response, second_params
):
final_result = read_json_body(client.invoke_endpoint(**second_params))

return {
"tool_call": tool_call,
"tool_result": tool_result,
"final_response": final_result,
}
finally:
client.close()


def run_invoke_endpoint_chat_tools() -> None:
respan = make_respan(EXAMPLE_NAME)
client = make_client()
custom_identifier = make_custom_identifier(EXAMPLE_NAME)
result: dict = {}

try:
with example_attributes(EXAMPLE_NAME, custom_identifier):
print_run_header(EXAMPLE_NAME, custom_identifier)
result = _invoke_chat_tools_workflow(client)
result = _invoke_chat_tools_workflow("Tokyo")
finally:
respan.shutdown()

Expand Down
32 changes: 17 additions & 15 deletions python/tracing/sagemaker/03_invoke_endpoint_stream.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,8 @@

import json

from respan import workflow

from _shared import (
collect_stream_text,
custom_attributes,
endpoint_name,
example_attributes,
Expand All @@ -16,15 +15,16 @@
print_run_header,
stubbed_response,
workflow_name,
collect_stream_text,
)
from respan import workflow

EXAMPLE_NAME = "invoke-endpoint-stream"


@workflow(name=workflow_name(EXAMPLE_NAME))
def _invoke_stream_workflow(client) -> dict:
request_body = json_bytes({"inputs": "Stream a concise SageMaker sentence."})
def _invoke_stream_workflow(prompt: str) -> dict:
client = make_client()
request_body = json_bytes({"inputs": prompt})
params = {
"EndpointName": endpoint_name(),
"Body": request_body,
Expand All @@ -48,26 +48,28 @@ def _invoke_stream_workflow(client) -> dict:
"ContentType": "application/json",
}

with stubbed_response(
client,
"invoke_endpoint_with_response_stream",
response,
params,
):
result = client.invoke_endpoint_with_response_stream(**params)
return {"stream_text": collect_stream_text(result)}
try:
with stubbed_response(
client,
"invoke_endpoint_with_response_stream",
response,
params,
):
result = client.invoke_endpoint_with_response_stream(**params)
return {"stream_text": collect_stream_text(result)}
finally:
client.close()


def run_invoke_endpoint_stream() -> None:
respan = make_respan(EXAMPLE_NAME)
client = make_client()
custom_identifier = make_custom_identifier(EXAMPLE_NAME)
result: dict = {}

try:
with example_attributes(EXAMPLE_NAME, custom_identifier):
print_run_header(EXAMPLE_NAME, custom_identifier)
result = _invoke_stream_workflow(client)
result = _invoke_stream_workflow("Stream a concise SageMaker sentence.")
finally:
respan.shutdown()

Expand Down
26 changes: 14 additions & 12 deletions python/tracing/sagemaker/04_invoke_endpoint_async.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,5 @@
from __future__ import annotations

from respan import workflow

from _shared import (
custom_attributes,
endpoint_name,
Expand All @@ -14,15 +12,17 @@
stubbed_response,
workflow_name,
)
from respan import workflow

EXAMPLE_NAME = "invoke-endpoint-async"


@workflow(name=workflow_name(EXAMPLE_NAME))
def _invoke_async_workflow(client) -> dict:
def _invoke_async_workflow(input_location: str) -> dict:
client = make_client()
params = {
"EndpointName": endpoint_name(),
"InputLocation": "s3://respan-sagemaker-example/input.json",
"InputLocation": input_location,
"ContentType": "application/json",
"Accept": "application/json",
"CustomAttributes": custom_attributes(),
Expand All @@ -32,24 +32,26 @@ def _invoke_async_workflow(client) -> dict:
"OutputLocation": "s3://respan-sagemaker-example/output.json",
}

with stubbed_response(client, "invoke_endpoint_async", response, params):
result = client.invoke_endpoint_async(**params)
return {
"inference_id": result.get("InferenceId"),
"output_location": result.get("OutputLocation"),
}
try:
with stubbed_response(client, "invoke_endpoint_async", response, params):
result = client.invoke_endpoint_async(**params)
return {
"inference_id": result.get("InferenceId"),
"output_location": result.get("OutputLocation"),
}
finally:
client.close()


def run_invoke_endpoint_async() -> None:
respan = make_respan(EXAMPLE_NAME)
client = make_client()
custom_identifier = make_custom_identifier(EXAMPLE_NAME)
result: dict = {}

try:
with example_attributes(EXAMPLE_NAME, custom_identifier):
print_run_header(EXAMPLE_NAME, custom_identifier)
result = _invoke_async_workflow(client)
result = _invoke_async_workflow("s3://respan-sagemaker-example/input.json")
finally:
respan.shutdown()

Expand Down
Loading