From b9f95472478923393ce45d56922cf4808ee61d84 Mon Sep 17 00:00:00 2001 From: Braelyn Boynton Date: Wed, 21 Aug 2024 13:39:25 -0700 Subject: [PATCH] event decorators have default names (#367) * record action uses function name as default * record tool uses function name as default * docstring --- agentops/decorators.py | 36 ++++++++++++++++++++++++++++-------- tests/test_record_action.py | 21 +++++++++++++++++++++ tests/test_record_tool.py | 21 +++++++++++++++++++++ 3 files changed, 70 insertions(+), 8 deletions(-) diff --git a/agentops/decorators.py b/agentops/decorators.py index 3c74ce8d..c1930b44 100644 --- a/agentops/decorators.py +++ b/agentops/decorators.py @@ -17,7 +17,7 @@ def record_function(event_name: str): return record_action(event_name) -def record_action(event_name: str): +def record_action(event_name: Optional[str] = None): """ Decorator to record an event before and after a function call. Usage: @@ -25,7 +25,7 @@ def record_action(event_name: str): function being decorated. Additionally, timing information about the action is recorded Args: - event_name (str): The name of the event to record. + event_name (optional, str): The name of the event to record. """ def decorator(func): @@ -53,11 +53,16 @@ async def async_wrapper(*args, session: Optional[Session] = None, **kwargs): arg_values.update(dict(zip(arg_names, args))) arg_values.update(kwargs) + if not event_name: + action_type = func.__name__ + else: + action_type = event_name + event = ActionEvent( params=arg_values, init_timestamp=init_time, agent_id=check_call_stack_for_agent_id(), - action_type=event_name, + action_type=action_type, ) try: @@ -114,11 +119,16 @@ def sync_wrapper(*args, session: Optional[Session] = None, **kwargs): arg_values.update(dict(zip(arg_names, args))) arg_values.update(kwargs) + if not event_name: + action_type = func.__name__ + else: + action_type = event_name + event = ActionEvent( params=arg_values, init_timestamp=init_time, agent_id=check_call_stack_for_agent_id(), - action_type=event_name, + action_type=action_type, ) try: @@ -153,7 +163,7 @@ def sync_wrapper(*args, session: Optional[Session] = None, **kwargs): return decorator -def record_tool(tool_name: str): +def record_tool(tool_name: Optional[str] = None): """ Decorator to record a tool use event before and after a function call. Usage: @@ -161,7 +171,7 @@ def record_tool(tool_name: str): function being decorated. Additionally, timing information about the action is recorded Args: - tool_name (str): The name of the event to record. + tool_name (optional, str): The name of the event to record. """ def decorator(func): @@ -189,11 +199,16 @@ async def async_wrapper(*args, session: Optional[Session] = None, **kwargs): arg_values.update(dict(zip(arg_names, args))) arg_values.update(kwargs) + if not tool_name: + name = func.__name__ + else: + name = tool_name + event = ToolEvent( params=arg_values, init_timestamp=init_time, agent_id=check_call_stack_for_agent_id(), - name=tool_name, + name=name, ) try: @@ -250,11 +265,16 @@ def sync_wrapper(*args, session: Optional[Session] = None, **kwargs): arg_values.update(dict(zip(arg_names, args))) arg_values.update(kwargs) + if not tool_name: + name = func.__name__ + else: + name = tool_name + event = ToolEvent( params=arg_values, init_timestamp=init_time, agent_id=check_call_stack_for_agent_id(), - name=tool_name, + name=name, ) try: diff --git a/tests/test_record_action.py b/tests/test_record_action.py index bc433c55..058e13dd 100644 --- a/tests/test_record_action.py +++ b/tests/test_record_action.py @@ -67,6 +67,27 @@ def add_two(x, y): agentops.end_session(end_state="Success") + def test_record_action_default_name(self, mock_req): + agentops.start_session() + + @record_action() + def add_two(x, y): + return x + y + + # Act + add_two(3, 4) + time.sleep(0.1) + + # Assert + assert len(mock_req.request_history) == 2 + assert mock_req.last_request.headers["X-Agentops-Api-Key"] == self.api_key + request_json = mock_req.last_request.json() + assert request_json["events"][0]["action_type"] == "add_two" + assert request_json["events"][0]["params"] == {"x": 3, "y": 4} + assert request_json["events"][0]["returns"] == 7 + + agentops.end_session(end_state="Success") + def test_record_action_decorator_multiple(self, mock_req): agentops.start_session() diff --git a/tests/test_record_tool.py b/tests/test_record_tool.py index b290e131..29a5d473 100644 --- a/tests/test_record_tool.py +++ b/tests/test_record_tool.py @@ -68,6 +68,27 @@ def add_two(x, y): agentops.end_session(end_state="Success") + def test_record_tool_default_name(self, mock_req): + agentops.start_session() + + @record_tool() + def add_two(x, y): + return x + y + + # Act + add_two(3, 4) + time.sleep(0.1) + + # Assert + assert len(mock_req.request_history) == 2 + assert mock_req.last_request.headers["X-Agentops-Api-Key"] == self.api_key + request_json = mock_req.last_request.json() + assert request_json["events"][0]["name"] == "add_two" + assert request_json["events"][0]["params"] == {"x": 3, "y": 4} + assert request_json["events"][0]["returns"] == 7 + + agentops.end_session(end_state="Success") + def test_record_tool_decorator_multiple(self, mock_req): agentops.start_session()