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

openlit/__helpers.py

@@ -7,8 +7,8 @@ import json
 import logging
 from urllib.parse import urlparse
 from typing import Any, Dict, List, Tuple
+import math
 import requests
-import tiktoken
 from opentelemetry.sdk.resources import SERVICE_NAME, TELEMETRY_SDK_NAME, DEPLOYMENT_ENVIRONMENT
 from opentelemetry.trace import Status, StatusCode
 from opentelemetry._events import Event
@@ -21,12 +21,13 @@ def response_as_dict(response):
     """
     Return parsed response as a dict
     """
+
     # pylint: disable=no-else-return
     if isinstance(response, dict):
         return response
-    if hasattr(response, "model_dump"):
+    if hasattr(response, 'model_dump'):
         return response.model_dump()
-    elif hasattr(response, "parse"):
+    elif hasattr(response, 'parse'):
         return response_as_dict(response.parse())
     else:
         return response
@@ -34,8 +35,8 @@ def response_as_dict(response):
 def get_env_variable(name, arg_value, error_message):
     """
     Retrieve an environment variable if the argument is not provided
-    and raise an error if both are not set.
     """
+
     if arg_value is not None:
         return arg_value
     value = os.getenv(name)
@@ -44,58 +45,21 @@ def get_env_variable(name, arg_value, error_message):
         raise RuntimeError(error_message)
     return value
 
-def openai_tokens(text, model):
-    """
-    Calculate the number of tokens a given text would take up for a specified model.
-    
-    Args:
-        text (str): The input text to be encoded.
-        model (str): The model identifier used for encoding.
-        
-    Returns:
-        int: The number of tokens the text is encoded into.
-    """
-    try:
-        encoding = tiktoken.encoding_for_model(model)
-    except:
-        encoding = tiktoken.get_encoding("cl100k_base")
-
-    num_tokens = len(encoding.encode(text))
-    return num_tokens
-
 def general_tokens(text):
     """
     Calculate the number of tokens a given text would take up.
-    
-    Args:
-        text (str): The input text to be encoded.
-        model (str): The model identifier used for encoding.
-        
-    Returns:
-        int: The number of tokens the text is encoded into.
     """
 
-    encoding = tiktoken.get_encoding("gpt2")
-
-    num_tokens = len(encoding.encode(text))
-    return num_tokens
+    return math.ceil(len(text) / 2)
 
 def get_chat_model_cost(model, pricing_info, prompt_tokens, completion_tokens):
     """
     Retrieve the cost of processing for a given model based on prompt and tokens.
-    
-    Args:
-        model (str): The model identifier.
-        pricing_info (dict): A dictionary containing pricing information for various models.
-        prompt_tokens (int): Number of tokens in the prompt.
-        completion_tokens (int): Number of tokens in the completion if applicable.
-        
-    Returns:
-        float: The calculated cost for the operation.
     """
+
     try:
-        cost = ((prompt_tokens / 1000) * pricing_info["chat"][model]["promptPrice"]) + \
-            ((completion_tokens / 1000) * pricing_info["chat"][model]["completionPrice"])
+        cost = ((prompt_tokens / 1000) * pricing_info['chat'][model]['promptPrice']) + \
+            ((completion_tokens / 1000) * pricing_info['chat'][model]['completionPrice'])
     except:
         cost = 0
     return cost
@@ -103,17 +67,10 @@ def get_chat_model_cost(model, pricing_info, prompt_tokens, completion_tokens):
 def get_embed_model_cost(model, pricing_info, prompt_tokens):
     """
     Retrieve the cost of processing for a given model based on prompt tokens.
-    
-    Args:
-        model (str): The model identifier.
-        pricing_info (dict): A dictionary containing pricing information for various models.
-        prompt_tokens (int): Number of tokens in the prompt.
-        
-    Returns:
-        float: The calculated cost for the operation.
     """
+
     try:
-        cost = (prompt_tokens / 1000) * pricing_info["embeddings"][model]
+        cost = (prompt_tokens / 1000) * pricing_info['embeddings'][model]
     except:
         cost = 0
     return cost
