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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (112) hide show
  1. openlit/__helpers.py +73 -0
  2. openlit/__init__.py +38 -11
  3. openlit/instrumentation/ag2/__init__.py +9 -10
  4. openlit/instrumentation/ag2/ag2.py +133 -68
  5. openlit/instrumentation/ai21/__init__.py +6 -5
  6. openlit/instrumentation/ai21/ai21.py +71 -534
  7. openlit/instrumentation/ai21/async_ai21.py +71 -534
  8. openlit/instrumentation/ai21/utils.py +407 -0
  9. openlit/instrumentation/anthropic/__init__.py +3 -3
  10. openlit/instrumentation/anthropic/anthropic.py +4 -4
  11. openlit/instrumentation/anthropic/async_anthropic.py +4 -4
  12. openlit/instrumentation/assemblyai/__init__.py +2 -2
  13. openlit/instrumentation/assemblyai/assemblyai.py +3 -3
  14. openlit/instrumentation/astra/__init__.py +25 -25
  15. openlit/instrumentation/astra/astra.py +2 -2
  16. openlit/instrumentation/astra/async_astra.py +2 -2
  17. openlit/instrumentation/azure_ai_inference/__init__.py +5 -5
  18. openlit/instrumentation/azure_ai_inference/async_azure_ai_inference.py +8 -8
  19. openlit/instrumentation/azure_ai_inference/azure_ai_inference.py +8 -8
  20. openlit/instrumentation/bedrock/__init__.py +2 -2
  21. openlit/instrumentation/bedrock/bedrock.py +3 -3
  22. openlit/instrumentation/chroma/__init__.py +9 -9
  23. openlit/instrumentation/chroma/chroma.py +2 -2
  24. openlit/instrumentation/cohere/__init__.py +7 -7
  25. openlit/instrumentation/cohere/async_cohere.py +9 -9
  26. openlit/instrumentation/cohere/cohere.py +9 -9
  27. openlit/instrumentation/controlflow/__init__.py +4 -4
  28. openlit/instrumentation/controlflow/controlflow.py +2 -2
  29. openlit/instrumentation/crawl4ai/__init__.py +3 -3
  30. openlit/instrumentation/crawl4ai/async_crawl4ai.py +2 -2
  31. openlit/instrumentation/crawl4ai/crawl4ai.py +2 -2
  32. openlit/instrumentation/crewai/__init__.py +3 -3
  33. openlit/instrumentation/crewai/crewai.py +2 -2
  34. openlit/instrumentation/dynamiq/__init__.py +5 -5
  35. openlit/instrumentation/dynamiq/dynamiq.py +2 -2
  36. openlit/instrumentation/elevenlabs/__init__.py +5 -5
  37. openlit/instrumentation/elevenlabs/async_elevenlabs.py +3 -3
  38. openlit/instrumentation/elevenlabs/elevenlabs.py +3 -3
  39. openlit/instrumentation/embedchain/__init__.py +2 -2
  40. openlit/instrumentation/embedchain/embedchain.py +4 -4
  41. openlit/instrumentation/firecrawl/__init__.py +3 -3
  42. openlit/instrumentation/firecrawl/firecrawl.py +2 -2
  43. openlit/instrumentation/google_ai_studio/__init__.py +3 -3
  44. openlit/instrumentation/google_ai_studio/async_google_ai_studio.py +3 -3
  45. openlit/instrumentation/google_ai_studio/google_ai_studio.py +3 -3
  46. openlit/instrumentation/gpt4all/__init__.py +3 -3
  47. openlit/instrumentation/gpt4all/gpt4all.py +7 -7
  48. openlit/instrumentation/groq/__init__.py +3 -3
  49. openlit/instrumentation/groq/async_groq.py +5 -5
  50. openlit/instrumentation/groq/groq.py +5 -5
  51. openlit/instrumentation/haystack/__init__.py +2 -2
  52. openlit/instrumentation/haystack/haystack.py +2 -2
  53. openlit/instrumentation/julep/__init__.py +7 -7
  54. openlit/instrumentation/julep/async_julep.py +3 -3
  55. openlit/instrumentation/julep/julep.py +3 -3
  56. openlit/instrumentation/langchain/__init__.py +2 -2
  57. openlit/instrumentation/langchain/async_langchain.py +13 -9
  58. openlit/instrumentation/langchain/langchain.py +13 -8
  59. openlit/instrumentation/letta/__init__.py +7 -7
  60. openlit/instrumentation/letta/letta.py +5 -5
  61. openlit/instrumentation/litellm/__init__.py +5 -5
  62. openlit/instrumentation/litellm/async_litellm.py +8 -8
  63. openlit/instrumentation/litellm/litellm.py +8 -8
  64. openlit/instrumentation/llamaindex/__init__.py +2 -2
  65. openlit/instrumentation/llamaindex/llamaindex.py +2 -2
  66. openlit/instrumentation/mem0/__init__.py +2 -2
  67. openlit/instrumentation/mem0/mem0.py +2 -2
  68. openlit/instrumentation/milvus/__init__.py +2 -2
  69. openlit/instrumentation/milvus/milvus.py +2 -2
  70. openlit/instrumentation/mistral/__init__.py +7 -7
  71. openlit/instrumentation/mistral/async_mistral.py +10 -10
  72. openlit/instrumentation/mistral/mistral.py +10 -10
  73. openlit/instrumentation/multion/__init__.py +7 -7
  74. openlit/instrumentation/multion/async_multion.py +5 -5
  75. openlit/instrumentation/multion/multion.py +5 -5
  76. openlit/instrumentation/ollama/__init__.py +11 -9
  77. openlit/instrumentation/ollama/async_ollama.py +71 -465
  78. openlit/instrumentation/ollama/ollama.py +71 -465
  79. openlit/instrumentation/ollama/utils.py +333 -0
  80. openlit/instrumentation/openai/__init__.py +11 -11
  81. openlit/instrumentation/openai/async_openai.py +18 -18
  82. openlit/instrumentation/openai/openai.py +18 -18
  83. openlit/instrumentation/phidata/__init__.py +2 -2
  84. openlit/instrumentation/phidata/phidata.py +2 -2
  85. openlit/instrumentation/pinecone/__init__.py +6 -6
  86. openlit/instrumentation/pinecone/pinecone.py +2 -2
  87. openlit/instrumentation/premai/__init__.py +3 -3
  88. openlit/instrumentation/premai/premai.py +7 -7
  89. openlit/instrumentation/qdrant/__init__.py +2 -2
  90. openlit/instrumentation/qdrant/async_qdrant.py +2 -2
  91. openlit/instrumentation/qdrant/qdrant.py +2 -2
  92. openlit/instrumentation/reka/__init__.py +3 -3
  93. openlit/instrumentation/reka/async_reka.py +3 -3
  94. openlit/instrumentation/reka/reka.py +3 -3
  95. openlit/instrumentation/together/__init__.py +5 -5
  96. openlit/instrumentation/together/async_together.py +8 -8
  97. openlit/instrumentation/together/together.py +8 -8
  98. openlit/instrumentation/transformers/__init__.py +2 -2
  99. openlit/instrumentation/transformers/transformers.py +4 -4
  100. openlit/instrumentation/vertexai/__init__.py +9 -9
  101. openlit/instrumentation/vertexai/async_vertexai.py +4 -4
  102. openlit/instrumentation/vertexai/vertexai.py +4 -4
  103. openlit/instrumentation/vllm/__init__.py +2 -2
  104. openlit/instrumentation/vllm/vllm.py +3 -3
  105. openlit/otel/events.py +85 -0
  106. openlit/otel/tracing.py +3 -13
  107. openlit/semcov/__init__.py +13 -1
  108. {openlit-1.33.10.dist-info → openlit-1.33.11.dist-info}/METADATA +2 -2
  109. openlit-1.33.11.dist-info/RECORD +125 -0
  110. openlit-1.33.10.dist-info/RECORD +0 -122
  111. {openlit-1.33.10.dist-info → openlit-1.33.11.dist-info}/LICENSE +0 -0
  112. {openlit-1.33.10.dist-info → openlit-1.33.11.dist-info}/WHEEL +0 -0
