openlit 1.33.9__py3-none-any.whl → 1.33.11__py3-none-any.whl

openlit/instrumentation/transformers/__init__.py

@@ -21,15 +21,15 @@ class TransformersInstrumentor(BaseInstrumentor):
         tracer = kwargs.get("tracer")
         metrics = kwargs.get("metrics_dict")
         pricing_info = kwargs.get("pricing_info")
-        trace_content = kwargs.get("trace_content")
+        capture_message_content = kwargs.get("capture_message_content")
         disable_metrics = kwargs.get("disable_metrics")
         version = importlib.metadata.version("transformers")
 
         wrap_function_wrapper(
             "transformers",  
             "TextGenerationPipeline.__call__",  
-            text_wrap("huggingface.text_generation", version, environment, application_name,
-                 tracer, pricing_info, trace_content, metrics, disable_metrics),
+            text_wrap(version, environment, application_name,
+                 tracer, pricing_info, capture_message_content, metrics, disable_metrics),
         )
 
     @staticmethod

openlit/instrumentation/transformers/transformers.py

@@ -1,19 +1,25 @@
-# pylint: disable=duplicate-code, broad-exception-caught, too-many-statements, unused-argument
 """
 Module for monitoring ChromaDB.
 """
 
 import logging
+import time
 from opentelemetry.trace import SpanKind, Status, StatusCode
-from opentelemetry.sdk.resources import TELEMETRY_SDK_NAME
-from openlit.__helpers import handle_exception, general_tokens
+from opentelemetry.sdk.resources import SERVICE_NAME, TELEMETRY_SDK_NAME, DEPLOYMENT_ENVIRONMENT
+from openlit.__helpers import (
+    get_chat_model_cost,
+    handle_exception,
+    general_tokens,
+    create_metrics_attributes,
+    set_server_address_and_port
+)
 from openlit.semcov import SemanticConvetion
 
 # Initialize logger for logging potential issues and operations
 logger = logging.getLogger(__name__)
 
-def text_wrap(gen_ai_endpoint, version, environment, application_name,
-                 tracer, pricing_info, trace_content, metrics, disable_metrics):
+def text_wrap(version, environment, application_name,
+                 tracer, pricing_info, capture_message_content, metrics, disable_metrics):
     """
     Creates a wrapper around a function call to trace and log its execution metrics.
 
@@ -21,13 +27,12 @@ def text_wrap(gen_ai_endpoint, version, environment, application_name,
     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.
+    - capture_message_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
@@ -55,8 +60,15 @@ def text_wrap(gen_ai_endpoint, version, environment, application_name,
         errors are handled and logged appropriately.
         """
 
-        with tracer.start_as_current_span(gen_ai_endpoint, kind= SpanKind.CLIENT) as span:
+        server_address, server_port = set_server_address_and_port(instance, "127.0.0.1", 80)
+        request_model = instance.model.config.name_or_path
+
+        span_name = f"{SemanticConvetion.GEN_AI_OPERATION_TYPE_CHAT} {request_model}"
+
+        with tracer.start_as_current_span(span_name, kind= SpanKind.CLIENT) as span:
+            start_time = time.time()
             response = wrapped(*args, **kwargs)
+            end_time = time.time()
 
             # pylint: disable=protected-access
             forward_params = instance._forward_params
@@ -67,45 +79,56 @@ def text_wrap(gen_ai_endpoint, version, environment, application_name,
                 else:
                     prompt = kwargs.get("args", "")
 
-                prompt_tokens = general_tokens(prompt[0])
+                input_tokens = general_tokens(prompt[0])
 
                 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_HUGGING_FACE)
-                span.set_attribute(SemanticConvetion.GEN_AI_ENVIRONMENT,
-                                   environment)
-                span.set_attribute(SemanticConvetion.GEN_AI_APPLICATION_NAME,
-                                   application_name)
                 span.set_attribute(SemanticConvetion.GEN_AI_OPERATION,
                                    SemanticConvetion.GEN_AI_OPERATION_TYPE_CHAT)
