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

openlit/instrumentation/bedrock/__init__.py

@@ -32,7 +32,7 @@ class BedrockInstrumentor(BaseInstrumentor):
         wrap_function_wrapper(
             "botocore.client",  
             "ClientCreator.create_client",  
-            converse("bedrock.converse", version, environment, application_name,
+            converse(version, environment, application_name,
                      tracer, pricing_info, trace_content, metrics, disable_metrics),
         )
 

openlit/instrumentation/bedrock/bedrock.py

@@ -1,17 +1,22 @@
-# pylint: disable=duplicate-code, broad-exception-caught, too-many-statements, unused-argument, protected-access, too-many-branches
 """
 Module for monitoring Amazon Bedrock API calls.
 """
 
 import logging
+import time
 from botocore.response import StreamingBody
 from botocore.exceptions import ReadTimeoutError, ResponseStreamingError
 from urllib3.exceptions import ProtocolError as URLLib3ProtocolError
 from urllib3.exceptions import ReadTimeoutError as URLLib3ReadTimeoutError
 from opentelemetry.trace import SpanKind, Status, StatusCode
-from opentelemetry.sdk.resources import TELEMETRY_SDK_NAME
-from openlit.__helpers import get_chat_model_cost
-from openlit.__helpers import handle_exception
+from opentelemetry.sdk.resources import SERVICE_NAME, TELEMETRY_SDK_NAME, DEPLOYMENT_ENVIRONMENT
+from openlit.__helpers import (
+    get_chat_model_cost,
+    handle_exception,
+    response_as_dict,
+    create_metrics_attributes,
+    set_server_address_and_port
+)
 from openlit.semcov import SemanticConvetion
 
 # Initialize logger for logging potential issues and operations
@@ -47,8 +52,7 @@ class CustomStreamWrapper(StreamingBody):
 
         return data_chunk
 
-
-def converse(gen_ai_endpoint, version, environment, application_name, tracer,
+def converse(version, environment, application_name, tracer,
          pricing_info, trace_content, metrics, disable_metrics):
     """
     Generates a telemetry wrapper for messages to collect metrics.
@@ -91,52 +95,104 @@ def converse(gen_ai_endpoint, version, environment, application_name, tracer,
             Returns:
                 The modified response with telemetry.
             """
-            with tracer.start_as_current_span(gen_ai_endpoint, kind=SpanKind.CLIENT) as span:
+
+            server_address, server_port = set_server_address_and_port(instance, 'aws.amazon.com', 443)
+            request_model = method_kwargs.get('modelId', 'amazon.titan-text-express-v1')
+
+            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 = original_method(*method_args, **method_kwargs)
+                end_time = time.time()
+
+                response_dict = response_as_dict(response)
 
                 try:
-                    message_prompt = method_kwargs.get("messages", "")
+                    message_prompt = method_kwargs.get('messages', '')
                     formatted_messages = []
                     for message in message_prompt:
-                        role = message["role"]
-                        content = message["content"]
+                        role = message['role']
+                        content = message['content']
 
                         if isinstance(content, list):
-                            # pylint: disable=line-too-long
                             content_str = ", ".join(f'text: {item["text"]}' for item in content if "text" in item)
-                            formatted_messages.append(f"{role}: {content_str}")
+                            formatted_messages.append(f'{role}: {content_str}')
                         else:
-                            formatted_messages.append(f"{role}: {content}")
-                    prompt = "\n".join(formatted_messages)
+                            formatted_messages.append(f'{role}: {content}')
+                    prompt = '\n'.join(formatted_messages)
+
+                    input_tokens = response_dict.get('usage').get('inputTokens')
+                    output_tokens = response_dict.get('usage').get('outputTokens')
 
-                    model = method_kwargs.get("modelId", "amazon.titan-text-express-v1")
-                    input_tokens = response["usage"]["inputTokens"]
-                    output_tokens = response["usage"]["outputTokens"]
+                    # Calculate cost of the operation
+                    cost = get_chat_model_cost(request_model, pricing_info,
+                                                input_tokens, output_tokens)
+
+                    llm_response = response_dict.get('output').get('message').get('content')[0].get('text')
 
-                    span.set_attribute(TELEMETRY_SDK_NAME, "openlit")
+                    # Set base span attribues (OTel Semconv)
+                    span.set_attribute(TELEMETRY_SDK_NAME, 'openlit')
+                    span.set_attribute(SemanticConvetion.GEN_AI_OPERATION,
+                                        SemanticConvetion.GEN_AI_OPERATION_TYPE_CHAT)
                     span.set_attribute(SemanticConvetion.GEN_AI_SYSTEM,
-                                        SemanticConvetion.GEN_AI_SYSTEM_BEDROCK)
-                    span.set_attribute(SemanticConvetion.GEN_AI_ENDPOINT,
-                                        gen_ai_endpoint)
-                    span.set_attribute(SemanticConvetion.GEN_AI_ENVIRONMENT,
-                                        environment)
-                    span.set_attribute(SemanticConvetion.GEN_AI_APPLICATION_NAME,
-                                        application_name)
+                                        SemanticConvetion.GEN_AI_SYSTEM_AWS_BEDROCK)
                     span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_MODEL,
-                                        model)
-                    span.set_attribute(SemanticConvetion.GEN_AI_USAGE_PROMPT_TOKENS,
+                                        request_model)
+                    span.set_attribute(SemanticConvetion.SERVER_PORT,
+                                        server_port)
+
+                    inference_config = method_kwargs.get('inferenceConfig', {})
+
+                    # List of attributes and their config keys
+                    attributes = [
+                        (SemanticConvetion.GEN_AI_REQUEST_FREQUENCY_PENALTY, 'frequencyPenalty'),
+                        (SemanticConvetion.GEN_AI_REQUEST_MAX_TOKENS, 'maxTokens'),
+                        (SemanticConvetion.GEN_AI_REQUEST_PRESENCE_PENALTY, 'presencePenalty'),
+                        (SemanticConvetion.GEN_AI_REQUEST_STOP_SEQUENCES, 'stopSequences'),
+                        (SemanticConvetion.GEN_AI_REQUEST_TEMPERATURE, 'temperature'),
+                        (SemanticConvetion.GEN_AI_REQUEST_TOP_P, 'topP'),
+                        (SemanticConvetion.GEN_AI_REQUEST_TOP_K, 'topK'),
+                    ]
+
+                    # Set each attribute if the corresponding value exists and is not None
+                    for attribute, key in attributes:
+                        value = inference_config.get(key)
+                        if value is not None:
+                            span.set_attribute(attribute, value)
+
+                    span.set_attribute(SemanticConvetion.GEN_AI_RESPONSE_ID,
+                                        response_dict.get('ResponseMetadata').get('RequestId'))
+                    span.set_attribute(SemanticConvetion.GEN_AI_RESPONSE_MODEL,
+                                        request_model)
+                    span.set_attribute(SemanticConvetion.GEN_AI_USAGE_INPUT_TOKENS,
                                         input_tokens)
