Add LLMAgent events to autogen instrumentation. (#1455)

Author: umaannamalai

newrelic/hooks/mlmodel_autogen.py

@@ -59,11 +59,39 @@ def wrap_on_messages_stream(wrapped, instance, args, kwargs):
     if not transaction:
         return wrapped(*args, **kwargs)
 
+    settings = transaction.settings or global_settings()
+    if not settings.ai_monitoring.enabled:
+        return wrapped(*args, **kwargs)
+
+    # Framework metric also used for entity tagging in the UI
+    transaction.add_ml_model_info("Autogen", AUTOGEN_VERSION)
+    transaction._add_agent_attribute("llm", True)
+
     agent_name = getattr(instance, "name", "agent")
+    agent_id = str(uuid.uuid4())
+    agent_event_dict = _construct_base_agent_event_dict(agent_name, agent_id, transaction)
     func_name = callable_name(wrapped)
     function_trace_name = f"{func_name}/{agent_name}"
-    with FunctionTrace(name=function_trace_name, group="Llm", source=wrapped):
-        return wrapped(*args, **kwargs)
+
+    ft = FunctionTrace(name=function_trace_name, group="Llm/agent/Autogen")
+    ft.__enter__()
+
+    try:
+        return_val = wrapped(*args, **kwargs)
+    except Exception:
+        ft.notice_error(attributes={"agent_id": agent_id})
+        ft.__exit__(*sys.exc_info())
+        # If we hit an exception, append the error attribute and duration from the exited function trace
+        agent_event_dict.update({"duration": ft.duration * 1000, "error": True})
+        transaction.record_custom_event("LlmAgent", agent_event_dict)
+        raise
+
+    ft.__exit__(None, None, None)
+    agent_event_dict.update({"duration": ft.duration * 1000})
+
+    transaction.record_custom_event("LlmAgent", agent_event_dict)
+
+    return return_val
 
 
 def _get_llm_metadata(transaction):
@@ -115,6 +143,26 @@ def _construct_base_tool_event_dict(bound_args, tool_call_data, tool_id, transac
     return tool_event_dict
 
 
+def _construct_base_agent_event_dict(agent_name, agent_id, transaction):
+    try:
+        linking_metadata = get_trace_linking_metadata()
+
+        agent_event_dict = {
+            "id": agent_id,
+            "name": agent_name,
+            "span_id": linking_metadata.get("span.id"),
+            "trace_id": linking_metadata.get("trace.id"),
+            "vendor": "autogen",
+            "ingest_source": "Python",
+        }
+        agent_event_dict.update(_get_llm_metadata(transaction))
+    except Exception:
+        agent_event_dict = {}
+        _logger.warning(RECORD_EVENTS_FAILURE_LOG_MESSAGE, exc_info=True)
+
+    return agent_event_dict
+
+
 async def wrap__execute_tool_call(wrapped, instance, args, kwargs):
     transaction = current_transaction()
     if not transaction:
@@ -150,7 +198,6 @@ async def wrap__execute_tool_call(wrapped, instance, args, kwargs):
         raise
 
     ft.__exit__(None, None, None)
-
     tool_event_dict.update({"duration": ft.duration * 1000})
 
     # If the tool was executed successfully, we can grab the tool output from the result

tests/mlmodel_autogen/conftest.py

@@ -13,21 +13,14 @@
 # limitations under the License.
 
 import json
-
 import pytest
-from autogen_agentchat.agents import AssistantAgent
-from autogen_agentchat.base import TaskResult
-from autogen_agentchat.teams import RoundRobinGroupChat
-from autogen_core import ComponentModel, FunctionCall, Image
+
+from autogen_core import FunctionCall
 from autogen_core.models import CreateResult, RequestUsage
-from autogen_core.models._model_client import ModelFamily
 from autogen_ext.models.replay import ReplayChatCompletionClient
-from pydantic import BaseModel, ValidationError
 from testing_support.fixture.event_loop import event_loop as loop
 from testing_support.fixtures import collector_agent_registration_fixture, collector_available_fixture
 
-from newrelic.common.object_names import callable_name
-
 _default_settings = {
     "package_reporting.enabled": False,  # Turn off package reporting for testing as it causes slowdowns.
     "transaction_tracer.explain_threshold": 0.0,
@@ -104,13 +97,13 @@ def multi_tool_model_client():
             ),
             CreateResult(
                 finish_reason="function_calls",
-                content=[FunctionCall(id="2", name="add_exclamation", arguments=json.dumps({"message": "Goodbye"}))],
+                content=[FunctionCall(id="3", name="compute_sum", arguments=json.dumps({"a": 5, "b": 3}))],
                 usage=RequestUsage(prompt_tokens=10, completion_tokens=5),
                 cached=False,
             ),
             CreateResult(
                 finish_reason="function_calls",
-                content=[FunctionCall(id="3", name="compute_sum", arguments=json.dumps({"a": 5, "b": 3}))],
+                content=[FunctionCall(id="2", name="add_exclamation", arguments=json.dumps({"message": "Goodbye"}))],
                 usage=RequestUsage(prompt_tokens=10, completion_tokens=5),
                 cached=False,
             ),

tests/mlmodel_autogen/test_assistant_agent.py

@@ -12,9 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import copy
 
 import pytest
+
 from autogen_agentchat.agents import AssistantAgent
 from autogen_agentchat.base import TaskResult
 from testing_support.fixtures import reset_core_stats_engine, validate_attributes
@@ -34,6 +34,10 @@
 from newrelic.api.background_task import background_task
 from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes
 from newrelic.common.object_names import callable_name
+from newrelic.common.package_version_utils import get_package_version_tuple
+
+
+AUTOGEN_VERSION = get_package_version_tuple("autogen-agentchat")
 
 tool_recorded_event = [
     (
@@ -74,25 +78,49 @@
 ]
 
 
+agent_recorded_event = [
+    (
+        {"type": "LlmAgent"},
+        {
+            "id": None,
+            "name": "pirate_agent",
+            "span_id": None,
+            "trace_id": "trace-id",
+            "vendor": "autogen",
+            "ingest_source": "Python",
+            "duration": None,
+        },
+    )
+]
+
+
 # Example tool for testing purposes
 def add_exclamation(message: str) -> str:
     return f"{message}!"
 
 
 @reset_core_stats_engine()
-@validate_custom_events(events_with_context_attrs(tool_recorded_event))
-@validate_custom_event_count(count=1)
+@validate_custom_events(
+    events_with_context_attrs(tool_recorded_event) + events_with_context_attrs(agent_recorded_event)
+)
+@validate_custom_event_count(count=2)
 @validate_transaction_metrics(
     "test_assistant_agent:test_run_assistant_agent",
     scoped_metrics=[
-        ("Llm/autogen_agentchat.agents._assistant_agent:AssistantAgent.on_messages_stream/pirate_agent", 1),
+        (
+            "Llm/agent/Autogen/autogen_agentchat.agents._assistant_agent:AssistantAgent.on_messages_stream/pirate_agent",
+            1,
+        ),
         (
             "Llm/tool/Autogen/autogen_agentchat.agents._assistant_agent:AssistantAgent._execute_tool_call/add_exclamation",
             1,
         ),
     ],
     rollup_metrics=[
-        ("Llm/autogen_agentchat.agents._assistant_agent:AssistantAgent.on_messages_stream/pirate_agent", 1),
+        (
+            "Llm/agent/Autogen/autogen_agentchat.agents._assistant_agent:AssistantAgent.on_messages_stream/pirate_agent",
+            1,
+        ),
         (
             "Llm/tool/Autogen/autogen_agentchat.agents._assistant_agent:AssistantAgent._execute_tool_call/add_exclamation",
             1,
@@ -117,19 +145,25 @@ async def _test():
 
 
 @reset_core_stats_engine()
-@validate_custom_events(tool_recorded_event)
-@validate_custom_event_count(count=1)
+@validate_custom_events(tool_recorded_event + agent_recorded_event)
+@validate_custom_event_count(count=2)
 @validate_transaction_metrics(
     "test_assistant_agent:test_run_stream_assistant_agent",
     scoped_metrics=[
-        ("Llm/autogen_agentchat.agents._assistant_agent:AssistantAgent.on_messages_stream/pirate_agent", 1),
+        (
+            "Llm/agent/Autogen/autogen_agentchat.agents._assistant_agent:AssistantAgent.on_messages_stream/pirate_agent",
+            1,
+        ),
         (
             "Llm/tool/Autogen/autogen_agentchat.agents._assistant_agent:AssistantAgent._execute_tool_call/add_exclamation",
             1,
         ),
     ],
     rollup_metrics=[
-        ("Llm/autogen_agentchat.agents._assistant_agent:AssistantAgent.on_messages_stream/pirate_agent", 1),
+        (
+            "Llm/agent/Autogen/autogen_agentchat.agents._assistant_agent:AssistantAgent.on_messages_stream/pirate_agent",
+            1,
+        ),
         (
             "Llm/tool/Autogen/autogen_agentchat.agents._assistant_agent:AssistantAgent._execute_tool_call/add_exclamation",
             1,
@@ -162,19 +196,25 @@ async def _test():
 
 @reset_core_stats_engine()
 @disabled_ai_monitoring_record_content_settings
-@validate_custom_events(tool_events_sans_content(tool_recorded_event))
-@validate_custom_event_count(count=1)
+@validate_custom_events(tool_events_sans_content(tool_recorded_event) + agent_recorded_event)
+@validate_custom_event_count(count=2)
 @validate_transaction_metrics(
     "test_assistant_agent:test_run_assistant_agent_no_content",
     scoped_metrics=[
-        ("Llm/autogen_agentchat.agents._assistant_agent:AssistantAgent.on_messages_stream/pirate_agent", 1),
+        (
+            "Llm/agent/Autogen/autogen_agentchat.agents._assistant_agent:AssistantAgent.on_messages_stream/pirate_agent",
+            1,
+        ),
         (
             "Llm/tool/Autogen/autogen_agentchat.agents._assistant_agent:AssistantAgent._execute_tool_call/add_exclamation",
             1,
         ),
     ],
     rollup_metrics=[
-        ("Llm/autogen_agentchat.agents._assistant_agent:AssistantAgent.on_messages_stream/pirate_agent", 1),
+        (
+            "Llm/agent/Autogen/autogen_agentchat.agents._assistant_agent:AssistantAgent.on_messages_stream/pirate_agent",
+            1,
+        ),
         (
             "Llm/tool/Autogen/autogen_agentchat.agents._assistant_agent:AssistantAgent._execute_tool_call/add_exclamation",
             1,
@@ -214,22 +254,35 @@ async def _test():
     loop.run_until_complete(_test())
 
 
+SKIP_IF_AUTOGEN_062 = pytest.mark.skipif(
+    AUTOGEN_VERSION > (0, 6, 1),
+    reason="Forcing invalid tool call arguments causes a hang on autogen versions above 0.6.1",
+)
+
+
+@SKIP_IF_AUTOGEN_062
 @reset_core_stats_engine()
 @validate_transaction_error_event_count(1)
 @validate_error_trace_attributes(callable_name(TypeError), exact_attrs={"agent": {}, "intrinsic": {}, "user": {}})
 @validate_custom_events(tool_recorded_event_error)
-@validate_custom_event_count(count=1)
+@validate_custom_event_count(count=2)
 @validate_transaction_metrics(
     "test_assistant_agent:test_run_assistant_agent_error",
     scoped_metrics=[
-        ("Llm/autogen_agentchat.agents._assistant_agent:AssistantAgent.on_messages_stream/pirate_agent", 1),
+        (
+            "Llm/agent/Autogen/autogen_agentchat.agents._assistant_agent:AssistantAgent.on_messages_stream/pirate_agent",
+            1,
+        ),
         (
             "Llm/tool/Autogen/autogen_agentchat.agents._assistant_agent:AssistantAgent._execute_tool_call/add_exclamation",
             1,
         ),
     ],
     rollup_metrics=[
-        ("Llm/autogen_agentchat.agents._assistant_agent:AssistantAgent.on_messages_stream/pirate_agent", 1),
+        (
+            "Llm/agent/Autogen/autogen_agentchat.agents._assistant_agent:AssistantAgent.on_messages_stream/pirate_agent",
+            1,
+        ),
         (
             "Llm/tool/Autogen/autogen_agentchat.agents._assistant_agent:AssistantAgent._execute_tool_call/add_exclamation",
             1,
@@ -250,7 +303,7 @@ def test_run_assistant_agent_error(loop, set_trace_info, single_tool_model_clien
 
     async def _test():
         with pytest.raises(TypeError):
-            response = await pirate_agent.run()
+            await pirate_agent.run()
 
     loop.run_until_complete(_test())