openlit/__helpers.py CHANGED
@@ -11,6 +11,7 @@ import requests
11
11
  import tiktoken
12
12
  from opentelemetry.sdk.resources import SERVICE_NAME, TELEMETRY_SDK_NAME, DEPLOYMENT_ENVIRONMENT
13
13
  from opentelemetry.trace import Status, StatusCode
14
+ from opentelemetry._events import Event
14
15
  from openlit.semcov import SemanticConvetion
15
16
 
16
17
  # Set up logging
@@ -285,3 +286,75 @@ def set_server_address_and_port(client_instance: Any,
285
286
  server_port = default_server_port
286
287
 
287
288
  return server_address, server_port
289
+
290
+ def otel_event(name, attributes, body):
291
+ """
292
+ Returns an OpenTelemetry Event object
293
+ """
294
+
295
+ return Event(
296
+ name=name,
297
+ attributes=attributes,
298
+ body=body,
299
+ )
300
+
301
+ def extract_and_format_input(messages):
302
+ """
303
+ Process a list of messages to extract content and categorize
304
+ them into fixed roles like 'user', 'assistant', 'system'.
305
+ """
306
+
307
+ fixed_roles = ['user', 'assistant', 'system', 'tool'] # Ensure these are your fixed keys
308
+ # Initialize the dictionary with fixed keys and empty structures
309
+ formatted_messages = {role_key: {"role": "", "content": ""} for role_key in fixed_roles}
310
+
311
+ for message in messages:
312
+ # Normalize the message structure
313
+ message = response_as_dict(message)
314
+
315
+ # Extract role and content
316
+ role = message.get("role")
317
+ if role not in fixed_roles:
318
+ continue # Skip any role not in our predefined roles
319
+
320
+ content = message.get("content", "")
321
+
322
+ # Prepare content as a string
323
+ if isinstance(content, list):
324
+ content_str = ", ".join(
325
+ # pylint: disable=line-too-long
326
+ 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", ""))}'
327
+ for item in content
328
+ )
329
+ else:
330
+ content_str = content
331
+
332
+ # Set the role in the formatted message and concatenate content
333
+ if not formatted_messages[role]["role"]:
334
+ formatted_messages[role]["role"] = role
335
+
336
+ if formatted_messages[role]["content"]:
337
+ formatted_messages[role]["content"] += " " + content_str
338
+ else:
339
+ formatted_messages[role]["content"] = content_str
340
+
341
+ return formatted_messages
342
+
343
+ # To be removed one the change to log events (from span events) is complete
344
+ def concatenate_all_contents(formatted_messages):
345
+ """
346
+ Concatenate all 'content' fields from the formatted messages
347
+ dictionary into a single string.
348
+
349
+ Parameters:
350
+ - formatted_messages: Dictionary with roles as keys and corresponding
351
+ role and content as values.
352
+
353
+ Returns:
354
+ - A single string with all content concatenated.
355
+ """
356
+ return " ".join(
357
+ message_data['content']
358
+ for message_data in formatted_messages.values()
359
+ if message_data['content']
360
+ )
openlit/__init__.py CHANGED
@@ -21,9 +21,9 @@ from opentelemetry.sdk.resources import SERVICE_NAME, DEPLOYMENT_ENVIRONMENT
21
21
  from openlit.semcov import SemanticConvetion
22
22
  from openlit.otel.tracing import setup_tracing
23
23
  from openlit.otel.metrics import setup_meter
24
+ from openlit.otel.events import setup_events
24
25
  from openlit.__helpers import fetch_pricing_info, get_env_variable
25
26
 
26
-
27
27
  # Instrumentors for various large language models.
28
28
  from openlit.instrumentation.openai import OpenAIInstrumentor
29
29
  from openlit.instrumentation.anthropic import AnthropicInstrumentor
@@ -85,10 +85,11 @@ class OpenlitConfig:
85
85
  application_name (str): Name of the application using openLIT.
86
86
  pricing_info (Dict[str, Any]): Pricing information.
87
87
  tracer (Optional[Any]): Tracer instance for OpenTelemetry.
88
+ event_provider (Optional[Any]): Event logger provider for OpenTelemetry.
88
89
  otlp_endpoint (Optional[str]): Endpoint for OTLP.
89
90
  otlp_headers (Optional[Dict[str, str]]): Headers for OTLP.
90
91
  disable_batch (bool): Flag to disable batch span processing in tracing.
91
- trace_content (bool): Flag to enable or disable tracing of content.
92
+ capture_message_content (bool): Flag to enable or disable tracing of content.
92
93
  """
93
94
 
94
95
  _instance = None
@@ -107,11 +108,12 @@ class OpenlitConfig:
107
108
  cls.application_name = "default"
108
109
  cls.pricing_info = {}
109
110
  cls.tracer = None
111
+ cls.event_provider = None
110
112
  cls.metrics_dict = {}
111
113
  cls.otlp_endpoint = None
112
114
  cls.otlp_headers = None
113
115
  cls.disable_batch = False
114
- cls.trace_content = True
116
+ cls.capture_message_content = True
115
117
  cls.disable_metrics = False
116
118
 
117
119
  @classmethod
@@ -120,10 +122,11 @@ class OpenlitConfig:
120
122
  environment,
121
123
  application_name,
122
124
  tracer,
125
+ event_provider,
123
126
  otlp_endpoint,
124
127
  otlp_headers,
125
128
  disable_batch,
126
- trace_content,
129
+ capture_message_content,
127
130
  metrics_dict,
128
131
  disable_metrics,
129
132
  pricing_json,
@@ -135,22 +138,26 @@ class OpenlitConfig:
135
138
  environment (str): Deployment environment.
136
139
  application_name (str): Application name.
137
140
  tracer: Tracer instance.
141
+ event_provider: Event logger provider instance.
138
142
  meter: Metric Instance
139
143
  otlp_endpoint (str): OTLP endpoint.
140
144
  otlp_headers (Dict[str, str]): OTLP headers.
141
145
  disable_batch (bool): Disable batch span processing flag.
142
- trace_content (bool): Enable or disable content tracing.
146
+ capture_message_content (bool): Enable or disable content tracing.
147
+ metrics_dict: Dictionary of metrics.
148
+ disable_metrics (bool): Flag to disable metrics.
143
149
  pricing_json(str): path or url to the pricing json file
144
150
  """
145
151
  cls.environment = environment
146
152
  cls.application_name = application_name
147
153
  cls.pricing_info = fetch_pricing_info(pricing_json)
148
154
  cls.tracer = tracer
155
+ cls.event_provider = event_provider
149
156
  cls.metrics_dict = metrics_dict
150
157
  cls.otlp_endpoint = otlp_endpoint
151
158
  cls.otlp_headers = otlp_headers
152
159
  cls.disable_batch = disable_batch
153
- cls.trace_content = trace_content
160
+ cls.capture_message_content = capture_message_content
154
161
  cls.disable_metrics = disable_metrics
155
162
 
156
163
 
@@ -187,8 +194,9 @@ def instrument_if_available(
187
194
  environment=config.environment,
188
195
  application_name=config.application_name,
189
196
  tracer=config.tracer,
197
+ event_provider=config.event_provider,
190
198
  pricing_info=config.pricing_info,
191
- trace_content=config.trace_content,
199
+ capture_message_content=config.capture_message_content,
192
200
  metrics_dict=config.metrics_dict,
193
201
  disable_metrics=config.disable_metrics,
194
202
  )
@@ -207,10 +215,11 @@ def init(
207
215
  environment="default",
208
216
  application_name="default",
209
217
  tracer=None,
218
+ event_logger=None,
210
219
  otlp_endpoint=None,
211
220
  otlp_headers=None,
212
221
  disable_batch=False,
213
- trace_content=True,
222
+ capture_message_content=True,
214
223
  disabled_instrumentors=None,
215
224
  meter=None,
216
225
  disable_metrics=False,
@@ -227,11 +236,12 @@ def init(
227
236
  environment (str): Deployment environment.
228
237
  application_name (str): Application name.
229
238
  tracer: Tracer instance (Optional).
239
+ event_logger: EventLoggerProvider instance (Optional).
230
240
  meter: OpenTelemetry Metrics Instance (Optional).
231
241
  otlp_endpoint (str): OTLP endpoint for exporter (Optional).
232
242
  otlp_headers (Dict[str, str]): OTLP headers for exporter (Optional).
233
243
  disable_batch (bool): Flag to disable batch span processing (Optional).
234
- trace_content (bool): Flag to trace content (Optional).
244
+ capture_message_content (bool): Flag to trace content (Optional).
235
245
  disabled_instrumentors (List[str]): Optional. List of instrumentor names to disable.
236
246
  disable_metrics (bool): Flag to disable metrics (Optional).
237
247
  pricing_json(str): File path or url to the pricing json (Optional).
@@ -308,9 +318,22 @@ def init(
308
318
  )
309
319
 
310
320
  if not tracer:
311
- logger.error("openLIT tracing setup failed. Tracing will not be available.")
321
+ logger.error("OpenLIT tracing setup failed. Tracing will not be available.")
312
322
  return
313
323
 
324
+ # Setup events based on the provided or default configuration.
325
+ event_provider = setup_events(
326
+ application_name=application_name,
327
+ environment=environment,
328
+ event_logger=event_logger,
329
+ otlp_endpoint=None,
330
+ otlp_headers=None,
331
+ disable_batch=disable_batch,
332
+ )
333
+
334
+ if not event_provider:
335
+ logger.error("OpenLIT events setup failed. Events will not be available")
336
+
314
337
  # Setup meter and receive metrics_dict instead of meter.
315
338
  metrics_dict, err = setup_meter(
316
339
  application_name=application_name,
@@ -326,15 +349,19 @@ def init(
326
349
  )
327
350
  return
328
351
 
352
+ if os.getenv("OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT", "").lower == "false":
353
+ capture_message_content=False
354
+
329
355
  # Update global configuration with the provided settings.
330
356
  config.update_config(
331
357
  environment,
332
358
  application_name,
333
359
  tracer,
360
+ event_provider,
334
361
  otlp_endpoint,
335
362
  otlp_headers,
336
363
  disable_batch,
337
- trace_content,
364
+ capture_message_content,
338
365
  metrics_dict,
339
366
  disable_metrics,
340
367
  pricing_json,
@@ -1,4 +1,3 @@
1
- # pylint: disable=useless-return, bad-staticmethod-argument, disable=duplicate-code
2
1
  """Initializer of Auto Instrumentation of AG2 Functions"""
3
2
 
4
3
  from typing import Collection
@@ -7,7 +6,7 @@ from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
7
6
  from wrapt import wrap_function_wrapper
8
7
 
9
8
  from openlit.instrumentation.ag2.ag2 import (
10
- wrap_ag2
9
+ conversable_agent, agent_run
11
10
  )
12
11
 
13
12
  _instruments = ("ag2 >= 0.3.2",)
@@ -24,27 +23,27 @@ class AG2Instrumentor(BaseInstrumentor):
24
23
  application_name = kwargs.get("application_name", "default_application")
25
24
  environment = kwargs.get("environment", "default_environment")
26
25
  tracer = kwargs.get("tracer")
26
+ event_provider = kwargs.get("event_provider")
27
27
  metrics = kwargs.get("metrics_dict")
28
28
  pricing_info = kwargs.get("pricing_info", {})
29
- trace_content = kwargs.get("trace_content", False)
29
+ capture_message_content = kwargs.get("capture_message_content", False)
30
30
  disable_metrics = kwargs.get("disable_metrics")
31
31
  version = importlib.metadata.version("ag2")
32
32
 
33
33
  wrap_function_wrapper(
34
34
  "autogen.agentchat.conversable_agent",
35
- "ConversableAgent.initiate_chat",
36
- wrap_ag2("ag2.initiate_chat", version, environment, application_name,
37
- tracer, pricing_info, trace_content, metrics, disable_metrics),
35
+ "ConversableAgent.__init__",
36
+ conversable_agent(version, environment, application_name,
37
+ tracer, event_provider, pricing_info, capture_message_content, metrics, disable_metrics),
38
38
  )
39
39
 
40
40
  wrap_function_wrapper(
41
41
  "autogen.agentchat.conversable_agent",
42
- "ConversableAgent.generate_reply",
43
- wrap_ag2("ag2.generate_reply", version, environment, application_name,
44
- tracer, pricing_info, trace_content, metrics, disable_metrics),
42
+ "ConversableAgent.run",
43
+ agent_run(version, environment, application_name,
44
+ tracer, event_provider, pricing_info, capture_message_content, metrics, disable_metrics),
45
45
  )
46
46
 
47
-
48
47
  def _uninstrument(self, **kwargs):
49
48
  # Proper uninstrumentation logic to revert patched methods
50
49
  pass
@@ -1,98 +1,163 @@
1
- # pylint: disable=duplicate-code, broad-exception-caught, too-many-statements, unused-argument
2
1
  """
3
- Module for monitoring AG2.
2
+ Module for monitoring AG2 API calls.
4
3
  """
5
4
 
6
5
  import logging
6
+ import time
7
7
  from opentelemetry.trace import SpanKind, Status, StatusCode
8
8
  from opentelemetry.sdk.resources import SERVICE_NAME, TELEMETRY_SDK_NAME, DEPLOYMENT_ENVIRONMENT
9
- from openlit.__helpers import handle_exception
9
+ from openlit.__helpers import (
10
+ handle_exception,
11
+ get_chat_model_cost,
12
+ otel_event,
13
+ )
10
14
  from openlit.semcov import SemanticConvetion
11
15
 
12
16
  # Initialize logger for logging potential issues and operations
13
17
  logger = logging.getLogger(__name__)
14
18
 
15
- def wrap_ag2(gen_ai_endpoint, version, environment, application_name,
16
- tracer, pricing_info, trace_content, metrics, disable_metrics):
19
+ AGENT_NAME = ''
20
+ REQUEST_MODEL = ''
21
+ SYSTEM_MESSAGE = ''
22
+ MODEL_AND_NAME_SET = False
23
+
24
+ def set_span_attributes(span, version, operation_name, environment,
25
+ application_name, server_address, server_port, request_model):
26
+ """
27
+ Set common attributes for the span.
28
+ """
29
+
30
+ # Set Span attributes (OTel Semconv)
31
+ span.set_attribute(TELEMETRY_SDK_NAME, "openlit")
32
+ span.set_attribute(SemanticConvetion.GEN_AI_OPERATION, operation_name)
33
+ span.set_attribute(SemanticConvetion.GEN_AI_SYSTEM, SemanticConvetion.GEN_AI_SYSTEM_AG2)
34
+ span.set_attribute(SemanticConvetion.GEN_AI_AGENT_NAME, AGENT_NAME)
35
+ span.set_attribute(SemanticConvetion.SERVER_ADDRESS, server_address)
36
+ span.set_attribute(SemanticConvetion.SERVER_PORT, server_port)
37
+ span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_MODEL, request_model)
38
+
39
+ # Set Span attributes (Extras)
40
+ span.set_attribute(DEPLOYMENT_ENVIRONMENT, environment)
41
+ span.set_attribute(SERVICE_NAME, application_name)
42
+ span.set_attribute(SemanticConvetion.GEN_AI_SDK_VERSION, version)
43
+
44
+ def calculate_tokens_and_cost(response, request_model, pricing_info):
17
45
  """
18
- Creates a wrapper around a function call to trace and log its execution metrics.
19
-
20
- This function wraps any given function to measure its execution time,
21
- log its operation, and trace its execution using OpenTelemetry.
22
-
23
- Parameters:
24
- - gen_ai_endpoint (str): A descriptor or name for the endpoint being traced.
25
- - version (str): The version of the Langchain application.
26
- - environment (str): The deployment environment (e.g., 'production', 'development').
27
- - application_name (str): Name of the Langchain application.
28
- - tracer (opentelemetry.trace.Tracer): The tracer object used for OpenTelemetry tracing.
29
- - pricing_info (dict): Information about the pricing for internal metrics (currently not used).
30
- - trace_content (bool): Flag indicating whether to trace the content of the response.
31
-
32
- Returns:
33
- - function: A higher-order function that takes a function 'wrapped' and returns
34
- a new function that wraps 'wrapped' with additional tracing and logging.
46
+ Calculate the input, output tokens, and their respective costs.
35
47
  """
48
+ input_tokens = 0
49
+ output_tokens = 0
36
50
 
51
+ for usage_data in response.cost.values():
52
+ if isinstance(usage_data, dict):
53
+ for model_data in usage_data.values():
54
+ if isinstance(model_data, dict):
55
+ input_tokens += model_data.get('prompt_tokens', 0)
56
+ output_tokens += model_data.get('completion_tokens', 0)
57
+
58
+ cost = get_chat_model_cost(request_model, pricing_info, input_tokens, output_tokens)
59
+ return input_tokens, output_tokens, cost
60
+
61
+ def emit_events(response, event_provider, capture_message_content):
62
+ """
63
+ Emit OpenTelemetry events for each chat history entry.
64
+ """
65
+ for chat in response.chat_history:
66
+ event_type = (
67
+ SemanticConvetion.GEN_AI_CHOICE if chat['role'] == 'user'
68
+ else SemanticConvetion.GEN_AI_USER_MESSAGE
69
+ )
70
+ choice_event = otel_event(
71
+ name=event_type,
72
+ attributes={
73
+ SemanticConvetion.GEN_AI_SYSTEM: SemanticConvetion.GEN_AI_SYSTEM_AG2
74
+ },
75
+ body={
76
+ "index": response.chat_history.index(chat),
77
+ "message": {
78
+ **({"content": chat['content']} if capture_message_content else {}),
79
+ "role": 'assistant' if chat['role'] == 'user' else 'user'
80
+ }
81
+ }
82
+ )
83
+ event_provider.emit(choice_event)
84
+
85
+ def conversable_agent(version, environment, application_name,
86
+ tracer, event_provider, pricing_info, capture_message_content, metrics, disable_metrics):
87
+ """
88
+ Generates a telemetry wrapper for GenAI function call
89
+ """
37
90
  def wrapper(wrapped, instance, args, kwargs):
38
- """
39
- An inner wrapper function that executes the wrapped function, measures execution
40
- time, and records trace data using OpenTelemetry.
41
-
42
- Parameters:
43
- - wrapped (Callable): The original function that this wrapper will execute.
44
- - instance (object): The instance to which the wrapped function belongs. This
45
- is used for instance methods. For static and classmethods,
46
- this may be None.
47
- - args (tuple): Positional arguments passed to the wrapped function.
48
- - kwargs (dict): Keyword arguments passed to the wrapped function.
49
-
50
- Returns:
51
- - The result of the wrapped function call.
52
-
53
- The wrapper initiates a span with the provided tracer, sets various attributes
54
- on the span based on the function's execution and response, and ensures
55
- errors are handled and logged appropriately.
56
- """
57
-
58
- with tracer.start_as_current_span(gen_ai_endpoint, kind= SpanKind.CLIENT) as span:
59
- response = wrapped(*args, **kwargs)
60
-
61
- if isinstance(instance.__dict__.get('llm_config'), dict):
62
- llm_model = instance.__dict__['llm_config'].get('model', 'gpt-4')
63
- else:
64
- # Fallback to default if 'llm_config' is not a dictionary
65
- llm_model = None
91
+ server_address, server_port = '127.0.0.1', 80
92
+ global AGENT_NAME, MODEL_AND_NAME_SET, REQUEST_MODEL, SYSTEM_MESSAGE
93
+
94
+ if not MODEL_AND_NAME_SET:
95
+ AGENT_NAME = kwargs.get("name", "NOT_FOUND")
96
+ REQUEST_MODEL = kwargs.get("llm_config", {}).get('model', 'gpt-4o')
97
+ SYSTEM_MESSAGE = kwargs.get('system_message', '')
98
+ MODEL_AND_NAME_SET = True
66
99
 
100
+ span_name = f"{SemanticConvetion.GEN_AI_OPERATION_TYPE_CREATE_AGENT} {AGENT_NAME}"
101
+
102
+ with tracer.start_as_current_span(span_name, kind=SpanKind.CLIENT) as span:
67
103
  try:
68
- span.set_attribute(TELEMETRY_SDK_NAME, "openlit")
69
- span.set_attribute(SemanticConvetion.GEN_AI_ENDPOINT,
70
- gen_ai_endpoint)
71
- span.set_attribute(SemanticConvetion.GEN_AI_SYSTEM,
72
- SemanticConvetion.GEN_AI_SYSTEM_AG2)
73
- span.set_attribute(DEPLOYMENT_ENVIRONMENT,
74
- environment)
75
- span.set_attribute(SERVICE_NAME,
76
- application_name)
77
- span.set_attribute(SemanticConvetion.GEN_AI_OPERATION,
78
- SemanticConvetion.GEN_AI_OPERATION_TYPE_AGENT)
79
- span.set_attribute(SemanticConvetion.GEN_AI_AGENT_ROLE,
80
- instance.name)
81
- if llm_model:
82
- span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_MODEL,
83
- llm_model)
104
+ start_time = time.time()
105
+ response = wrapped(*args, **kwargs)
106
+ end_time = time.time()
84
107
 
108
+ set_span_attributes(span, version, SemanticConvetion.GEN_AI_OPERATION_TYPE_CREATE_AGENT,
109
+ environment, application_name, server_address, server_port, REQUEST_MODEL)
110
+ span.set_attribute(SemanticConvetion.GEN_AI_AGENT_DESCRIPTION, SYSTEM_MESSAGE)
111
+ span.set_attribute(SemanticConvetion.GEN_AI_RESPONSE_MODEL, REQUEST_MODEL)
112
+ span.set_attribute(SemanticConvetion.GEN_AI_SERVER_TTFT, end_time - start_time)
85
113
 
86
114
  span.set_status(Status(StatusCode.OK))
87
115
 
88
- # Return original response
89
116
  return response
90
117
 
91
118
  except Exception as e:
92
119
  handle_exception(span, e)
93
120
  logger.error("Error in trace creation: %s", e)
121
+ return response
122
+
123
+ return wrapper
124
+
125
+ def agent_run(version, environment, application_name,
126
+ tracer, event_provider, pricing_info, capture_message_content, metrics, disable_metrics):
127
+ """
128
+ Generates a telemetry wrapper for GenAI function call
129
+ """
130
+ def wrapper(wrapped, instance, args, kwargs):
131
+ server_address, server_port = '127.0.0.1', 80
132
+
133
+ span_name = f"{SemanticConvetion.GEN_AI_OPERATION_TYPE_EXECUTE_AGENT_TASK} {AGENT_NAME}"
134
+
135
+ with tracer.start_as_current_span(span_name, kind=SpanKind.CLIENT) as span:
136
+ try:
137
+ start_time = time.time()
138
+ response = wrapped(*args, **kwargs)
139
+ end_time = time.time()
140
+
141
+ input_tokens, output_tokens, cost = calculate_tokens_and_cost(response, REQUEST_MODEL, pricing_info)
142
+ response_model = list(response.cost.get('usage_including_cached_inference', {}).keys())[1]
143
+
144
+ set_span_attributes(span, version, SemanticConvetion.GEN_AI_OPERATION_TYPE_EXECUTE_AGENT_TASK,
145
+ environment, application_name, server_address, server_port, REQUEST_MODEL)
146
+ span.set_attribute(SemanticConvetion.GEN_AI_RESPONSE_MODEL, response_model)
147
+ span.set_attribute(SemanticConvetion.GEN_AI_USAGE_INPUT_TOKENS, input_tokens)
148
+ span.set_attribute(SemanticConvetion.GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens)
149
+ span.set_attribute(SemanticConvetion.GEN_AI_CLIENT_TOKEN_USAGE, input_tokens + output_tokens)
150
+ span.set_attribute(SemanticConvetion.GEN_AI_USAGE_COST, cost)
151
+ span.set_attribute(SemanticConvetion.GEN_AI_SERVER_TTFT, end_time - start_time)
152
+
153
+ emit_events(response, event_provider, capture_message_content)
154
+ span.set_status(Status(StatusCode.OK))
155
+
156
+ return response
94
157
 
95
- # Return original response
158
+ except Exception as e:
159
+ handle_exception(span, e)
160
+ logger.error("Error in trace creation: %s", e)
96
161
  return response
97
162
 
98
163
  return wrapper
@@ -27,9 +27,10 @@ class AI21Instrumentor(BaseInstrumentor):
27
27
  application_name = kwargs.get("application_name", "default_application")
28
28
  environment = kwargs.get("environment", "default_environment")
29
29
  tracer = kwargs.get("tracer")
30
+ event_provider = kwargs.get("event_provider")
30
31
  metrics = kwargs.get("metrics_dict")
31
32
  pricing_info = kwargs.get("pricing_info", {})
32
- trace_content = kwargs.get("trace_content", False)
33
+ capture_message_content = kwargs.get("capture_message_content", False)
33
34
  disable_metrics = kwargs.get("disable_metrics")
34
35
  version = importlib.metadata.version("ai21")
35
36
 
@@ -38,13 +39,13 @@ class AI21Instrumentor(BaseInstrumentor):
38
39
  "ai21.clients.studio.resources.chat.chat_completions",
39
40
  "ChatCompletions.create",
40
41
  chat(version, environment, application_name,
41
- tracer, pricing_info, trace_content, metrics, disable_metrics),
42
+ tracer, event_provider, pricing_info, capture_message_content, metrics, disable_metrics),
42
43
  )
43
44
  wrap_function_wrapper(
44
45
  "ai21.clients.studio.resources.studio_conversational_rag",
45
46
  "StudioConversationalRag.create",
46
47
  chat_rag(version, environment, application_name,
47
- tracer, pricing_info, trace_content, metrics, disable_metrics),
48
+ tracer, event_provider, pricing_info, capture_message_content, metrics, disable_metrics),
48
49
  )
49
50
 
50
51
  #Async
@@ -52,13 +53,13 @@ class AI21Instrumentor(BaseInstrumentor):
52
53
  "ai21.clients.studio.resources.chat.async_chat_completions",
53
54
  "AsyncChatCompletions.create",
54
55
  async_chat(version, environment, application_name,
55
- tracer, pricing_info, trace_content, metrics, disable_metrics),
56
+ tracer, event_provider, pricing_info, capture_message_content, metrics, disable_metrics),
56
57
  )
57
58
  wrap_function_wrapper(
58
59
  "ai21.clients.studio.resources.studio_conversational_rag",
59
60
  "AsyncStudioConversationalRag.create",
60
61
  async_chat_rag(version, environment, application_name,
61
- tracer, pricing_info, trace_content, metrics, disable_metrics),
62
+ tracer, event_provider, pricing_info, capture_message_content, metrics, disable_metrics),
62
63
  )
63
64
 
64
65
  def _uninstrument(self, **kwargs):