@@ -121,18 +78,10 @@ def get_embed_model_cost(model, pricing_info, prompt_tokens):
 def get_image_model_cost(model, pricing_info, size, quality):
     """
     Retrieve the cost of processing for a given model based on image size and quailty.
-    
-    Args:
-        model (str): The model identifier.
-        pricing_info (dict): A dictionary containing pricing information for various models.
-        size (str): Size of the Image.
-        quality (int): Quality of the Image.
-        
-    Returns:
-        float: The calculated cost for the operation.
     """
+
     try:
-        cost = pricing_info["images"][model][quality][size]
+        cost = pricing_info['images'][model][quality][size]
     except:
         cost = 0
     return cost
@@ -140,20 +89,13 @@ def get_image_model_cost(model, pricing_info, size, quality):
 def get_audio_model_cost(model, pricing_info, prompt, duration=None):
     """
     Retrieve the cost of processing for a given model based on prompt.
-    
-    Args:
-        model (str): The model identifier.
-        pricing_info (dict): A dictionary containing pricing information for various models.
-        prompt (str): Prompt to the LLM Model
-        
-    Returns:
-        float: The calculated cost for the operation.
     """
+
     try:
         if prompt:
-            cost = (len(prompt) / 1000) * pricing_info["audio"][model]
+            cost = (len(prompt) / 1000) * pricing_info['audio'][model]
         else:
-            cost = duration * pricing_info["audio"][model]
+            cost = duration * pricing_info['audio'][model]
     except:
         cost = 0
     return cost
@@ -161,15 +103,10 @@ def get_audio_model_cost(model, pricing_info, prompt, duration=None):
 def fetch_pricing_info(pricing_json=None):
     """
     Fetches pricing information from a specified URL or File Path.
-
-    Args:
-        pricing_json(str): path or url to the pricing json file
-
-    Returns:
-        dict: The pricing json
     """
+
     if pricing_json:
-        is_url = urlparse(pricing_json).scheme != ""
+        is_url = urlparse(pricing_json).scheme != ''
         if is_url:
             pricing_url = pricing_json
         else:
@@ -177,39 +114,36 @@ def fetch_pricing_info(pricing_json=None):
                 with open(pricing_json, mode='r', encoding='utf-8') as f:
                     return json.load(f)
             except FileNotFoundError:
-                logger.error("Pricing information file not found: %s", pricing_json)
+                logger.error('Pricing information file not found: %s', pricing_json)
             except json.JSONDecodeError:
-                logger.error("Error decoding JSON from file: %s", pricing_json)
+                logger.error('Error decoding JSON from file: %s', pricing_json)
             except Exception as file_err:
-                logger.error("Unexpected error occurred while reading file: %s", file_err)
+                logger.error('Unexpected error occurred while reading file: %s', file_err)
             return {}
     else:
-        pricing_url = "https://raw.githubusercontent.com/openlit/openlit/main/assets/pricing.json"
+        pricing_url = 'https://raw.githubusercontent.com/openlit/openlit/main/assets/pricing.json'
     try:
         # Set a timeout of 10 seconds for both the connection and the read
         response = requests.get(pricing_url, timeout=20)
         response.raise_for_status()
         return response.json()
     except requests.HTTPError as http_err:
-        logger.error("HTTP error occured while fetching pricing info: %s", http_err)
+        logger.error('HTTP error occured while fetching pricing info: %s', http_err)
     except Exception as err:
-        logger.error("Unexpected error occurred while fetching pricing info: %s", err)
+        logger.error('Unexpected error occurred while fetching pricing info: %s', err)
     return {}
 
 def handle_exception(span,e):
     """Handles Exception when LLM Function fails or trace creation fails."""
-    # Record the exception details within the span
+
     span.record_exception(e)
     span.set_status(Status(StatusCode.ERROR))
 
 def calculate_ttft(timestamps: List[float], start_time: float) -> float:
     """
     Calculate the time to the first tokens.
-    
-    :param timestamps: List of timestamps for received tokens
-    :param start_time: The start time of the streaming process
-    :return: Time to the first tokens
     """
