Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
Original file line number Diff line number Diff line change
@@ -1,43 +1,50 @@
import os
from typing import Dict, List, override
import json
from typing import Dict, List, override, Any
from dotenv import load_dotenv

from dotenv import load_dotenv
from agentex.lib.utils.model_utils import BaseModel
from mcp import StdioServerParameters
from temporalio import workflow
from agents import ModelSettings
from agents import ModelSettings, RunContextWrapper
from openai.types.shared import Reasoning

from agentex.lib import adk
from agentex.lib.types.acp import CreateTaskParams, SendEventParams
from agentex.lib.core.temporal.workflows.workflow import BaseWorkflow
from agentex.lib.core.temporal.types.workflow import SignalName
from agentex.lib.utils.logging import make_logger
from agentex.lib.core.tracing.tracing_processor_manager import add_tracing_processor_config
from agentex.lib.core.tracing.tracing_processor_manager import (
add_tracing_processor_config,
)
from agentex.lib.types.tracing import SGPTracingProcessorConfig
from agentex.lib.environment_variables import EnvironmentVariables
from agentex.types.text_content import TextContent
from agentex.lib.core.temporal.activities.adk.providers.openai_activities import (
FunctionTool,
)

environment_variables = EnvironmentVariables.refresh()
load_dotenv(dotenv_path=".env")

add_tracing_processor_config(SGPTracingProcessorConfig(
sgp_api_key=os.environ.get("SCALE_GP_API_KEY", ""),
sgp_account_id=os.environ.get("SCALE_GP_ACCOUNT_ID", ""),
))
add_tracing_processor_config(
SGPTracingProcessorConfig(
sgp_api_key=os.environ.get("SCALE_GP_API_KEY", ""),
sgp_account_id=os.environ.get("SCALE_GP_ACCOUNT_ID", ""),
)
)

if environment_variables.WORKFLOW_NAME is None:
if not environment_variables.WORKFLOW_NAME:
raise ValueError("Environment variable WORKFLOW_NAME is not set")

if environment_variables.AGENT_NAME is None:
if not environment_variables.AGENT_NAME:
raise ValueError("Environment variable AGENT_NAME is not set")

logger = make_logger(__name__)


class StateModel(BaseModel):
input_list: List[Dict]
input_list: List[Dict[str, Any]]
turn_number: int


Expand All @@ -49,44 +56,139 @@ class StateModel(BaseModel):
StdioServerParameters(
command="uvx",
args=["openai-websearch-mcp"],
env={
"OPENAI_API_KEY": os.environ.get("OPENAI_API_KEY", "")
}
env={"OPENAI_API_KEY": os.environ.get("OPENAI_API_KEY", "")},
),
]


async def calculator(context: RunContextWrapper, args: str) -> str:
"""
Simple calculator that can perform basic arithmetic operations.

Args:
context: The run context wrapper
args: JSON string containing the operation and operands

Returns:
String representation of the calculation result
"""
try:
# Parse the JSON arguments
parsed_args = json.loads(args)
operation = parsed_args.get("operation")
a = parsed_args.get("a")
b = parsed_args.get("b")

if operation is None or a is None or b is None:
return (
"Error: Missing required parameters. "
"Please provide 'operation', 'a', and 'b'."
)

# Convert to numbers
try:
a = float(a)
b = float(b)
except (ValueError, TypeError):
return "Error: 'a' and 'b' must be valid numbers."

# Perform the calculation
if operation == "add":
result = a + b
elif operation == "subtract":
result = a - b
elif operation == "multiply":
result = a * b
elif operation == "divide":
if b == 0:
return "Error: Division by zero is not allowed."
result = a / b
else:
supported_ops = "add, subtract, multiply, divide"
return (
f"Error: Unknown operation '{operation}'. "
f"Supported operations: {supported_ops}."
)

# Format the result nicely
if result == int(result):
return f"The result of {a} {operation} {b} is {int(result)}"
else:
formatted = f"{result:.6f}".rstrip("0").rstrip(".")
return f"The result of {a} {operation} {b} is {formatted}"

except json.JSONDecodeError:
return "Error: Invalid JSON format in arguments."
except Exception as e:
return f"Error: An unexpected error occurred: {str(e)}"


# Create the calculator tool
CALCULATOR_TOOL = FunctionTool(
name="calculator",
description=(
"Performs basic arithmetic operations (add, subtract, multiply, divide) "
"on two numbers."
),
params_json_schema={
"type": "object",
"properties": {
"operation": {
"type": "string",
"enum": ["add", "subtract", "multiply", "divide"],
"description": "The arithmetic operation to perform",
},
"a": {"type": "number", "description": "The first number"},
"b": {"type": "number", "description": "The second number"},
},
"required": ["operation", "a", "b"],
"additionalProperties": False,
},
strict_json_schema=True,
on_invoke_tool=calculator,
)


@workflow.defn(name=environment_variables.WORKFLOW_NAME)
class At010AgentChatWorkflow(BaseWorkflow):
"""
Minimal async workflow template for AgentEx Temporal agents.
"""

def __init__(self):
super().__init__(display_name=environment_variables.AGENT_NAME)
self._complete_task = False
self._state = None
self._state: StateModel | None = None

@workflow.signal(name=SignalName.RECEIVE_EVENT)
@override
async def on_task_event_send(self, params: SendEventParams) -> None:
logger.info(f"Received task message instruction: {params}")