+                span.set_attribute(SemanticConvetion.GEN_AI_SYSTEM,
+                                   SemanticConvetion.GEN_AI_SYSTEM_HUGGING_FACE)
                 span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_MODEL,
-                                   instance.model.config.name_or_path)
+                                   request_model)
                 span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_TEMPERATURE,
                                    forward_params.get("temperature", "null"))
                 span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_TOP_P,
                                    forward_params.get("top_p", "null"))
                 span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_MAX_TOKENS,
                                    forward_params.get("max_length", -1))
-                if trace_content:
+                span.set_attribute(SemanticConvetion.GEN_AI_USAGE_INPUT_TOKENS,
+                                   input_tokens)
+                span.set_attribute(SemanticConvetion.SERVER_ADDRESS,
+                                    server_address)
+                span.set_attribute(SemanticConvetion.SERVER_PORT,
+                                    server_port)
+                span.set_attribute(SemanticConvetion.GEN_AI_RESPONSE_MODEL,
+                                    request_model)
+
+                span.set_attribute(DEPLOYMENT_ENVIRONMENT,
+                                   environment)
+                span.set_attribute(SERVICE_NAME,
+                                   application_name)
+                span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_IS_STREAM,
+                                    False)
+                span.set_attribute(SemanticConvetion.GEN_AI_SERVER_TTFT,
+                                    end_time - start_time)
+                span.set_attribute(SemanticConvetion.GEN_AI_SDK_VERSION,
+                                    version)
+                if capture_message_content:
                     span.add_event(
                         name=SemanticConvetion.GEN_AI_CONTENT_PROMPT_EVENT,
                         attributes={
                             SemanticConvetion.GEN_AI_CONTENT_PROMPT: prompt,
                         },
                     )
-                span.set_attribute(SemanticConvetion.GEN_AI_USAGE_INPUT_TOKENS,
-                                prompt_tokens)
 
                 i = 0
-                completion_tokens = 0
+                output_tokens = 0
                 for completion in response:
                     if len(response) > 1:
                         attribute_name = f"gen_ai.content.completion.{i}"
                     else:
                         attribute_name = SemanticConvetion.GEN_AI_CONTENT_COMPLETION_EVENT
-                    if trace_content:
+                    if capture_message_content:
                         # pylint: disable=bare-except
                         try:
                             llm_response = completion.get('generated_text', '')
@@ -115,43 +138,51 @@ def text_wrap(gen_ai_endpoint, version, environment, application_name,
                         span.add_event(
                             name=attribute_name,
                             attributes={
-                                # pylint: disable=line-too-long
                                 SemanticConvetion.GEN_AI_CONTENT_COMPLETION: llm_response,
                             },
                         )
-                    completion_tokens += general_tokens(llm_response)
+                    output_tokens += general_tokens(llm_response)
 
                     i=i+1
                 span.set_attribute(SemanticConvetion.GEN_AI_USAGE_OUTPUT_TOKENS,
-                                   completion_tokens)
+                                   output_tokens)
                 span.set_attribute(SemanticConvetion.GEN_AI_USAGE_TOTAL_TOKENS,
-                                   prompt_tokens + completion_tokens)
+                                   input_tokens + output_tokens)
+
+                # Calculate cost of the operation
+                cost = get_chat_model_cost(request_model,
+                                            pricing_info, input_tokens,
+                                            output_tokens)
+                span.set_attribute(SemanticConvetion.GEN_AI_USAGE_COST,
+                                    cost)
+
                 span.set_status(Status(StatusCode.OK))
 
                 if disable_metrics is False:
