From 1af0228fe59a8e7bb50c23b23f7bd6d2a18bc08c Mon Sep 17 00:00:00 2001 From: Chris Park Date: Fri, 17 Jan 2025 18:23:17 -0500 Subject: [PATCH 1/4] mypy fix --- .../examples/ex.py | 61 ------- .../instrumentation/autogen/__init__.py | 172 +++++++++++------- 2 files changed, 108 insertions(+), 125 deletions(-) delete mode 100644 python/instrumentation/openinference-instrumentation-autogen/examples/ex.py diff --git a/python/instrumentation/openinference-instrumentation-autogen/examples/ex.py b/python/instrumentation/openinference-instrumentation-autogen/examples/ex.py deleted file mode 100644 index 716a8d422..000000000 --- a/python/instrumentation/openinference-instrumentation-autogen/examples/ex.py +++ /dev/null @@ -1,61 +0,0 @@ -import os -from importlib import import_module - -import autogen -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() - - endpoint = "http://127.0.0.1:6006/v1/traces" - exporter = OTLPSpanExporter(endpoint=endpoint) - span_processor = SimpleSpanProcessor(exporter) - trace_provider.add_span_processor(span_processor) - - trace.set_tracer_provider(trace_provider) - - openai_instrumentation = import_module("openinference.instrumentation.openai") - openai_instrumentation.OpenAIInstrumentor().instrument(tracer_provider=trace_provider) - AutogenInstrumentor().instrument() - - config_list = [ - { - "model": "gpt-4o", - "api_key": os.environ["OPENAI_API_KEY"], - } - ] - - assistant = autogen.AssistantAgent( - name="assistant", - llm_config={ - "seed": 42, # seed for caching and reproducibility - "config_list": config_list, - "temperature": 0, # temperature for sampling - }, - ) - - user_proxy = autogen.UserProxyAgent( - name="user_proxy", - human_input_mode="NEVER", - max_consecutive_auto_reply=10, - is_termination_msg=lambda x: x.get("content", "").rstrip().endswith("TERMINATE"), - code_execution_config={ - "work_dir": "coding", - "use_docker": False, # set to True or an image name like "python:3" to use docker - }, - ) - - user_proxy.initiate_chat( - assistant, - message="""What date is today? Compare the year-to-date gain for META and TESLA.""", - ) - - -if __name__ == "__main__": - main() 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 c01b369ef..da03f79f1 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,104 +1,137 @@ import json +from typing import Any, Callable, Dict, Optional, Union from opentelemetry import trace -from opentelemetry.trace import Status, StatusCode +from opentelemetry.trace import Link, Span, SpanContext, Status, StatusCode -from autogen import ConversableAgent +from autogen import ConversableAgent # type: ignore class AutogenInstrumentor: - def __init__(self): + def __init__(self) -> None: self.tracer = trace.get_tracer(__name__) - self._original_generate = None - self._original_initiate_chat = None - self._original_execute_function = None + self._original_generate: Optional[Callable[..., Any]] = None + self._original_initiate_chat: Optional[Callable[..., Any]] = None + self._original_execute_function: Optional[Callable[..., Any]] = None - def _safe_json_dumps(self, obj): + def _safe_json_dumps(self, obj: Any) -> str: try: return json.dumps(obj) except (TypeError, ValueError): return json.dumps(str(obj)) - def instrument(self): + def instrument(self) -> "AutogenInstrumentor": + # Save original methods 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): + def wrapped_generate( + agent_self: ConversableAgent, + messages: Optional[Any] = None, + sender: Optional[str] = None, + **kwargs: Any, + ) -> Any: + span: Optional[Span] = None try: - current_context = trace.get_current_span().get_span_context() + current_span = trace.get_current_span() + current_context: SpanContext = current_span.get_span_context() with instrumentor.tracer.start_as_current_span( - self.__class__.__name__, - context=trace.set_span_in_context(trace.get_current_span()), - links=[trace.Link(current_context)], + agent_self.__class__.__name__, + context=trace.set_span_in_context(current_span), + links=[Link(current_context)], ) as span: span.set_attribute(SpanAttributes.OPENINFERENCE_SPAN_KIND, "AGENT") span.set_attribute( - SpanAttributes.INPUT_VALUE, instrumentor._safe_json_dumps(messages) + SpanAttributes.INPUT_VALUE, + instrumentor._safe_json_dumps(messages), ) span.set_attribute(SpanAttributes.INPUT_MIME_TYPE, "application/json") - span.set_attribute("agent.type", self.__class__.__name__) + span.set_attribute("agent.type", agent_self.__class__.__name__) - response = instrumentor._original_generate( - self, messages=messages, sender=sender, **kwargs - ) + if instrumentor._original_generate is not None: + response = instrumentor._original_generate( + agent_self, messages=messages, sender=sender, **kwargs + ) + else: + # Fallback or raise an error if needed + response = None span.set_attribute( - SpanAttributes.OUTPUT_VALUE, instrumentor._safe_json_dumps(response) + SpanAttributes.OUTPUT_VALUE, + instrumentor._safe_json_dumps(response), ) span.set_attribute(SpanAttributes.OUTPUT_MIME_TYPE, "application/json") return response - except Exception as e: - if "span" in locals(): + if span is not None: span.set_status(Status(StatusCode.ERROR)) span.record_exception(e) raise - def wrapped_initiate_chat(self, recipient, *args, **kwargs): + def wrapped_initiate_chat( + agent_self: ConversableAgent, recipient: Any, *args: Any, **kwargs: Any + ) -> Any: + span: Optional[Span] = None try: message = kwargs.get("message", args[0] if args else None) - current_context = trace.get_current_span().get_span_context() + current_span = trace.get_current_span() + current_context: SpanContext = current_span.get_span_context() with instrumentor.tracer.start_as_current_span( "Autogen", - context=trace.set_span_in_context(trace.get_current_span()), - links=[trace.Link(current_context)], + context=trace.set_span_in_context(current_span), + links=[Link(current_context)], ) as span: span.set_attribute(SpanAttributes.OPENINFERENCE_SPAN_KIND, "AGENT") span.set_attribute( - SpanAttributes.INPUT_VALUE, instrumentor._safe_json_dumps(message) + SpanAttributes.INPUT_VALUE, + instrumentor._safe_json_dumps(message), ) span.set_attribute(SpanAttributes.INPUT_MIME_TYPE, "application/json") - result = instrumentor._original_initiate_chat(self, recipient, *args, **kwargs) + if instrumentor._original_initiate_chat is not None: + result = instrumentor._original_initiate_chat( + agent_self, recipient, *args, **kwargs + ) + else: + result = None if hasattr(result, "chat_history") and result.chat_history: last_message = result.chat_history[-1]["content"] span.set_attribute( - SpanAttributes.OUTPUT_VALUE, instrumentor._safe_json_dumps(last_message) + SpanAttributes.OUTPUT_VALUE, + instrumentor._safe_json_dumps(last_message), ) else: span.set_attribute( - SpanAttributes.OUTPUT_VALUE, instrumentor._safe_json_dumps(result) + 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(): + if span is not None: span.set_status(Status(StatusCode.ERROR)) span.record_exception(e) raise - def wrapped_execute_function(self, func_call, call_id=None, verbose=False): + def wrapped_execute_function( + agent_self: ConversableAgent, + func_call: Union[str, Dict[str, Any]], + call_id: Optional[str] = None, + verbose: bool = False, + ) -> Any: + span: Optional[Span] = None try: - current_context = trace.get_current_span().get_span_context() + current_span = trace.get_current_span() + current_context: SpanContext = current_span.get_span_context() # Handle both dictionary and string inputs if isinstance(func_call, str): @@ -109,26 +142,29 @@ def wrapped_execute_function(self, func_call, call_id=None, verbose=False): 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)], + context=trace.set_span_in_context(current_span), + links=[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) + 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__), - ) + # If the agent stores a function map, you can store annotations + if hasattr(agent_self, "_function_map"): + function_map = getattr(agent_self, "_function_map", {}) + if function_name in function_map: + func = 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): @@ -143,34 +179,42 @@ def wrapped_execute_function(self, func_call, call_id=None, verbose=False): 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 - ) + if instrumentor._original_execute_function is not None: + result = instrumentor._original_execute_function( + agent_self, func_call, call_id=call_id, verbose=verbose + ) + else: + result = None # Record output span.set_attribute( - SpanAttributes.OUTPUT_VALUE, instrumentor._safe_json_dumps(result) + 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(): + if span is not None: span.set_status(Status(StatusCode.ERROR)) span.record_exception(e) raise - # Replace methods + # Replace methods on ConversableAgent with wrapped versions ConversableAgent.generate_reply = wrapped_generate ConversableAgent.initiate_chat = wrapped_initiate_chat ConversableAgent.execute_function = wrapped_execute_function return self - def uninstrument(self): - """Restore original behavior""" - if self._original_generate and self._original_initiate_chat: + def uninstrument(self) -> "AutogenInstrumentor": + """Restore original behavior.""" + if ( + self._original_generate + and self._original_initiate_chat + and self._original_execute_function + ): ConversableAgent.generate_reply = self._original_generate ConversableAgent.initiate_chat = self._original_initiate_chat ConversableAgent.execute_function = self._original_execute_function @@ -181,14 +225,14 @@ def uninstrument(self): class SpanAttributes: - OPENINFERENCE_SPAN_KIND = "openinference.span.kind" - INPUT_VALUE = "input.value" - INPUT_MIME_TYPE = "input.mime_type" - OUTPUT_VALUE = "output.value" - OUTPUT_MIME_TYPE = "output.mime_type" - TOOL_NAME = "tool.name" - TOOL_ARGS = "tool.args" - TOOL_KWARGS = "tool.kwargs" - TOOL_PARAMETERS = "tool.parameters" - TOOL_CALL_FUNCTION_ARGUMENTS = "tool_call.function.arguments" - TOOL_CALL_FUNCTION_NAME = "tool_call.function.name" + OPENINFERENCE_SPAN_KIND: str = "openinference.span.kind" + INPUT_VALUE: str = "input.value" + INPUT_MIME_TYPE: str = "input.mime_type" + OUTPUT_VALUE: str = "output.value" + OUTPUT_MIME_TYPE: str = "output.mime_type" + TOOL_NAME: str = "tool.name" + TOOL_ARGS: str = "tool.args" + TOOL_KWARGS: str = "tool.kwargs" + TOOL_PARAMETERS: str = "tool.parameters" + TOOL_CALL_FUNCTION_ARGUMENTS: str = "tool_call.function.arguments" + TOOL_CALL_FUNCTION_NAME: str = "tool_call.function.name" From d42fadf5dbbae10752cdd4f81725320bb7c5f157 Mon Sep 17 00:00:00 2001 From: Chris Park Date: Fri, 17 Jan 2025 18:32:55 -0500 Subject: [PATCH 2/4] typing fix --- .../src/openinference/instrumentation/autogen/__init__.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) 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 da03f79f1..e224d008e 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 @@ -2,7 +2,7 @@ from typing import Any, Callable, Dict, Optional, Union from opentelemetry import trace -from opentelemetry.trace import Link, Span, SpanContext, Status, StatusCode +from opentelemetry.trace import Link, SpanContext, Status, StatusCode from autogen import ConversableAgent # type: ignore @@ -34,7 +34,6 @@ def wrapped_generate( sender: Optional[str] = None, **kwargs: Any, ) -> Any: - span: Optional[Span] = None try: current_span = trace.get_current_span() current_context: SpanContext = current_span.get_span_context() @@ -76,7 +75,6 @@ def wrapped_generate( def wrapped_initiate_chat( agent_self: ConversableAgent, recipient: Any, *args: Any, **kwargs: Any ) -> Any: - span: Optional[Span] = None try: message = kwargs.get("message", args[0] if args else None) current_span = trace.get_current_span() @@ -128,7 +126,6 @@ def wrapped_execute_function( call_id: Optional[str] = None, verbose: bool = False, ) -> Any: - span: Optional[Span] = None try: current_span = trace.get_current_span() current_context: SpanContext = current_span.get_span_context() From 25e63b6c71762561afe3c1be6b89aa3399e52268 Mon Sep 17 00:00:00 2001 From: Chris Park Date: Fri, 17 Jan 2025 18:38:34 -0500 Subject: [PATCH 3/4] test --- .../tests/test_instrumentor.py | 2 ++ 1 file changed, 2 insertions(+) create mode 100644 python/instrumentation/openinference-instrumentation-autogen/tests/test_instrumentor.py diff --git a/python/instrumentation/openinference-instrumentation-autogen/tests/test_instrumentor.py b/python/instrumentation/openinference-instrumentation-autogen/tests/test_instrumentor.py new file mode 100644 index 000000000..769d4c253 --- /dev/null +++ b/python/instrumentation/openinference-instrumentation-autogen/tests/test_instrumentor.py @@ -0,0 +1,2 @@ +def test(): + return From e0904ccfdebe06bb584dd7c4a25ddd60c84a3156 Mon Sep 17 00:00:00 2001 From: Chris Park Date: Fri, 17 Jan 2025 18:40:53 -0500 Subject: [PATCH 4/4] test mypy --- .../tests/test_instrumentor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/instrumentation/openinference-instrumentation-autogen/tests/test_instrumentor.py b/python/instrumentation/openinference-instrumentation-autogen/tests/test_instrumentor.py index 769d4c253..884bfc7b9 100644 --- a/python/instrumentation/openinference-instrumentation-autogen/tests/test_instrumentor.py +++ b/python/instrumentation/openinference-instrumentation-autogen/tests/test_instrumentor.py @@ -1,2 +1,2 @@ -def test(): +def test() -> None: return