+
     if timestamps:
         return timestamps[0] - start_time
     return 0.0
@@ -217,10 +151,8 @@ def calculate_ttft(timestamps: List[float], start_time: float) -> float:
 def calculate_tbt(timestamps: List[float]) -> float:
     """
     Calculate the average time between tokens.
-    
-    :param timestamps: List of timestamps for received tokens
-    :return: Average time between tokens
     """
+
     if len(timestamps) > 1:
         time_diffs = [timestamps[i] - timestamps[i - 1] for i in range(1, len(timestamps))]
         return sum(time_diffs) / len(time_diffs)
@@ -239,8 +171,9 @@ def create_metrics_attributes(
     """
     Returns OTel metrics attributes
     """
+
     return {
-        TELEMETRY_SDK_NAME: "openlit",
+        TELEMETRY_SDK_NAME: 'openlit',
         SERVICE_NAME: service_name,
         DEPLOYMENT_ENVIRONMENT: deployment_environment,
         SemanticConvetion.GEN_AI_OPERATION: operation,
@@ -259,18 +192,18 @@ def set_server_address_and_port(client_instance: Any,
     """
 
     # Try getting base_url from multiple potential attributes
-    base_client = getattr(client_instance, "_client", None)
-    base_url = getattr(base_client, "base_url", None)
+    base_client = getattr(client_instance, '_client', None)
+    base_url = getattr(base_client, 'base_url', None)
 
     if not base_url:
         # Attempt to get endpoint from instance._config.endpoint if base_url is not set
-        config = getattr(client_instance, "_config", None)
-        base_url = getattr(config, "endpoint", None)
+        config = getattr(client_instance, '_config', None)
+        base_url = getattr(config, 'endpoint', None)
 
     if not base_url:
         # Attempt to get server_url from instance.sdk_configuration.server_url
-        config = getattr(client_instance, "sdk_configuration", None)
-        base_url = getattr(config, "server_url", None)
+        config = getattr(client_instance, 'sdk_configuration', None)
+        base_url = getattr(config, 'server_url', None)
 
     if base_url:
         if isinstance(base_url, str):
@@ -278,8 +211,8 @@ def set_server_address_and_port(client_instance: Any,
             server_address = url.hostname or default_server_address
             server_port = url.port if url.port is not None else default_server_port
         else:  # base_url might not be a str; handle as an object.
-            server_address = getattr(base_url, "host", None) or default_server_address
-            port_attr = getattr(base_url, "port", None)
+            server_address = getattr(base_url, 'host', None) or default_server_address
+            port_attr = getattr(base_url, 'port', None)
             server_port = port_attr if port_attr is not None else default_server_port
     else:  # no base_url or endpoint provided; use defaults.
         server_address = default_server_address
@@ -301,59 +234,74 @@ def otel_event(name, attributes, body):
 def extract_and_format_input(messages):
     """
     Process a list of messages to extract content and categorize
-    them into fixed roles like 'user', 'assistant', 'system'.
+    them into fixed roles like 'user', 'assistant', 'system', 'tool'.
     """
 
     fixed_roles = ['user', 'assistant', 'system', 'tool']  # Ensure these are your fixed keys
     # Initialize the dictionary with fixed keys and empty structures
-    formatted_messages = {role_key: {"role": "", "content": ""} for role_key in fixed_roles}
+    formatted_messages = {role_key: {'role': '', 'content': ''} for role_key in fixed_roles}
 
     for message in messages:
         # Normalize the message structure
         message = response_as_dict(message)
 
         # Extract role and content
-        role = message.get("role")
+        role = message.get('role')
         if role not in fixed_roles:
             continue  # Skip any role not in our predefined roles
 
-        content = message.get("content", "")
+        content = message.get('content', '')
 
         # Prepare content as a string
         if isinstance(content, list):
             content_str = ", ".join(
-                # pylint: disable=line-too-long
-                f'{item.get("type", "text")}: {item.get("text", item.get("image_url", "").get("url", "") if isinstance(item.get("image_url", ""), dict) else item.get("image_url", ""))}'
+                f'{item.get("type", "text")}: {extract_text_from_item(item)}'
                 for item in content
             )
         else:
             content_str = content
 
         # Set the role in the formatted message and concatenate content
-        if not formatted_messages[role]["role"]:
-            formatted_messages[role]["role"] = role
+        if not formatted_messages[role]['role']:
+            formatted_messages[role]['role'] = role
 
-        if formatted_messages[role]["content"]:
-            formatted_messages[role]["content"] += " " + content_str
+        if formatted_messages[role]['content']:
+            formatted_messages[role]['content'] += ' ' + content_str
         else:
-            formatted_messages[role]["content"] = content_str
+            formatted_messages[role]['content'] = content_str
 
     return formatted_messages
 
+def extract_text_from_item(item):
+    """
+    Extract text from inpit message
+    """
+
+    #pylint: disable=no-else-return
+    if item.get('type') == 'text':
+        return item.get('text', '')
+    elif item.get('type') == 'image':
+        # Handle image content specifically checking for 'url' or 'base64'
+        source = item.get('source', {})
+        if isinstance(source, dict):
+            if source.get('type') == 'base64':
+                # Return the actual base64 data if present
+                return source.get('data', '[Missing base64 data]')
+            elif source.get('type') == 'url':
+                return source.get('url', '[Missing URL]')
+    elif item.get('type') == 'image_url':
+        # New format: Handle the 'image_url' type
+        image_url = item.get('image_url', {})
+        if isinstance(image_url, dict):
+            return image_url.get('url', '[Missing image URL]')
+    return ''
+
 # To be removed one the change to log events (from span events) is complete
 def concatenate_all_contents(formatted_messages):
     """
-    Concatenate all 'content' fields from the formatted messages
-    dictionary into a single string.
-
-    Parameters:
-    - formatted_messages: Dictionary with roles as keys and corresponding
-      role and content as values.
-
-    Returns:
-    - A single string with all content concatenated.
+    Concatenate all 'content' fields into a single strin
     """
-    return " ".join(
+    return ' '.join(
         message_data['content']
         for message_data in formatted_messages.values()
         if message_data['content']

openlit/instrumentation/ag2/__init__.py

@@ -9,7 +9,7 @@ from openlit.instrumentation.ag2.ag2 import (
     conversable_agent, agent_run
 )
 
-_instruments = ("ag2 >= 0.3.2",)
+_instruments = ('ag2 >= 0.3.2',)
 
 class AG2Instrumentor(BaseInstrumentor):
     """
@@ -20,26 +20,26 @@ class AG2Instrumentor(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")
-        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("ag2")
+        application_name = kwargs.get('application_name', 'default_application')
+        environment = kwargs.get('environment', 'default_environment')
+        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('ag2')
 
         wrap_function_wrapper(
-            "autogen.agentchat.conversable_agent",
-            "ConversableAgent.__init__",
+            'autogen.agentchat.conversable_agent',
+            'ConversableAgent.__init__',
             conversable_agent(version, environment, application_name,
                   tracer, event_provider, pricing_info, capture_message_content, metrics, disable_metrics),
         )
 
         wrap_function_wrapper(
-            "autogen.agentchat.conversable_agent",
-            "ConversableAgent.run",
+            'autogen.agentchat.conversable_agent',
+            'ConversableAgent.run',
             agent_run(version, environment, application_name,
                   tracer, event_provider, pricing_info, capture_message_content, metrics, disable_metrics),
         )

openlit/instrumentation/ag2/ag2.py

@@ -28,7 +28,7 @@ def set_span_attributes(span, version, operation_name, environment,
     """
 
     # Set Span attributes (OTel Semconv)
-    span.set_attribute(TELEMETRY_SDK_NAME, "openlit")
+    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)
@@ -73,10 +73,10 @@ def emit_events(response, event_provider, capture_message_content):
                 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'
+                'index': response.chat_history.index(chat),
+                'message': {
+                    **({'content': chat['content']} if capture_message_content else {}),
+                    'role': 'assistant' if chat['role'] == 'user' else 'user'
                 }
             }
         )
@@ -92,12 +92,12 @@ def conversable_agent(version, environment, application_name,
         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')
+            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}"
+        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:
@@ -117,7 +117,7 @@ def conversable_agent(version, environment, application_name,
 
             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 wrapper
@@ -130,7 +130,7 @@ def agent_run(version, environment, application_name,
     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}"
+        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:
@@ -157,7 +157,7 @@ def agent_run(version, environment, application_name,
 
             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 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,40 +24,40 @@ 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")
-        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")
+        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, 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, 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, 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, event_provider, pricing_info, capture_message_content, metrics, disable_metrics),
         )

openlit/instrumentation/ai21/ai21.py

@@ -7,7 +7,6 @@ import time
 from opentelemetry.trace import SpanKind
 from openlit.__helpers import (
     handle_exception,
-    response_as_dict,
     set_server_address_and_port,
 )
 from openlit.instrumentation.ai21.utils import (
@@ -47,9 +46,9 @@ def chat(version, environment, application_name,
             self._span = span
             self._span_name = span_name
             # Placeholder for aggregating streaming response
-            self._llmresponse = ""
-            self._response_id = ""
-            self._finish_reason = ""
+            self._llmresponse = ''
+            self._response_id = ''
+            self._finish_reason = ''
             self._input_tokens = 0
             self._output_tokens = 0
             self._choices = []
@@ -100,7 +99,7 @@ def chat(version, environment, application_name,
                         )
                 except Exception as e:
                     handle_exception(self._span, e)
-                    logger.error("Error in trace creation: %s", e)
+                    logger.error('Error in trace creation: %s', e)
                 raise
 
     def wrapper(wrapped, instance, args, kwargs):
@@ -109,12 +108,12 @@ def chat(version, environment, application_name,
         """
 
         # Check if streaming is enabled for the API call
-        streaming = kwargs.get("stream", False)
+        streaming = kwargs.get('stream', False)
 
-        server_address, server_port = set_server_address_and_port(instance, "api.ai21.com", 443)
-        request_model = kwargs.get("model", "jamba-1.5-mini")
+        server_address, server_port = set_server_address_and_port(instance, 'api.ai21.com', 443)
+        request_model = kwargs.get('model', 'jamba-1.5-mini')
 
-        span_name = f"{SemanticConvetion.GEN_AI_OPERATION_TYPE_CHAT} {request_model}"
+        span_name = f'{SemanticConvetion.GEN_AI_OPERATION_TYPE_CHAT} {request_model}'
 
         # pylint: disable=no-else-return
         if streaming:
@@ -129,7 +128,7 @@ def chat(version, environment, application_name,
                 start_time = time.time()
                 response = wrapped(*args, **kwargs)
                 response = process_chat_response(
-                    response=response_as_dict(response),
+                    response=response,
                     request_model=request_model,
                     pricing_info=pricing_info,
                     server_port=server_port,
@@ -161,16 +160,16 @@ def chat_rag(version, environment, application_name,
         Wraps the GenAI function call.
         """
 
-        server_address, server_port = set_server_address_and_port(instance, "api.ai21.com", 443)
-        request_model = kwargs.get("model", "jamba-1.5-mini")
+        server_address, server_port = set_server_address_and_port(instance, 'api.ai21.com', 443)
+        request_model = kwargs.get('model', 'jamba-1.5-mini')
 
-        span_name = f"{SemanticConvetion.GEN_AI_OPERATION_TYPE_CHAT} {request_model}"
+        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)
             response = process_chat_rag_response(
-                response=response_as_dict(response),
+                response=response,
                 request_model=request_model,
                 pricing_info=pricing_info,
                 server_port=server_port,