openlit 1.33.10__py3-none-any.whl → 1.33.12__py3-none-any.whl

openlit/instrumentation/ag2/ag2.py

@@ -1,98 +1,163 @@
-# pylint: disable=duplicate-code, broad-exception-caught, too-many-statements, unused-argument
 """
-Module for monitoring AG2.
+Module for monitoring AG2 API calls.
 """
 
 import logging
+import time
 from opentelemetry.trace import SpanKind, Status, StatusCode
 from opentelemetry.sdk.resources import SERVICE_NAME, TELEMETRY_SDK_NAME, DEPLOYMENT_ENVIRONMENT
-from openlit.__helpers import handle_exception
+from openlit.__helpers import (
+    handle_exception,
+    get_chat_model_cost,
+    otel_event,
+)
 from openlit.semcov import SemanticConvetion
 
 # Initialize logger for logging potential issues and operations
 logger = logging.getLogger(__name__)
 
-def wrap_ag2(gen_ai_endpoint, version, environment, application_name,
-                 tracer, pricing_info, trace_content, metrics, disable_metrics):
+AGENT_NAME = ''
+REQUEST_MODEL = ''
+SYSTEM_MESSAGE = ''
+MODEL_AND_NAME_SET = False
+
+def set_span_attributes(span, version, operation_name, environment,
+        application_name, server_address, server_port, request_model):
+    """
+    Set common attributes for the span.
+    """
+
+    # Set Span attributes (OTel Semconv)
+    span.set_attribute(TELEMETRY_SDK_NAME, 'openlit')
+    span.set_attribute(SemanticConvetion.GEN_AI_OPERATION, operation_name)
+    span.set_attribute(SemanticConvetion.GEN_AI_SYSTEM, SemanticConvetion.GEN_AI_SYSTEM_AG2)
+    span.set_attribute(SemanticConvetion.GEN_AI_AGENT_NAME, AGENT_NAME)
+    span.set_attribute(SemanticConvetion.SERVER_ADDRESS, server_address)
+    span.set_attribute(SemanticConvetion.SERVER_PORT, server_port)
+    span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_MODEL, request_model)
+
+    # Set Span attributes (Extras)
+    span.set_attribute(DEPLOYMENT_ENVIRONMENT, environment)
+    span.set_attribute(SERVICE_NAME, application_name)
+    span.set_attribute(SemanticConvetion.GEN_AI_SDK_VERSION, version)
+
+def calculate_tokens_and_cost(response, request_model, pricing_info):
     """
-    Creates a wrapper around a function call to trace and log its execution metrics.
-
-    This function wraps any given function to measure its execution time,
-    log its operation, and trace its execution using OpenTelemetry.
-    
-    Parameters:
-    - gen_ai_endpoint (str): A descriptor or name for the endpoint being traced.
-    - version (str): The version of the Langchain application.
-    - environment (str): The deployment environment (e.g., 'production', 'development').
-    - application_name (str): Name of the Langchain application.
-    - tracer (opentelemetry.trace.Tracer): The tracer object used for OpenTelemetry tracing.
-    - pricing_info (dict): Information about the pricing for internal metrics (currently not used).
-    - trace_content (bool): Flag indicating whether to trace the content of the response.
-
-    Returns:
-    - function: A higher-order function that takes a function 'wrapped' and returns
-                a new function that wraps 'wrapped' with additional tracing and logging.
+    Calculate the input, output tokens, and their respective costs.
     """
+    input_tokens = 0
+    output_tokens = 0
 
+    for usage_data in response.cost.values():
+        if isinstance(usage_data, dict):
+            for model_data in usage_data.values():
+                if isinstance(model_data, dict):
+                    input_tokens += model_data.get('prompt_tokens', 0)
+                    output_tokens += model_data.get('completion_tokens', 0)
+
+    cost = get_chat_model_cost(request_model, pricing_info, input_tokens, output_tokens)
+    return input_tokens, output_tokens, cost
+
+def emit_events(response, event_provider, capture_message_content):
+    """
+    Emit OpenTelemetry events for each chat history entry.
+    """
+    for chat in response.chat_history:
+        event_type = (
+            SemanticConvetion.GEN_AI_CHOICE if chat['role'] == 'user'
+            else SemanticConvetion.GEN_AI_USER_MESSAGE
+        )
+        choice_event = otel_event(
+            name=event_type,
+            attributes={
+                SemanticConvetion.GEN_AI_SYSTEM: SemanticConvetion.GEN_AI_SYSTEM_AG2
+            },
+            body={
+                'index': response.chat_history.index(chat),
+                'message': {
+                    **({'content': chat['content']} if capture_message_content else {}),
+                    'role': 'assistant' if chat['role'] == 'user' else 'user'
+                }
+            }
+        )
+        event_provider.emit(choice_event)
+
+def conversable_agent(version, environment, application_name,
+                      tracer, event_provider, pricing_info, capture_message_content, metrics, disable_metrics):
+    """
+    Generates a telemetry wrapper for GenAI function call
+    """
     def wrapper(wrapped, instance, args, kwargs):
