From 2a638f7ca51a2d77b27a556bed75aa6318aa805b Mon Sep 17 00:00:00 2001 From: Chris Park Date: Fri, 17 Jan 2025 14:08:55 -0500 Subject: [PATCH] feat(autogen2): basic tool calling (#1216) --- .../examples/ex.py | 3 +- .../instrumentation/autogen/__init__.py | 77 ++++++++++++++++++- 2 files changed, 75 insertions(+), 5 deletions(-) diff --git a/python/instrumentation/openinference-instrumentation-autogen/examples/ex.py b/python/instrumentation/openinference-instrumentation-autogen/examples/ex.py index dfef678f3..716a8d422 100644 --- a/python/instrumentation/openinference-instrumentation-autogen/examples/ex.py +++ b/python/instrumentation/openinference-instrumentation-autogen/examples/ex.py @@ -2,12 +2,13 @@ from importlib import import_module import autogen -from openinference.instrumentation.autogen import AutogenInstrumentor from opentelemetry import trace from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from openinference.instrumentation.autogen import AutogenInstrumentor + def main(): trace_provider = TracerProvider() diff --git a/python/instrumentation/openinference-instrumentation-autogen/src/openinference/instrumentation/autogen/__init__.py b/python/instrumentation/openinference-instrumentation-autogen/src/openinference/instrumentation/autogen/__init__.py index be2a0f6e4..04c0f0b39 100644 --- a/python/instrumentation/openinference-instrumentation-autogen/src/openinference/instrumentation/autogen/__init__.py +++ b/python/instrumentation/openinference-instrumentation-autogen/src/openinference/instrumentation/autogen/__init__.py @@ -1,10 +1,9 @@ import json -import autogen -from autogen import ConversableAgent from opentelemetry import trace -from opentelemetry.sdk.trace import TracerProvider -from opentelemetry.trace import Context, Status, StatusCode +from opentelemetry.trace import Status, StatusCode + +from autogen import ConversableAgent class AutogenInstrumentor: @@ -12,6 +11,7 @@ def __init__(self): self.tracer = trace.get_tracer(__name__) self._original_generate = None self._original_initiate_chat = None + self._original_execute_function = None def _safe_json_dumps(self, obj): try: @@ -22,6 +22,7 @@ def _safe_json_dumps(self, obj): def instrument(self): self._original_generate = ConversableAgent.generate_reply self._original_initiate_chat = ConversableAgent.initiate_chat + self._original_execute_function = ConversableAgent.execute_function instrumentor = self def wrapped_generate(self, messages=None, sender=None, **kwargs): @@ -95,9 +96,75 @@ def wrapped_initiate_chat(self, recipient, *args, **kwargs): span.record_exception(e) raise + def wrapped_execute_function(self, func_call, call_id=None, verbose=False): + try: + current_context = trace.get_current_span().get_span_context() + + # Handle both dictionary and string inputs + if isinstance(func_call, str): + function_name = func_call + func_call = {"name": function_name} + else: + function_name = func_call.get("name", "unknown") + + with instrumentor.tracer.start_as_current_span( + f"{function_name}", + context=trace.set_span_in_context(trace.get_current_span()), + links=[trace.Link(current_context)], + ) as span: + span.set_attribute(SpanAttributes.OPENINFERENCE_SPAN_KIND, "TOOL") + span.set_attribute(SpanAttributes.TOOL_NAME, function_name) + + # Record input + span.set_attribute( + SpanAttributes.INPUT_VALUE, instrumentor._safe_json_dumps(func_call) + ) + span.set_attribute(SpanAttributes.INPUT_MIME_TYPE, "application/json") + + # Record tool-specific attributes + if hasattr(self, "_function_map") and function_name in self._function_map: + func = self._function_map[function_name] + if hasattr(func, "__annotations__"): + span.set_attribute( + SpanAttributes.TOOL_PARAMETERS, + instrumentor._safe_json_dumps(func.__annotations__), + ) + + # Record function call details + if isinstance(func_call, dict): + # Record function arguments + if "arguments" in func_call: + span.set_attribute( + SpanAttributes.TOOL_CALL_FUNCTION_ARGUMENTS, + instrumentor._safe_json_dumps(func_call["arguments"]), + ) + + # Record function name + span.set_attribute(SpanAttributes.TOOL_CALL_FUNCTION_NAME, function_name) + + # Execute function + result = instrumentor._original_execute_function( + self, func_call, call_id=call_id, verbose=verbose + ) + + # Record output + span.set_attribute( + SpanAttributes.OUTPUT_VALUE, instrumentor._safe_json_dumps(result) + ) + span.set_attribute(SpanAttributes.OUTPUT_MIME_TYPE, "application/json") + + return result + + except Exception as e: + if "span" in locals(): + span.set_status(Status(StatusCode.ERROR)) + span.record_exception(e) + raise + # Replace methods ConversableAgent.generate_reply = wrapped_generate ConversableAgent.initiate_chat = wrapped_initiate_chat + ConversableAgent.execute_function = wrapped_execute_function return self @@ -106,8 +173,10 @@ def uninstrument(self): if self._original_generate and self._original_initiate_chat: ConversableAgent.generate_reply = self._original_generate ConversableAgent.initiate_chat = self._original_initiate_chat + ConversableAgent.execute_function = self._original_execute_function self._original_generate = None self._original_initiate_chat = None + self._original_execute_function = None return self