-                    attributes = {
-                        TELEMETRY_SDK_NAME:
-                            "openlit",
-                        SemanticConvetion.GEN_AI_APPLICATION_NAME:
-                            application_name,
-                        SemanticConvetion.GEN_AI_SYSTEM:
-                            SemanticConvetion.GEN_AI_SYSTEM_HUGGING_FACE,
-                        SemanticConvetion.GEN_AI_ENVIRONMENT:
-                            environment,
-                        SemanticConvetion.GEN_AI_OPERATION:
-                            SemanticConvetion.GEN_AI_OPERATION_TYPE_CHAT,
-                        SemanticConvetion.GEN_AI_REQUEST_MODEL:
-                            instance.model.config.name_or_path
-                    }
+                    attributes = create_metrics_attributes(
+                        service_name=application_name,
+                        deployment_environment=environment,
+                        operation=SemanticConvetion.GEN_AI_OPERATION_TYPE_CHAT,
+                        system=SemanticConvetion.GEN_AI_SYSTEM_HUGGING_FACE,
+                        request_model=request_model,
+                        server_address=server_address,
+                        server_port=server_port,
+                        response_model=request_model,
+                    )
 
+                    metrics["genai_client_usage_tokens"].record(
+                        input_tokens + output_tokens, attributes
+                    )
+                    metrics["genai_client_operation_duration"].record(
+                        end_time - start_time, attributes
+                    )
+                    metrics["genai_server_ttft"].record(
+                        end_time - start_time, attributes
+                    )
                     metrics["genai_requests"].add(1, attributes)
-                    metrics["genai_total_tokens"].add(
-                        prompt_tokens +
-                        completion_tokens, attributes)
-                    metrics["genai_completion_tokens"].add(
-                        completion_tokens, attributes)
-                    metrics["genai_prompt_tokens"].add(
-                        prompt_tokens, attributes)
+                    metrics["genai_completion_tokens"].add(output_tokens, attributes)
+                    metrics["genai_prompt_tokens"].add(input_tokens, attributes)
+                    metrics["genai_cost"].record(cost, attributes)
 
                 # Return original response
                 return response

openlit/instrumentation/vertexai/__init__.py

@@ -7,16 +7,10 @@ from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from wrapt import wrap_function_wrapper
 
 from openlit.instrumentation.vertexai.vertexai import (
-    generate_content, predict, predict_streaming,
-    send_message, start_chat, start_chat_streaming,
-    embeddings
+    send_message
 )
 from openlit.instrumentation.vertexai.async_vertexai import (
-    generate_content_async, predict_async,
-    predict_streaming_async,
-    send_message_async,
-    start_chat_async, start_chat_streaming_async,
-    embeddings_async
+    async_send_message
 )
 
 
@@ -36,7 +30,7 @@ class VertexAIInstrumentor(BaseInstrumentor):
         tracer = kwargs.get("tracer")
         metrics = kwargs.get("metrics_dict")
         pricing_info = kwargs.get("pricing_info", {})
-        trace_content = kwargs.get("trace_content", False)
+        capture_message_content = kwargs.get("capture_message_content", False)
         disable_metrics = kwargs.get("disable_metrics")
         version = importlib.metadata.version("google-cloud-aiplatform")
 