-        """
-        An inner wrapper function that executes the wrapped function, measures execution
-        time, and records trace data using OpenTelemetry.
-
-        Parameters:
-        - wrapped (Callable): The original function that this wrapper will execute.
-        - instance (object): The instance to which the wrapped function belongs. This
-                             is used for instance methods. For static and classmethods,
-                             this may be None.
-        - args (tuple): Positional arguments passed to the wrapped function.
-        - kwargs (dict): Keyword arguments passed to the wrapped function.
-
-        Returns:
-        - The result of the wrapped function call.
-        
-        The wrapper initiates a span with the provided tracer, sets various attributes
-        on the span based on the function's execution and response, and ensures
-        errors are handled and logged appropriately.
-        """
-
-        with tracer.start_as_current_span(gen_ai_endpoint, kind= SpanKind.CLIENT) as span:
-            response = wrapped(*args, **kwargs)
-
-            if isinstance(instance.__dict__.get('llm_config'), dict):
-                llm_model = instance.__dict__['llm_config'].get('model', 'gpt-4')
-            else:
-                # Fallback to default if 'llm_config' is not a dictionary
-                llm_model = None
+        server_address, server_port = '127.0.0.1', 80
+        global AGENT_NAME, MODEL_AND_NAME_SET, REQUEST_MODEL, SYSTEM_MESSAGE
+
+        if not MODEL_AND_NAME_SET:
+            AGENT_NAME = kwargs.get('name', 'NOT_FOUND')
+            REQUEST_MODEL = kwargs.get('llm_config', {}).get('model', 'gpt-4o')
+            SYSTEM_MESSAGE = kwargs.get('system_message', '')
+            MODEL_AND_NAME_SET = True
 
+        span_name = f'{SemanticConvetion.GEN_AI_OPERATION_TYPE_CREATE_AGENT} {AGENT_NAME}'
+
+        with tracer.start_as_current_span(span_name, kind=SpanKind.CLIENT) as span:
             try:
-                span.set_attribute(TELEMETRY_SDK_NAME, "openlit")
-                span.set_attribute(SemanticConvetion.GEN_AI_ENDPOINT,
-                                   gen_ai_endpoint)
-                span.set_attribute(SemanticConvetion.GEN_AI_SYSTEM,
-                                   SemanticConvetion.GEN_AI_SYSTEM_AG2)
-                span.set_attribute(DEPLOYMENT_ENVIRONMENT,
-                                   environment)
-                span.set_attribute(SERVICE_NAME,
-                                   application_name)
-                span.set_attribute(SemanticConvetion.GEN_AI_OPERATION,
-                                   SemanticConvetion.GEN_AI_OPERATION_TYPE_AGENT)
-                span.set_attribute(SemanticConvetion.GEN_AI_AGENT_ROLE,
-                                   instance.name)
-                if llm_model:
-                    span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_MODEL,
-                                   llm_model)
+                start_time = time.time()
+                response = wrapped(*args, **kwargs)
+                end_time = time.time()
 
+                set_span_attributes(span, version, SemanticConvetion.GEN_AI_OPERATION_TYPE_CREATE_AGENT,
+                    environment, application_name, server_address, server_port, REQUEST_MODEL)
+                span.set_attribute(SemanticConvetion.GEN_AI_AGENT_DESCRIPTION, SYSTEM_MESSAGE)
+                span.set_attribute(SemanticConvetion.GEN_AI_RESPONSE_MODEL, REQUEST_MODEL)
+                span.set_attribute(SemanticConvetion.GEN_AI_SERVER_TTFT, end_time - start_time)
 
                 span.set_status(Status(StatusCode.OK))
 
-                # Return original response
                 return response
 
             except Exception as e:
                 handle_exception(span, e)
-                logger.error("Error in trace creation: %s", e)
+                logger.error('Error in trace creation: %s', e)
+                return response
 
