here is the script for all 1-3 only difference in 2 and 3 is instead of litellm.completion i do cached_litellm_obj.completion / cached_litellm_obj.handle_completion
# %%
import asyncio
from hx_platform.tools.local_settings_util import load_settings_to_env
load_settings_to_env()
# %%
import os
import logging
from openinference.instrumentation.litellm import LiteLLMInstrumentor
from openinference.semconv.resource import ResourceAttributes
from opentelemetry import trace as trace_api
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from opentelemetry.sdk import trace as trace_sdk
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from openinference.instrumentation.openai import OpenAIInstrumentor
import litellm
# LlamaIndexInstrumentor().instrument()
LiteLLMInstrumentor().instrument()
# OpenAIInstrumentor().instrument()
# Use service_name if provided, otherwise fall back to project_name
resource_attributes = {}
resource_attributes[ResourceAttributes.PROJECT_NAME] = 'test4'
resource = Resource.create(resource_attributes)
tracer_provider = trace_sdk.TracerProvider(resource=resource)
multi_span_processor = trace_sdk.ConcurrentMultiSpanProcessor()
if os.environ.get("PHOENIX_COLLECTOR_ENDPOINT"):
# Add OTLP exporter
multi_span_processor.add_span_processor(
SimpleSpanProcessor(
OTLPSpanExporter("http://localhost:6006/v1/traces")
)
)
else:
logging.info(
"PHOENIX_COLLECTOR_ENDPOINT not set. Phoenix Tracing will not be enabled."
)
tracer_provider.add_span_processor(multi_span_processor)
trace_api.set_tracer_provider(tracer_provider)
MODEL_NAME = 'bedrock/anthropic.claude-3-7-sonnet-20250219-v1:0'
MODEL_ID = 'arn:aws:bedrock:<YOUR_AWS_REGION>:<YOUR_AWS_ACCOUNT_ID>:inference-profile/<MODEL_NAME>:<MODEL_VERSION>'
AWS_ACCESS_KEY_ID = os.environ.get("AWS_ACCESS_KEY_ID")
AWS_SECRET_ACCESS_KEY = os.environ.get("AWS_SECRET_ACCESS_KEY")
AWS_REGION = os.environ.get("AWS_REGION_NAME")
# %%
TOOLS = [
{
"type": "function",
"function": {
"name": "search_knowledge_base",
"description": "Search the knowledge base for relevant information",
"parameters": {
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "The search query"
},
"limit": {
"type": "integer",
"description": "Maximum number of results to return",
"default": 5
}
},
"required": ["query"]
}
}
},
{
"type": "function",
"function": {
"name": "get_incident_details",
"description": "Get details about a specific incident",
"parameters": {
"type": "object",
"properties": {
"incident_id": {
"type": "string",
"description": "The unique identifier for the incident"
}
},
"required": ["incident_id"]
}
}
}
]
INPUT_MESSAGES = [
{"role": "system", "content": "You are an AI assistant helping with incident management. You can search for information and retrieve incident details to help users."},
{"role": "user", "content": "What are the recent critical incidents in our production environment?"}
]
tracer = trace_api.get_tracer(__name__)
# %%
## No Stream Sync Completion Claude With Thinking + Tools
with tracer.start_as_current_span("litellm.completion_sync_with_tools_and_thinking") as span:
out = litellm.completion(
model=MODEL_NAME,
model_id=MODEL_ID,
api_key=None,
aws_access_key_id=AWS_ACCESS_KEY_ID,
aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
aws_region_name=AWS_REGION,
messages=INPUT_MESSAGES,
thinking={"type": "enabled", "budget_tokens": 2000},
temperature=1
)
print(out)
# %%
# No Stream Async Completion Claude With Thinking + Tools
with tracer.start_as_current_span("litellm.acompletion_async_with_tools_and_thinking") as span:
out2 = asyncio.run(
litellm.acompletion(
model=MODEL_NAME,
model_id=MODEL_ID,
api_key=None,
aws_access_key_id=AWS_ACCESS_KEY_ID,
aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
aws_region_name=AWS_REGION,
messages=INPUT_MESSAGES,
thinking={"type": "enabled", "budget_tokens": 2000},
temperature=1,
tools=TOOLS
)
)
print(out2)
# %%
with tracer.start_as_current_span("litellm.completion_sync_with_tools_and_thinking") as span:
out3 = litellm.completion(
model=MODEL_NAME,
model_id=MODEL_ID,
api_key=None,
aws_access_key_id=AWS_ACCESS_KEY_ID,
aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
aws_region_name=AWS_REGION,
messages=INPUT_MESSAGES,
temperature=1,
thinking={"type": "enabled", "budget_tokens": 2000},
stream=True,
tools=TOOLS
)
for chunk in out3:
print(chunk)
# %%
async def run():
with tracer.start_as_current_span("litellm.acompletion_async_with_tools_and_thinking") as span:
out4 = await litellm.acompletion(
model=MODEL_NAME,
model_id=MODEL_ID,
api_key=None,
aws_access_key_id=AWS_ACCESS_KEY_ID,
aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
aws_region_name=AWS_REGION,
messages=INPUT_MESSAGES,
temperature=1,
thinking={"type": "enabled", "budget_tokens": 2000},
stream=True,
tools=TOOLS
)
async for chunk in out4:
print(chunk)
asyncio.run(run())
# %% 