if not params.event.content:
return
if params.event.content.type != "text":
raise ValueError(f"Expected text message, got {params.event.content.type}")

if params.event.content.author != "user":
raise ValueError(f"Expected user message, got {params.event.content.author}")

raise ValueError(
f"Expected user message, got {params.event.content.author}"
)

if self._state is None:
raise ValueError("State is not initialized")

# Increment the turn number
self._state.turn_number += 1
# Add the new user message to the message history
self._state.input_list.append({"role": "user", "content": params.event.content.content})
self._state.input_list.append(
{"role": "user", "content": params.event.content.content}
)

async with adk.tracing.span(
trace_id=params.task.id,
name=f"Turn {self._state.turn_number}",
input=self._state
input=self._state,
) as span:
# Echo back the user's message so it shows up in the UI. This is not done by default so the agent developer has full control over what is shown to the user.
await adk.messages.create(
Expand All @@ -102,7 +204,15 @@ async def on_task_event_send(self, params: SendEventParams) -> None:
trace_id=params.task.id,
content=TextContent(
author="agent",
content="Hey, sorry I'm unable to respond to your message because you're running this example without an OpenAI API key. Please set the OPENAI_API_KEY environment variable to run this example. Do this by either by adding a .env file to the project/ directory or by setting the environment variable in your terminal.",
content=(
"Hey, sorry I'm unable to respond to your message "
"because you're running this example without an "
"OpenAI API key. Please set the OPENAI_API_KEY "
"environment variable to run this example. Do this "
"by either by adding a .env file to the project/ "
"directory or by setting the environment variable "
"in your terminal."
),
),
parent_span_id=span.id if span else None,
)
Expand All @@ -115,22 +225,33 @@ async def on_task_event_send(self, params: SendEventParams) -> None:
input_list=self._state.input_list,
mcp_server_params=MCP_SERVERS,
agent_name="Tool-Enabled Assistant",
agent_instructions="""You are a helpful assistant that can answer questions using various tools.
You have access to sequential thinking and web search capabilities through MCP servers.
Use these tools when appropriate to provide accurate and well-reasoned responses.""",
agent_instructions=(
"You are a helpful assistant that can answer questions "
"using various tools. You have access to sequential "
"thinking and web search capabilities through MCP servers, "
"as well as a calculator tool for performing basic "
"arithmetic operations. Use these tools when appropriate "
"to provide accurate and well-reasoned responses."
),
parent_span_id=span.id if span else None,
model="o4-mini",
model_settings=ModelSettings(
# Include reasoning items in the response (IDs, summaries)
# response_include=["reasoning.encrypted_content"],
# Ask the model to include a short reasoning summary
reasoning=Reasoning(effort="medium", summary="auto"),
)
),
tools=[CALCULATOR_TOOL],
)
self._state.input_list = run_result.final_input_list
if self._state:
# Update the state with the final input list if available
final_list = getattr(run_result, "final_input_list", None)
if final_list is not None:
self._state.input_list = final_list

# Set the span output to the state for the next turn
span.output = self._state
if span and self._state:
span.output = self._state.model_dump()

@workflow.run
@override
Expand All @@ -151,5 +272,5 @@ async def on_task_create(self, params: CreateTaskParams) -> None:

await workflow.wait_condition(
lambda: self._complete_task,
timeout=None, # Set a timeout if you want to prevent the task from running indefinitely. Generally this is not needed. Temporal can run hundreds of millions of workflows in parallel and more. Only do this if you have a specific reason to do so.
timeout=None, # Set a timeout if you want to prevent the task from running indefinitely. Generally this is not needed. Temporal can run hundreds of millions of workflows in parallel and more. Only do this if you have a specific reason to do so.
)
3 changes: 2 additions & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,8 @@ dependencies = [
"pytest-asyncio>=1.0.0",
"scale-gp-beta==0.1.0a20",
"ipykernel>=6.29.5",
"openai>=1.99.9",
"openai==1.99.9", # anything higher than 1.99.9 breaks litellm - https://github.com/BerriAI/litellm/issues/13711
"cloudpickle>=3.1.1",
]
requires-python = ">= 3.12,<4"
classifiers = [
Expand Down
5 changes: 4 additions & 1 deletion requirements-dev.lock
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,8 @@ click==8.2.1
# via litellm
# via typer
# via uvicorn
cloudpickle==3.1.1
# via agentex-sdk
colorama==0.4.6
# via griffe
colorlog==6.7.0
Expand Down Expand Up @@ -179,9 +181,10 @@ oauthlib==3.3.1
# via kubernetes
# via requests-oauthlib
openai==1.99.9
# via agentex-sdk
# via litellm
# via openai-agents
openai-agents==0.2.6
openai-agents==0.2.7
# via agentex-sdk
packaging==23.2
# via huggingface-hub
Expand Down
5 changes: 4 additions & 1 deletion requirements.lock
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,8 @@ click==8.2.1
# via litellm
# via typer
# via uvicorn
cloudpickle==3.1.1
# via agentex-sdk
colorama==0.4.6
# via griffe
comm==0.2.3
Expand Down Expand Up @@ -162,9 +164,10 @@ oauthlib==3.3.1
# via kubernetes
# via requests-oauthlib
openai==1.99.9
# via agentex-sdk
# via litellm
# via openai-agents
openai-agents==0.2.6
openai-agents==0.2.7
# via agentex-sdk
packaging==25.0
# via huggingface-hub
Expand Down
Loading
Loading