-                # Return original response
+    return wrapper
+
+def agent_run(version, environment, application_name,
+              tracer, event_provider, pricing_info, capture_message_content, metrics, disable_metrics):
+    """
+    Generates a telemetry wrapper for GenAI function call
+    """
+    def wrapper(wrapped, instance, args, kwargs):
+        server_address, server_port = '127.0.0.1', 80
+
+        span_name = f'{SemanticConvetion.GEN_AI_OPERATION_TYPE_EXECUTE_AGENT_TASK} {AGENT_NAME}'
+
+        with tracer.start_as_current_span(span_name, kind=SpanKind.CLIENT) as span:
+            try:
+                start_time = time.time()
+                response = wrapped(*args, **kwargs)
+                end_time = time.time()
+
+                input_tokens, output_tokens, cost = calculate_tokens_and_cost(response, REQUEST_MODEL, pricing_info)
+                response_model = list(response.cost.get('usage_including_cached_inference', {}).keys())[1]
+
+                set_span_attributes(span, version, SemanticConvetion.GEN_AI_OPERATION_TYPE_EXECUTE_AGENT_TASK,
+                    environment, application_name, server_address, server_port, REQUEST_MODEL)
+                span.set_attribute(SemanticConvetion.GEN_AI_RESPONSE_MODEL, response_model)
+                span.set_attribute(SemanticConvetion.GEN_AI_USAGE_INPUT_TOKENS, input_tokens)
+                span.set_attribute(SemanticConvetion.GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens)
+                span.set_attribute(SemanticConvetion.GEN_AI_CLIENT_TOKEN_USAGE, input_tokens + output_tokens)
+                span.set_attribute(SemanticConvetion.GEN_AI_USAGE_COST, cost)
+                span.set_attribute(SemanticConvetion.GEN_AI_SERVER_TTFT, end_time - start_time)
+
+                emit_events(response, event_provider, capture_message_content)
+                span.set_status(Status(StatusCode.OK))
+
+                return response
+
+            except Exception as e:
+                handle_exception(span, e)
+                logger.error('Error in trace creation: %s', e)
                 return response
 
     return wrapper

openlit/instrumentation/ai21/__init__.py

@@ -13,7 +13,7 @@ from openlit.instrumentation.ai21.async_ai21 import (
     async_chat, async_chat_rag
 )
 
-_instruments = ("ai21 >= 3.0.0",)
+_instruments = ('ai21 >= 3.0.0',)
 
 class AI21Instrumentor(BaseInstrumentor):
     """
@@ -24,41 +24,42 @@ class AI21Instrumentor(BaseInstrumentor):
         return _instruments
 
     def _instrument(self, **kwargs):
-        application_name = kwargs.get("application_name", "default_application")
-        environment = kwargs.get("environment", "default_environment")
-        tracer = kwargs.get("tracer")
-        metrics = kwargs.get("metrics_dict")
-        pricing_info = kwargs.get("pricing_info", {})
-        trace_content = kwargs.get("trace_content", False)
-        disable_metrics = kwargs.get("disable_metrics")
-        version = importlib.metadata.version("ai21")
+        application_name = kwargs.get('application_name', 'default')
+        environment = kwargs.get('environment', 'default')
+        tracer = kwargs.get('tracer')
+        event_provider = kwargs.get('event_provider')
+        metrics = kwargs.get('metrics_dict')
+        pricing_info = kwargs.get('pricing_info', {})
+        capture_message_content = kwargs.get('capture_message_content', False)
+        disable_metrics = kwargs.get('disable_metrics')
+        version = importlib.metadata.version('ai21')
 
         #sync
         wrap_function_wrapper(
-            "ai21.clients.studio.resources.chat.chat_completions",
-            "ChatCompletions.create",
+            'ai21.clients.studio.resources.chat.chat_completions',
+            'ChatCompletions.create',
             chat(version, environment, application_name,
-                  tracer, pricing_info, trace_content, metrics, disable_metrics),
+                  tracer, event_provider, pricing_info, capture_message_content, metrics, disable_metrics),
         )
         wrap_function_wrapper(
-            "ai21.clients.studio.resources.studio_conversational_rag",
-            "StudioConversationalRag.create",
+            'ai21.clients.studio.resources.studio_conversational_rag',
+            'StudioConversationalRag.create',
             chat_rag(version, environment, application_name,
-                  tracer, pricing_info, trace_content, metrics, disable_metrics),
+                  tracer, event_provider, pricing_info, capture_message_content, metrics, disable_metrics),
         )
 
         #Async
         wrap_function_wrapper(
-            "ai21.clients.studio.resources.chat.async_chat_completions",
-            "AsyncChatCompletions.create",
+            'ai21.clients.studio.resources.chat.async_chat_completions',
+            'AsyncChatCompletions.create',
             async_chat(version, environment, application_name,
-                  tracer, pricing_info, trace_content, metrics, disable_metrics),
+                  tracer, event_provider, pricing_info, capture_message_content, metrics, disable_metrics),
         )
         wrap_function_wrapper(
-            "ai21.clients.studio.resources.studio_conversational_rag",
-            "AsyncStudioConversationalRag.create",
+            'ai21.clients.studio.resources.studio_conversational_rag',
+            'AsyncStudioConversationalRag.create',
             async_chat_rag(version, environment, application_name,
-                  tracer, pricing_info, trace_content, metrics, disable_metrics),
+                  tracer, event_provider, pricing_info, capture_message_content, metrics, disable_metrics),
         )
 
     def _uninstrument(self, **kwargs):