Skip to content

Commit

Permalink
Set agent name in constructor (#322)
Browse files Browse the repository at this point in the history
* allow naming in agent constructor

* allow naming in agent constructor

* formatting

* moved to decorators

* doc kwarg update

* remove empty file
  • Loading branch information
bboynton97 authored Jul 31, 2024
1 parent 6dd5a01 commit 60ae045
Show file tree
Hide file tree
Showing 3 changed files with 44 additions and 0 deletions.
5 changes: 5 additions & 0 deletions agentops/decorators.py
Original file line number Diff line number Diff line change
Expand Up @@ -156,6 +156,11 @@ def decorator(obj):

def new_init(self, *args, **kwargs):
try:
kwarg_name = kwargs.get("agentops_name", None)
if kwarg_name is not None:
self.agent_ops_agent_name = kwarg_name
del kwargs["agentops_name"]

original_init(self, *args, **kwargs)

if not Client().is_initialized:
Expand Down
5 changes: 5 additions & 0 deletions docs/v1/usage/tracking-agents.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,11 @@ class MyAgent:
...
```

Alternatively, to name an agent in runtime, the `@track_agent` decorator creates an additional keyword argument in the class constructor. Using the above example, you can create an agent with a dynamic name with:
```python
research_agent = MyAgent(agentops_name='ResearchAgent')
```

`trackagent.name` is optional. <br></br>
If omitted, agent name defaults to the name of the class (e.g. MyAgent). <br></br>
If an event does not originate from a tracked agent, agent name defaults to "Default Agent".
Expand Down
34 changes: 34 additions & 0 deletions tests/test_agent.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
from unittest import TestCase
from unittest.mock import patch, MagicMock
from uuid import uuid4

from agentops import track_agent
import agentops


class TrackAgentTests(TestCase):
def test_track_agent_with_class(self):
agentops.init()

@track_agent(name="agent_name")
class TestAgentClass:
t = "a"
pass

obj = TestAgentClass()
self.assertTrue(isinstance(obj, TestAgentClass))
self.assertEqual(getattr(obj, "agent_ops_agent_name"), "agent_name")
self.assertIsNotNone(getattr(obj, "agent_ops_agent_id"))

def test_track_agent_with_class_name(self):
agentops.init()

@track_agent(name="agent_name")
class TestAgentClass:
t = "a"
pass

obj = TestAgentClass(agentops_name="agent1")
self.assertTrue(isinstance(obj, TestAgentClass))
self.assertEqual(getattr(obj, "agent_ops_agent_name"), "agent1")
self.assertIsNotNone(getattr(obj, "agent_ops_agent_id"))

0 comments on commit 60ae045

Please sign in to comment.