-                    span.set_attribute(SemanticConvetion.GEN_AI_USAGE_COMPLETION_TOKENS,
+                    span.set_attribute(SemanticConvetion.GEN_AI_USAGE_OUTPUT_TOKENS,
                                         output_tokens)
+                    span.set_attribute(SemanticConvetion.SERVER_ADDRESS,
+                                        server_address)
+                    if isinstance(llm_response, str):
+                        span.set_attribute(SemanticConvetion.GEN_AI_OUTPUT_TYPE,
+                                        'text')
+                    else:
+                        span.set_attribute(SemanticConvetion.GEN_AI_OUTPUT_TYPE,
+                                        'json')
+
+                    # Set base span attribues (Extras)
+                    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_USAGE_TOTAL_TOKENS,
                                         input_tokens + output_tokens)
-
-                    # Calculate cost of the operation
-                    cost = get_chat_model_cost(model,
-                                            pricing_info, input_tokens,
-                                            output_tokens)
                     span.set_attribute(SemanticConvetion.GEN_AI_USAGE_COST,
-                                cost)
+                                        cost)
+                    span.set_attribute(SemanticConvetion.GEN_AI_SERVER_TTFT,
+                                        end_time - start_time)
+                    span.set_attribute(SemanticConvetion.GEN_AI_SDK_VERSION,
+                                        version)
 
                     if trace_content:
                         span.add_event(
@@ -148,42 +204,43 @@ def converse(gen_ai_endpoint, version, environment, application_name, tracer,
                         span.add_event(
                             name=SemanticConvetion.GEN_AI_CONTENT_COMPLETION_EVENT,
                             attributes={
-                                # pylint: disable=line-too-long
-                                SemanticConvetion.GEN_AI_CONTENT_COMPLETION: response["output"]["message"]["content"][0]["text"],
+                                SemanticConvetion.GEN_AI_CONTENT_COMPLETION: llm_response,
                             },
                         )
 
                     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_BEDROCK,
-                            SemanticConvetion.GEN_AI_ENVIRONMENT:
-                                environment,
-                            SemanticConvetion.GEN_AI_TYPE:
-                                SemanticConvetion.GEN_AI_TYPE_CHAT,
-                            SemanticConvetion.GEN_AI_REQUEST_MODEL:
-                                model
-                        }
-
-                        metrics["genai_requests"].add(1, attributes)
-                        metrics["genai_total_tokens"].add(
+                        attributes = create_metrics_attributes(
+                            service_name=application_name,
+                            deployment_environment=environment,
+                            operation=SemanticConvetion.GEN_AI_OPERATION_TYPE_CHAT,
+                            system=SemanticConvetion.GEN_AI_SYSTEM_AWS_BEDROCK,
+                            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_completion_tokens"].add(output_tokens, attributes)
-                        metrics["genai_prompt_tokens"].add(input_tokens, attributes)
-                        metrics["genai_cost"].record(cost, 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_completion_tokens'].add(output_tokens, attributes)
+                        metrics['genai_prompt_tokens'].add(input_tokens, attributes)
+                        metrics['genai_cost'].record(cost, attributes)
 
                     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 original response
                     return response
@@ -192,7 +249,7 @@ def converse(gen_ai_endpoint, version, environment, application_name, tracer,
         client = wrapped(*args, **kwargs)
 
         # Replace the original method with the instrumented one
-        if kwargs.get("service_name") == "bedrock-runtime":
+        if kwargs.get('service_name') == 'bedrock-runtime':
             original_invoke_model = client.converse
             client.converse = lambda *args, **kwargs: converse_wrapper(original_invoke_model,
                                                                             *args, **kwargs)

openlit/instrumentation/chroma/chroma.py

@@ -77,8 +77,8 @@ def general_wrap(gen_ai_endpoint, version, environment, application_name,
                                    environment)
                 span.set_attribute(SemanticConvetion.GEN_AI_APPLICATION_NAME,
                                    application_name)
-                span.set_attribute(SemanticConvetion.GEN_AI_TYPE,
-                                   SemanticConvetion.GEN_AI_TYPE_VECTORDB)
+                span.set_attribute(SemanticConvetion.GEN_AI_OPERATION,
+                                   SemanticConvetion.GEN_AI_OPERATION_TYPE_VECTORDB)
                 span.set_attribute(SemanticConvetion.DB_SYSTEM,
                                    SemanticConvetion.DB_SYSTEM_CHROMA)
                 span.set_attribute(SemanticConvetion.DB_COLLECTION_NAME,
@@ -179,8 +179,8 @@ def general_wrap(gen_ai_endpoint, version, environment, application_name,
                             SemanticConvetion.DB_SYSTEM_CHROMA,
                         SemanticConvetion.GEN_AI_ENVIRONMENT:
                             environment,
-                        SemanticConvetion.GEN_AI_TYPE:
-                            SemanticConvetion.GEN_AI_TYPE_VECTORDB,
+                        SemanticConvetion.GEN_AI_OPERATION:
+                            SemanticConvetion.GEN_AI_OPERATION_TYPE_VECTORDB,
                         SemanticConvetion.DB_OPERATION:
                             db_operation
                     }

openlit/instrumentation/cohere/__init__.py

@@ -6,8 +6,10 @@ from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from wrapt import wrap_function_wrapper
 
 from openlit.instrumentation.cohere.cohere import chat, chat_stream, embed
+from openlit.instrumentation.cohere.async_cohere import async_chat, async_chat_stream, async_embed
 
-_instruments = ("cohere >= 5.3.2",)
+
+_instruments = ("cohere >= 5.14.0",)
 
 class CohereInstrumentor(BaseInstrumentor):
     """An instrumentor for Cohere's client library."""
@@ -25,24 +27,45 @@ class CohereInstrumentor(BaseInstrumentor):
         disable_metrics = kwargs.get("disable_metrics")
         version = importlib.metadata.version("cohere")
 
+        # Sync Client
+        wrap_function_wrapper(
+            "cohere.client_v2",  
+            "ClientV2.chat",  
+            chat(version, environment, application_name,
+                 tracer, pricing_info, trace_content, metrics, disable_metrics),
+        )
+        wrap_function_wrapper(
+            "cohere.client_v2",  
+            "ClientV2.chat_stream",  
+            chat_stream(version, environment, application_name,
+                        tracer, pricing_info, trace_content, metrics, disable_metrics),
+        )
+        wrap_function_wrapper(
+            "cohere.client_v2",  
+            "ClientV2.embed",  
+            embed(version, environment, application_name,
+                  tracer, pricing_info, trace_content, metrics, disable_metrics),
+        )
+
+        # Async Client
         wrap_function_wrapper(
-            "cohere.client",  
-            "Client.chat",  
-            chat("cohere.chat", version, environment, application_name,
+            "cohere.client_v2",  
+            "AsyncClientV2.chat",  
+            async_chat(version, environment, application_name,
                  tracer, pricing_info, trace_content, metrics, disable_metrics),
         )
 
         wrap_function_wrapper(
-            "cohere.client",  
-            "Client.chat_stream",  
-            chat_stream("cohere.chat", version, environment, application_name,
+            "cohere.client_v2",  
+            "AsyncClientV2.chat_stream",  
+            async_chat_stream(version, environment, application_name,
                         tracer, pricing_info, trace_content, metrics, disable_metrics),
         )
 
         wrap_function_wrapper(
-            "cohere.client",  
-            "Client.embed",  
-            embed("cohere.embed", version, environment, application_name,
+            "cohere.client_v2",  
+            "AsyncClientV2.embed",  
+            async_embed(version, environment, application_name,
                   tracer, pricing_info, trace_content, metrics, disable_metrics),
         )