Skip to content
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.

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

event decorators have default names #367

Merged
merged 3 commits into from
Aug 21, 2024
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
36 changes: 28 additions & 8 deletions agentops/decorators.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,15 +17,15 @@
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:
- Actions: Records function parameters and return statements of the
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):
Expand Down Expand Up @@ -53,11 +53,16 @@
arg_values.update(dict(zip(arg_names, args)))
arg_values.update(kwargs)

if not event_name:
action_type = func.__name__

Check warning on line 57 in agentops/decorators.py

View check run for this annotation

Codecov / codecov/patch

agentops/decorators.py#L57

Added line #L57 was not covered by tests
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:
Expand Down Expand Up @@ -114,11 +119,16 @@
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:
Expand Down Expand Up @@ -153,15 +163,15 @@
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:
- Tools: Records function parameters and return statements of the
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):
Expand Down Expand Up @@ -189,11 +199,16 @@
arg_values.update(dict(zip(arg_names, args)))
arg_values.update(kwargs)

if not tool_name:
name = func.__name__

Check warning on line 203 in agentops/decorators.py

View check run for this annotation

Codecov / codecov/patch

agentops/decorators.py#L203

Added line #L203 was not covered by tests
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:
Expand Down Expand Up @@ -250,11 +265,16 @@
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:
Expand Down
21 changes: 21 additions & 0 deletions tests/test_record_action.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down
21 changes: 21 additions & 0 deletions tests/test_record_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down
Loading