diff --git a/agentops/decorators.py b/agentops/decorators.py index 860b7c73f..62266b9e7 100644 --- a/agentops/decorators.py +++ b/agentops/decorators.py @@ -3,11 +3,12 @@ from typing import Optional, Union from uuid import uuid4 +from .client import Client +from .descriptor import agentops_property from .event import ActionEvent, ErrorEvent, ToolEvent from .helpers import check_call_stack_for_agent_id, get_ISO_time -from .session import Session -from .client import Client from .log_config import logger +from .session import Session def record_function(event_name: str): @@ -303,45 +304,56 @@ def sync_wrapper(*args, session: Optional[Session] = None, **kwargs): def track_agent(name: Union[str, None] = None): def decorator(obj): - if name: - obj.agent_ops_agent_name = name - if inspect.isclass(obj): + # Set up the descriptors on the class + setattr(obj, "agentops_agent_id", agentops_property()) + setattr(obj, "agentops_agent_name", agentops_property()) + original_init = obj.__init__ def new_init(self, *args, **kwargs): + """ + WIthin the __init__ method, we set agentops_ properties via the private, internal descriptor + """ 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"] + # Handle name from kwargs first + name_ = kwargs.pop("agentops_name", None) + # Call original init original_init(self, *args, **kwargs) - self.agent_ops_agent_id = str(uuid4()) + # Set the agent ID + self._agentops_agent_id = str(uuid4()) + + # Force set the private name directly to bypass potential Pydantic interference + if name_ is not None: + setattr(self, "_agentops_agent_name", name_) + elif name is not None: + setattr(self, "_agentops_agent_name", name) + elif hasattr(self, "role"): + setattr(self, "_agentops_agent_name", self.role) session = kwargs.get("session", None) if session is not None: - self.agent_ops_session_id = session.session_id + self._agentops_session_id = session.session_id Client().create_agent( - name=self.agent_ops_agent_name, - agent_id=self.agent_ops_agent_id, + name=self.agentops_agent_name, + agent_id=self.agentops_agent_id, session=session, ) - except AttributeError as e: + + except AttributeError as ex: + logger.debug(ex) Client().add_pre_init_warning(f"Failed to track an agent {name} with the @track_agent decorator.") logger.warning("Failed to track an agent with the @track_agent decorator.") - original_init(self, *args, **kwargs) obj.__init__ = new_init elif inspect.isfunction(obj): - obj.agent_ops_agent_id = str(uuid4()) # type: ignore - Client().create_agent( - name=obj.agent_ops_agent_name, - agent_id=obj.agent_ops_agent_id, # type: ignore - ) + obj.agentops_agent_id = str(uuid4()) + obj.agentops_agent_name = name + Client().create_agent(name=obj.agentops_agent_name, agent_id=obj.agentops_agent_id) else: raise Exception("Invalid input, 'obj' must be a class or a function") diff --git a/agentops/descriptor.py b/agentops/descriptor.py new file mode 100644 index 000000000..020804cbe --- /dev/null +++ b/agentops/descriptor.py @@ -0,0 +1,187 @@ +import inspect +import logging +from typing import Union +from uuid import UUID + + +class agentops_property: + """ + A descriptor that provides a standardized way to handle agent property access and storage. + Properties are automatically stored with an '_agentops_' prefix to avoid naming conflicts. + + The descriptor can be used in two ways: + 1. As a class attribute directly + 2. Added dynamically through a decorator (like @track_agent) + + Attributes: + private_name (str): The internal name used for storing the property value, + prefixed with '_agentops_'. Set either through __init__ or __set_name__. + + Example: + ```python + # Direct usage in a class + class Agent: + name = agentops_property() + id = agentops_property() + + def __init__(self): + self.name = "Agent1" # Stored as '_agentops_name' + self.id = "123" # Stored as '_agentops_id' + + # Usage with decorator + @track_agent() + class Agent: + pass + # agentops_agent_id and agentops_agent_name are added automatically + ``` + + Notes: + - Property names with 'agentops_' prefix are automatically stripped when creating + the internal storage name + - Returns None if the property hasn't been set + - The descriptor will attempt to resolve property names even when added dynamically + """ + + def __init__(self, name=None): + """ + Initialize the descriptor. + + Args: + name (str, optional): The name for the property. Used as fallback when + the descriptor is added dynamically and __set_name__ isn't called. + """ + self.private_name = None + if name: + self.private_name = f"_agentops_{name.replace('agentops_', '')}" + + def __set_name__(self, owner, name): + """ + Called by Python when the descriptor is defined directly in a class. + Sets up the private name used for attribute storage. + + Args: + owner: The class that owns this descriptor + name: The name given to this descriptor in the class + """ + self.private_name = f"_agentops_{name.replace('agentops_', '')}" + + def __get__(self, obj, objtype=None): + """ + Get the property value. + + Args: + obj: The instance to get the property from + objtype: The class of the instance + + Returns: + The property value, or None if not set + The descriptor itself if accessed on the class rather than an instance + + Raises: + AttributeError: If the property name cannot be determined + """ + if obj is None: + return self + + # Handle case where private_name wasn't set by __set_name__ + if self.private_name is None: + # Try to find the name by looking through the class dict + for name, value in type(obj).__dict__.items(): + if value is self: + self.private_name = f"_agentops_{name.replace('agentops_', '')}" + break + if self.private_name is None: + raise AttributeError("Property name could not be determined") + + # First try getting from object's __dict__ (for Pydantic) + if hasattr(obj, "__dict__"): + dict_value = obj.__dict__.get(self.private_name[1:]) + if dict_value is not None: + return dict_value + + # Fall back to our private storage + return getattr(obj, self.private_name, None) + + def __set__(self, obj, value): + """ + Set the property value. + + Args: + obj: The instance to set the property on + value: The value to set + + Raises: + AttributeError: If the property name cannot be determined + """ + if self.private_name is None: + # Same name resolution as in __get__ + for name, val in type(obj).__dict__.items(): + if val is self: + self.private_name = f"_agentops_{name.replace('agentops_', '')}" + break + if self.private_name is None: + raise AttributeError("Property name could not be determined") + + # Set in both object's __dict__ (for Pydantic) and our private storage + if hasattr(obj, "__dict__"): + obj.__dict__[self.private_name[1:]] = value + setattr(obj, self.private_name, value) + + def __delete__(self, obj): + """ + Delete the property value. + + Args: + obj: The instance to delete the property from + + Raises: + AttributeError: If the property name cannot be determined + """ + if self.private_name is None: + raise AttributeError("Property name could not be determined") + try: + delattr(obj, self.private_name) + except AttributeError: + pass + + @staticmethod + def stack_lookup() -> Union[UUID, None]: + """ + Look through the call stack to find an agent ID. + + This method searches the call stack for objects that have agentops_property + descriptors and returns the agent_id if found. + + Returns: + UUID: The agent ID if found in the call stack + None: If no agent ID is found or if "__main__" is encountered + """ + for frame_info in inspect.stack(): + local_vars = frame_info.frame.f_locals + + for var_name, var in local_vars.items(): + # Stop at main + if var == "__main__": + return None + + try: + # Check if object has our AgentOpsDescriptor descriptors + var_type = type(var) + + # Get all class attributes + class_attrs = {name: getattr(var_type, name, None) for name in dir(var_type)} + + agent_id_desc = class_attrs.get("agentops_agent_id") + + if isinstance(agent_id_desc, agentops_property): + agent_id = agent_id_desc.__get__(var, var_type) + + if agent_id: + agent_name_desc = class_attrs.get("agentops_agent_name") + if isinstance(agent_name_desc, agentops_property): + agent_name = agent_name_desc.__get__(var, var_type) + return agent_id + except Exception: + continue + + return None diff --git a/agentops/helpers.py b/agentops/helpers.py index 025302135..ca0c4f0e3 100644 --- a/agentops/helpers.py +++ b/agentops/helpers.py @@ -1,14 +1,16 @@ -from pprint import pformat -from functools import wraps -from datetime import datetime, timezone import inspect -from typing import Union -import requests import json -from importlib.metadata import version, PackageNotFoundError +from datetime import datetime, timezone +from functools import wraps +from importlib.metadata import PackageNotFoundError, version +from pprint import pformat +from typing import Any, Optional, Union +from uuid import UUID +from .descriptor import agentops_property + +import requests from .log_config import logger -from uuid import UUID def get_ISO_time(): @@ -100,20 +102,7 @@ def remove_unwanted_items(value): def check_call_stack_for_agent_id() -> Union[UUID, None]: - for frame_info in inspect.stack(): - # Look through the call stack for the class that called the LLM - local_vars = frame_info.frame.f_locals - for var in local_vars.values(): - # We stop looking up the stack at main because after that we see global variables - if var == "__main__": - return None - if hasattr(var, "agent_ops_agent_id") and getattr(var, "agent_ops_agent_id"): - logger.debug( - "LLM call from agent named: %s", - getattr(var, "agent_ops_agent_name"), - ) - return getattr(var, "agent_ops_agent_id") - return None + return agentops_property.stack_lookup() def get_agentops_version(): diff --git a/tests/test_agent.py b/tests/test_agent.py index 9005dc759..95617ca6a 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -1,5 +1,8 @@ from unittest import TestCase +from uuid import uuid4 + from agentops import track_agent +from agentops.descriptor import agentops_property class TrackAgentTests(TestCase): @@ -11,8 +14,8 @@ class TestAgentClass: 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")) + self.assertEqual(getattr(obj, "agentops_agent_name", None), "agent_name") + self.assertIsNotNone(getattr(obj, "agentops_agent_id", None)) def test_track_agent_with_class_name(self): @track_agent(name="agent_name") @@ -22,5 +25,226 @@ class TestAgentClass: 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")) + self.assertEqual(getattr(obj, "agentops_agent_name"), "agent1") + self.assertIsNotNone(getattr(obj, "agentops_agent_id")) + + def test_track_agent_with_post_init_name_assignment(self): + """Test setting agentops_agent_name after initialization""" + + @track_agent() + class TestAgentClass: + def __init__(self): + self.role = "test_role" + # Simulate post_init behavior like in CrewAI + self.agentops_agent_name = self.role + + obj = TestAgentClass() + self.assertEqual(getattr(obj, "agentops_agent_name"), "test_role") + self.assertIsNotNone(getattr(obj, "agentops_agent_id")) + + def test_track_agent_with_property_override(self): + """Test overriding agentops properties after initialization""" + + @track_agent() + class TestAgentClass: + def __init__(self): + self.role = "initial_role" + self.agentops_agent_name = self.role + + @property + def role(self): + return self._role + + @role.setter + def role(self, value): + self._role = value + # Update agentops_agent_name when role changes + if hasattr(self, "agentops_agent_name"): + self.agentops_agent_name = value + + # Test initial setting + obj = TestAgentClass() + self.assertEqual(getattr(obj, "agentops_agent_name"), "initial_role") + + # Test property update + obj.role = "updated_role" + self.assertEqual(getattr(obj, "agentops_agent_name"), "updated_role") + self.assertIsNotNone(getattr(obj, "agentops_agent_id")) + + def test_track_agent_with_none_values(self): + """Test handling of None values for agentops properties""" + + @track_agent() + class TestAgentClass: + def __init__(self): + self.role = None + self.agentops_agent_name = None + self._model_validate() + + def _model_validate(self): + # Simulate setting name after validation + if self.role is not None: + self.agentops_agent_name = self.role + + # Test initialization with None + obj = TestAgentClass() + self.assertIsNone(getattr(obj, "agentops_agent_name")) + self.assertIsNotNone(getattr(obj, "agentops_agent_id")) # ID should still be set + + # Test updating from None + obj.role = "new_role" + obj._model_validate() + self.assertEqual(getattr(obj, "agentops_agent_name"), "new_role") + + def test_track_agent_with_pydantic_model(self): + """Test setting agentops_agent_name with actual Pydantic BaseModel""" + try: + from pydantic import BaseModel, Field, model_validator + except ImportError: + self.skipTest("Pydantic not installed, skipping Pydantic model test") + + @track_agent() + class TestAgentModel(BaseModel): + role: str = Field(default="test_role") + agentops_agent_name: str | None = None + agentops_agent_id: str | None = None + + @model_validator(mode="after") + def set_agent_name(self): + # Simulate CrewAI's post_init_setup behavior + self.agentops_agent_name = self.role + return self + + # Test basic initialization + obj = TestAgentModel() + self.assertEqual(obj.agentops_agent_name, "test_role") + self.assertIsNotNone(obj.agentops_agent_id) + + # Test with custom role + obj2 = TestAgentModel(role="custom_role") + self.assertEqual(obj2.agentops_agent_name, "custom_role") + self.assertIsNotNone(obj2.agentops_agent_id) + + # Test model update + obj.role = "updated_role" + obj.set_agent_name() + self.assertEqual(obj.agentops_agent_name, "updated_role") + + +class TestAgentOpsDescriptor(TestCase): + def test_agent_property_get_set(self): + """Test basic get/set functionality of agentops_property""" + + class TestAgent: + agent_id = agentops_property() + agent_name = agentops_property() + + agent = TestAgent() + test_id = str(uuid4()) + test_name = "TestAgent" + + # Test setting values + agent.agent_id = test_id + agent.agent_name = test_name + + # Test getting values + self.assertEqual(agent.agent_id, test_id) + self.assertEqual(agent.agent_name, test_name) + + # Test getting non-existent value returns None + self.assertIsNone(TestAgent().agent_id) + + def test_from_stack_direct_call(self): + """Test from_stack when called directly from a method with an agent""" + + @track_agent(name="TestAgent") + class TestAgent: + def get_my_id(self): + return agentops_property.stack_lookup() + + agent = TestAgent() + detected_id = agent.get_my_id() + self.assertEqual(detected_id, agent.agentops_agent_id) + + def test_from_stack_nested_call(self): + """Test from_stack when called through nested function calls""" + + @track_agent(name="TestAgent") + class TestAgent: + def get_my_id(self): + def nested_func(): + return agentops_property.stack_lookup() + + return nested_func() + + agent = TestAgent() + detected_id = agent.get_my_id() + self.assertEqual(detected_id, agent.agentops_agent_id) + + def test_from_stack_multiple_agents(self): + """Test from_stack with multiple agents in different stack frames""" + + @track_agent(name="Agent1") + class Agent1: + def get_other_agent_id(self, other_agent): + return other_agent.get_my_id() + + @track_agent(name="Agent2") + class Agent2: + def get_my_id(self): + return agentops_property.stack_lookup() + + agent1 = Agent1() + agent2 = Agent2() + + # Should return agent2's ID since it's the closest in the call stack + detected_id = agent1.get_other_agent_id(agent2) + self.assertEqual(detected_id, agent2.agentops_agent_id) + self.assertNotEqual(detected_id, agent1.agentops_agent_id) + + def test_from_stack_no_agent(self): + """Test from_stack when no agent is in the call stack""" + + class NonAgent: + def get_id(self): + return agentops_property.stack_lookup() + + non_agent = NonAgent() + self.assertIsNone(non_agent.get_id()) + + def test_from_stack_with_exception(self): + """Test from_stack's behavior when exceptions occur during stack inspection""" + + class ProblemAgent: + agentops_agent_id = agentops_property() + + @property + def problematic_attr(self): + raise Exception("Simulated error") + + def get_id(self): + return agentops_property.stack_lookup() + + agent = ProblemAgent() + # Should return None and not raise exception + self.assertIsNone(agent.get_id()) + + def test_from_stack_inheritance(self): + """Test from_stack with inheritance hierarchy""" + + @track_agent(name="BaseAgent") + class BaseAgent: + def get_id_from_base(self): + return agentops_property.stack_lookup() + + @track_agent(name="DerivedAgent") + class DerivedAgent(BaseAgent): + def get_id_from_derived(self): + return agentops_property.stack_lookup() + + derived = DerivedAgent() + base_call_id = derived.get_id_from_base() + derived_call_id = derived.get_id_from_derived() + + self.assertEqual(base_call_id, derived.agentops_agent_id) + self.assertEqual(derived_call_id, derived.agentops_agent_id)