@@ -44,102 +38,58 @@ class VertexAIInstrumentor(BaseInstrumentor):
         wrap_function_wrapper(
             "vertexai.generative_models",
             "GenerativeModel.generate_content",
-            generate_content("vertexai.generate_content", version, environment, application_name,
-                     tracer, pricing_info, trace_content, metrics, disable_metrics),
+            send_message(version, environment, application_name,
+                     tracer, pricing_info, capture_message_content, metrics, disable_metrics),
         )
 
         wrap_function_wrapper(
             "vertexai.generative_models",
             "ChatSession.send_message",
-            send_message("vertexai.send_message", version, environment, application_name,
-                     tracer, pricing_info, trace_content, metrics, disable_metrics),
-        )
-
-        wrap_function_wrapper(
-            "vertexai.language_models",  
-            "TextGenerationModel.predict",  
-            predict("vertexai.predict", version, environment, application_name,
-                     tracer, pricing_info, trace_content, metrics, disable_metrics),
-        )
-
-        wrap_function_wrapper(
-            "vertexai.language_models",
-            "TextGenerationModel.predict_streaming",
-            predict_streaming("vertexai.predict", version, environment, application_name,
-                     tracer, pricing_info, trace_content, metrics, disable_metrics),
+            send_message(version, environment, application_name,
+                     tracer, pricing_info, capture_message_content, metrics, disable_metrics),
         )
 
         wrap_function_wrapper(
             "vertexai.language_models",
             "ChatSession.send_message",
-            start_chat("vertexai.send_message", version, environment, application_name,
-                     tracer, pricing_info, trace_content, metrics, disable_metrics),
+            send_message(version, environment, application_name,
+                     tracer, pricing_info, capture_message_content, metrics, disable_metrics),
         )
 
         wrap_function_wrapper(
             "vertexai.language_models",
             "ChatSession.send_message_streaming",
-            start_chat_streaming("vertexai.send_message", version, environment, application_name,
-                     tracer, pricing_info, trace_content, metrics, disable_metrics),
-        )
-
-        wrap_function_wrapper(
-            "vertexai.language_models",
-            "TextEmbeddingModel.get_embeddings",
-            embeddings("vertexai.get_embeddings", version, environment, application_name,
-                     tracer, pricing_info, trace_content, metrics, disable_metrics),
+            send_message(version, environment, application_name,
+                     tracer, pricing_info, capture_message_content, metrics, disable_metrics),
         )
 
         #async
         wrap_function_wrapper(
             "vertexai.generative_models",
             "GenerativeModel.generate_content_async",
-            generate_content_async("vertexai.generate_content", version, environment,
-                                   application_name, tracer, pricing_info, trace_content,
-                                   metrics, disable_metrics),
+            async_send_message(version, environment, application_name,
+                    tracer, pricing_info, capture_message_content, metrics, disable_metrics),
         )
 
         wrap_function_wrapper(
             "vertexai.generative_models",
             "ChatSession.send_message_async",
-            send_message_async("vertexai.send_message", version, environment, application_name,
-                     tracer, pricing_info, trace_content, metrics, disable_metrics),
-        )
-
-        wrap_function_wrapper(
-            "vertexai.language_models",
-            "TextGenerationModel.predict_async",
-            predict_async("vertexai.predict", version, environment, application_name,
-                     tracer, pricing_info, trace_content, metrics, disable_metrics),
-        )
-
-        wrap_function_wrapper(
-            "vertexai.language_models",
-            "TextGenerationModel.predict_streaming_async",
-            predict_streaming_async("vertexai.predict", version, environment, application_name,
-                     tracer, pricing_info, trace_content, metrics, disable_metrics),
+            async_send_message(version, environment, application_name,
+                    tracer, pricing_info, capture_message_content, metrics, disable_metrics),
         )
 
         wrap_function_wrapper(
             "vertexai.language_models",
             "ChatSession.send_message_async",
-            start_chat_async("vertexai.send_message", version, environment, application_name,
-                     tracer, pricing_info, trace_content, metrics, disable_metrics),
+            async_send_message(version, environment, application_name,
+                    tracer, pricing_info, capture_message_content, metrics, disable_metrics),
         )
 
         wrap_function_wrapper(
             "vertexai.language_models",
             "ChatSession.send_message_streaming_async",
-            start_chat_streaming_async("vertexai.send_message", version, environment,
-                                       application_name, tracer, pricing_info, trace_content,
-                                       metrics, disable_metrics),
-        )
-
-        wrap_function_wrapper(
-            "vertexai.language_models",
-            "TextEmbeddingModel.get_embeddings_async",
-            embeddings_async("vertexai.get_embeddings", version, environment, application_name,
-                     tracer, pricing_info, trace_content, metrics, disable_metrics),
+            async_send_message(version, environment, application_name,
+                    tracer, pricing_info, capture_message_content, metrics, disable_metrics),
         )
 
     def _uninstrument(self, **kwargs):