From 60ae0456e7f44806d407d68278f11d64c5d62115 Mon Sep 17 00:00:00 2001
From: Braelyn Boynton <bboynton97@gmail.com>
Date: Wed, 31 Jul 2024 13:31:10 -0700
Subject: [PATCH] Set agent name in constructor (#322)

* allow naming in agent constructor

* allow naming in agent constructor

* formatting

* moved to decorators

* doc kwarg update

* remove empty file
---
 agentops/decorators.py            |  5 +++++
 docs/v1/usage/tracking-agents.mdx |  5 +++++
 tests/test_agent.py               | 34 +++++++++++++++++++++++++++++++
 3 files changed, 44 insertions(+)
 create mode 100644 tests/test_agent.py

diff --git a/agentops/decorators.py b/agentops/decorators.py
index 0ea961e8..ce0d4050 100644
--- a/agentops/decorators.py
+++ b/agentops/decorators.py
@@ -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:
diff --git a/docs/v1/usage/tracking-agents.mdx b/docs/v1/usage/tracking-agents.mdx
index 04c4ec41..a261951d 100644
--- a/docs/v1/usage/tracking-agents.mdx
+++ b/docs/v1/usage/tracking-agents.mdx
@@ -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".
diff --git a/tests/test_agent.py b/tests/test_agent.py
new file mode 100644
index 00000000..abf7b12a
--- /dev/null
+++ b/tests/test_agent.py
@@ -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"))