openlit 1.34.30__py3-none-any.whl → 1.34.31__py3-none-any.whl

openlit/guard/utils.py

@@ -16,6 +16,7 @@ from openlit.semcov import SemanticConvention
 # Initialize logger for logging potential issues and operations
 logger = logging.getLogger(__name__)
 
+
 class JsonOutput(BaseModel):
     """
     A model representing the structure of JSON output for prompt injection detection.
@@ -34,13 +35,17 @@ class JsonOutput(BaseModel):
     classification: str
     explanation: str
 
-def setup_provider(provider: Optional[str], api_key: Optional[str],
-                   model: Optional[str],
-                   base_url: Optional[str]) -> Tuple[Optional[str], Optional[str], Optional[str]]:
+
+def setup_provider(
+    provider: Optional[str],
+    api_key: Optional[str],
+    model: Optional[str],
+    base_url: Optional[str],
+) -> Tuple[Optional[str], Optional[str], Optional[str]]:
     """Function to setup LLM provider"""
     provider_configs = {
         "openai": {"env_var": "OPENAI_API_KEY"},
-        "anthropic": {"env_var": "ANTHROPIC_API_KEY"}
+        "anthropic": {"env_var": "ANTHROPIC_API_KEY"},
     }
 
     if provider is None:
@@ -60,7 +65,9 @@ def setup_provider(provider: Optional[str], api_key: Optional[str],
 
     if not api_key:
         # pylint: disable=line-too-long
-        raise ValueError(f"API key required via 'api_key' parameter or '{env_var}' environment variable")
+        raise ValueError(
+            f"API key required via 'api_key' parameter or '{env_var}' environment variable"
+        )
 
     return api_key, model, base_url
 
@@ -69,6 +76,7 @@ def format_prompt(system_prompt: str, text: str) -> str:
     """Function to format the prompt"""
     return system_prompt.replace("{{prompt}}", text)
 
+
 def llm_response(provider: str, prompt: str, model: str, base_url: str) -> str:
     """Function to get LLM response based on provider"""
     # pylint: disable=no-else-return
@@ -79,6 +87,7 @@ def llm_response(provider: str, prompt: str, model: str, base_url: str) -> str:
     else:
         raise ValueError(f"Unsupported provider: {provider}")
 
+
 def llm_response_openai(prompt: str, model: str, base_url: str) -> str:
     """Function to make LLM call to OpenAI"""
     client = OpenAI(base_url=base_url)
@@ -95,10 +104,11 @@ def llm_response_openai(prompt: str, model: str, base_url: str) -> str:
             {"role": "user", "content": prompt},
         ],
         temperature=0.0,
-        response_format=JsonOutput
+        response_format=JsonOutput,
     )
     return response.choices[0].message.content
 
+
 def llm_response_anthropic(prompt: str, model: str) -> str:
     """Function to make LLM call to Anthropic"""
     client = Anthropic()
@@ -113,26 +123,42 @@ def llm_response_anthropic(prompt: str, model: str) -> str:
             "input_schema": {
                 "type": "object",
                 "properties": {
-                    "verdict": {"type": "string", "description": "Verdict of guardrail"},
+                    "verdict": {
+                        "type": "string",
+                        "description": "Verdict of guardrail",
+                    },
                     "guard": {"type": "string", "description": "Type of guard"},
-                    "score": {"type": "number", "description": "Prompt score from Guard."},
-                    "classification": {"type": "string", "description": "Incorrect prompt type"},
-                    "explanation": {"type": "string", "description": "Reason for classification"}
+                    "score": {
+                        "type": "number",
+                        "description": "Prompt score from Guard.",
+                    },
+                    "classification": {
+                        "type": "string",
+                        "description": "Incorrect prompt type",
+                    },
+                    "explanation": {
+                        "type": "string",
+                        "description": "Reason for classification",
+                    },
                 },
-                "required": ["verdict", "guard", "score", "classification", "explanation"]
-            }
+                "required": [
+                    "verdict",
+                    "guard",
+                    "score",
+                    "classification",
+                    "explanation",
+                ],
+            },
         }
     ]
 
     response = client.messages.create(
         model=model,
-        messages=[
-            {"role": "user", "content": prompt}
-        ],
+        messages=[{"role": "user", "content": prompt}],
         max_tokens=2000,
         temperature=0.0,
         tools=tools,
-        stream=False
+        stream=False,
     )
 
     for content in response.content:
@@ -142,6 +168,7 @@ def llm_response_anthropic(prompt: str, model: str) -> str:
 
     return response
 
+
 def parse_llm_response(response) -> JsonOutput:
     """
     Parses the LLM response into a JsonOutput object.
@@ -163,8 +190,14 @@ def parse_llm_response(response) -> JsonOutput:
         return JsonOutput(**data)
     except (json.JSONDecodeError, TypeError) as e:
         logger.error("Error parsing LLM response: '%s'", e)
-        return JsonOutput(score=0, classification="none", explanation="none",
-                          verdict="none", guard="none")
+        return JsonOutput(
+            score=0,
+            classification="none",
+            explanation="none",
+            verdict="none",
+            guard="none",
+        )
+
 
 def custom_rule_detection(text: str, custom_rules: list) -> JsonOutput:
     """
@@ -183,10 +216,12 @@ def custom_rule_detection(text: str, custom_rules: list) -> JsonOutput:
                 guard=rule.get("guard", "prompt_injection"),
                 score=rule.get("score", 0.5),
                 classification=rule.get("classification", "custom"),
-                explanation=rule.get("explanation", "Matched custom rule pattern.")
+                explanation=rule.get("explanation", "Matched custom rule pattern."),
             )
-    return JsonOutput(score=0, classification="none", explanation="none",
-                      verdict="none", guard="none")
+    return JsonOutput(
+        score=0, classification="none", explanation="none", verdict="none", guard="none"
+    )
+
 
 def guard_metrics():
     """
@@ -204,11 +239,12 @@ def guard_metrics():
     guard_requests = meter.create_counter(
         name=SemanticConvention.GUARD_REQUESTS,
         description="Counter for Guard requests",
-        unit="1"
+        unit="1",
     )
 
     return guard_requests
 
+
 def guard_metric_attributes(verdict, score, validator, classification, explanation):
     """
     Initializes OpenTelemetry attributes for metrics.

openlit/instrumentation/ag2/__init__.py

@@ -5,15 +5,15 @@ import importlib.metadata
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from wrapt import wrap_function_wrapper
 
-from openlit.instrumentation.ag2.ag2 import (
-    conversable_agent, agent_run
-)
+from openlit.instrumentation.ag2.ag2 import conversable_agent, agent_run
 from openlit.instrumentation.ag2.async_ag2 import (
-    async_conversable_agent, async_agent_run
+    async_conversable_agent,
+    async_agent_run,
 )
 
 _instruments = ("ag2 >= 0.3.2",)
 
+
 class AG2Instrumentor(BaseInstrumentor):
     """
     An instrumentor for AG2 client library.
@@ -36,16 +36,32 @@ class AG2Instrumentor(BaseInstrumentor):
         wrap_function_wrapper(
             "autogen.agentchat.conversable_agent",
             "ConversableAgent.__init__",
-            conversable_agent(version, environment, application_name,
-                tracer, pricing_info, capture_message_content, metrics, disable_metrics),
+            conversable_agent(
+                version,
+                environment,
+                application_name,
+                tracer,
+                pricing_info,
+                capture_message_content,
+                metrics,
+                disable_metrics,
+            ),
         )
 
         # sync agent run
         wrap_function_wrapper(
             "autogen.agentchat.conversable_agent",
             "ConversableAgent.run",
-            agent_run(version, environment, application_name,
-                tracer, pricing_info, capture_message_content, metrics, disable_metrics),
+            agent_run(
+                version,
+                environment,
+                application_name,
+                tracer,
+                pricing_info,
+                capture_message_content,
+                metrics,
+                disable_metrics,
+            ),
         )
 
     def _uninstrument(self, **kwargs):

openlit/instrumentation/ag2/ag2.py

@@ -4,18 +4,24 @@ Module for monitoring AG2 API calls.
 
 import time
 from opentelemetry.trace import SpanKind
-from openlit.__helpers import (
-    handle_exception,
-    set_server_address_and_port
-)
+from openlit.__helpers import handle_exception, set_server_address_and_port
 from openlit.instrumentation.ag2.utils import (
     process_agent_creation,
     process_agent_run,
 )
 from openlit.semcov import SemanticConvention
 
-def conversable_agent(version, environment, application_name, tracer, pricing_info,
-    capture_message_content, metrics, disable_metrics):
+
+def conversable_agent(
+    version,
+    environment,
+    application_name,
+    tracer,
+    pricing_info,
+    capture_message_content,
+    metrics,
+    disable_metrics,
+):
     """
     Generates a telemetry wrapper for AG2 conversable agent creation.
     """
@@ -25,12 +31,16 @@ def conversable_agent(version, environment, application_name, tracer, pricing_in
         Wraps the AG2 conversable agent creation call.
         """
 
-        server_address, server_port = set_server_address_and_port(instance, "127.0.0.1", 80)
+        server_address, server_port = set_server_address_and_port(
+            instance, "127.0.0.1", 80
+        )
         agent_name = kwargs.get("name", "NOT_FOUND")
         llm_config = kwargs.get("llm_config", {})
         system_message = kwargs.get("system_message", "")
 
-        span_name = f"{SemanticConvention.GEN_AI_OPERATION_TYPE_CREATE_AGENT} {agent_name}"
+        span_name = (
+            f"{SemanticConvention.GEN_AI_OPERATION_TYPE_CREATE_AGENT} {agent_name}"
+        )
 
         with tracer.start_as_current_span(span_name, kind=SpanKind.CLIENT) as span:
             start_time = time.time()
@@ -51,7 +61,7 @@ def conversable_agent(version, environment, application_name, tracer, pricing_in
                     span=span,
                     capture_message_content=capture_message_content,
                     disable_metrics=disable_metrics,
-                    version=version
+                    version=version,
                 )
 
             except Exception as e:
@@ -61,8 +71,17 @@ def conversable_agent(version, environment, application_name, tracer, pricing_in
 
     return wrapper
 
-def agent_run(version, environment, application_name, tracer, pricing_info,
-    capture_message_content, metrics, disable_metrics):
+
+def agent_run(
+    version,
+    environment,
+    application_name,
+    tracer,
+    pricing_info,
+    capture_message_content,
+    metrics,
+    disable_metrics,
+):
     """
     Generates a telemetry wrapper for AG2 agent run execution.
     """
@@ -72,7 +91,9 @@ def agent_run(version, environment, application_name, tracer, pricing_info,
         Wraps the AG2 agent run execution call.
         """
 
-        server_address, server_port = set_server_address_and_port(instance, "127.0.0.1", 80)
+        server_address, server_port = set_server_address_and_port(
+            instance, "127.0.0.1", 80
+        )
 
         # Extract agent name from instance
         agent_name = getattr(instance, "name", "NOT_FOUND")
@@ -103,7 +124,7 @@ def agent_run(version, environment, application_name, tracer, pricing_info,
                     span=span,
                     capture_message_content=capture_message_content,
                     disable_metrics=disable_metrics,
-                    version=version
+                    version=version,
                 )
 
             except Exception as e:

openlit/instrumentation/ag2/async_ag2.py

@@ -4,18 +4,24 @@ Module for monitoring AG2 API calls (async version).
 
 import time
 from opentelemetry.trace import SpanKind
-from openlit.__helpers import (
-    handle_exception,
-    set_server_address_and_port
-)
+from openlit.__helpers import handle_exception, set_server_address_and_port
 from openlit.instrumentation.ag2.utils import (
     process_agent_creation,
     process_agent_run,
 )
 from openlit.semcov import SemanticConvention
 
-def async_conversable_agent(version, environment, application_name, tracer, pricing_info,
-    capture_message_content, metrics, disable_metrics):
+
+def async_conversable_agent(
+    version,
+    environment,
+    application_name,
+    tracer,
+    pricing_info,
+    capture_message_content,
+    metrics,
+    disable_metrics,
+):
     """
     Generates a telemetry wrapper for AG2 async conversable agent creation.
     """
@@ -25,12 +31,16 @@ def async_conversable_agent(version, environment, application_name, tracer, pric
         Wraps the AG2 async conversable agent creation call.
         """
 
-        server_address, server_port = set_server_address_and_port(instance, "127.0.0.1", 80)
+        server_address, server_port = set_server_address_and_port(
+            instance, "127.0.0.1", 80
+        )
         agent_name = kwargs.get("name", "NOT_FOUND")
         llm_config = kwargs.get("llm_config", {})
         system_message = kwargs.get("system_message", "")
 
-        span_name = f"{SemanticConvention.GEN_AI_OPERATION_TYPE_CREATE_AGENT} {agent_name}"
+        span_name = (
+            f"{SemanticConvention.GEN_AI_OPERATION_TYPE_CREATE_AGENT} {agent_name}"
+        )
 
         with tracer.start_as_current_span(span_name, kind=SpanKind.CLIENT) as span:
             start_time = time.time()
@@ -51,7 +61,7 @@ def async_conversable_agent(version, environment, application_name, tracer, pric
                     span=span,
                     capture_message_content=capture_message_content,
                     disable_metrics=disable_metrics,
-                    version=version
+                    version=version,
                 )
 
             except Exception as e:
@@ -61,8 +71,17 @@ def async_conversable_agent(version, environment, application_name, tracer, pric
 
     return wrapper
 
-def async_agent_run(version, environment, application_name, tracer, pricing_info,
-    capture_message_content, metrics, disable_metrics):
+
+def async_agent_run(
+    version,
+    environment,
+    application_name,
+    tracer,
+    pricing_info,
+    capture_message_content,
+    metrics,
+    disable_metrics,
+):
     """
     Generates a telemetry wrapper for AG2 async agent run execution.
     """
@@ -72,7 +91,9 @@ def async_agent_run(version, environment, application_name, tracer, pricing_info
         Wraps the AG2 async agent run execution call.
         """
 
-        server_address, server_port = set_server_address_and_port(instance, "127.0.0.1", 80)
+        server_address, server_port = set_server_address_and_port(
+            instance, "127.0.0.1", 80
+        )
 
         # Extract agent name from instance
         agent_name = getattr(instance, "name", "NOT_FOUND")
@@ -103,7 +124,7 @@ def async_agent_run(version, environment, application_name, tracer, pricing_info
                     span=span,
                     capture_message_content=capture_message_content,
                     disable_metrics=disable_metrics,
-                    version=version
+                    version=version,
                 )
 
             except Exception as e:

openlit/instrumentation/ag2/utils.py

@@ -1,6 +1,7 @@
 """
 AG2 OpenTelemetry instrumentation utility functions
 """
+
 import time
 
 from opentelemetry.trace import Status, StatusCode
@@ -12,6 +13,7 @@ from openlit.__helpers import (
 )
 from openlit.semcov import SemanticConvention
 
+
 def calculate_tokens_and_cost(response, request_model, pricing_info):
     """
     Calculate the input, output tokens, and their respective costs from AG2 response.
@@ -21,7 +23,9 @@ def calculate_tokens_and_cost(response, request_model, pricing_info):
 
     # Early return if response doesn't have cost data
     if not hasattr(response, "cost") or response.cost is None:
-        cost = get_chat_model_cost(request_model, pricing_info, input_tokens, output_tokens)
+        cost = get_chat_model_cost(
+            request_model, pricing_info, input_tokens, output_tokens
+        )
         return input_tokens, output_tokens, cost
 
     try:
@@ -34,6 +38,7 @@ def calculate_tokens_and_cost(response, request_model, pricing_info):
     cost = get_chat_model_cost(request_model, pricing_info, input_tokens, output_tokens)
     return input_tokens, output_tokens, cost
 
+
 def _extract_tokens_from_cost(cost_data):
     """
     Extract input and output tokens from AG2 cost data structure.
@@ -52,6 +57,7 @@ def _extract_tokens_from_cost(cost_data):
 
     return input_tokens, output_tokens
 
+
 def format_content(chat_history):
     """
     Format the chat history into a string for span events.
@@ -67,32 +73,62 @@ def format_content(chat_history):
 
     return "\n".join(formatted_messages)
 
-def common_agent_logic(scope, pricing_info, environment, application_name, metrics,
-    capture_message_content, disable_metrics, version, operation_type):
+
+def common_agent_logic(
+    scope,
+    pricing_info,
+    environment,
+    application_name,
+    metrics,
+    capture_message_content,
+    disable_metrics,
+    version,
+    operation_type,
+):
     """
     Process agent request and generate Telemetry
     """
 
     # Common Span Attributes
-    common_span_attributes(scope,
-        operation_type, SemanticConvention.GEN_AI_SYSTEM_AG2,
-        scope._server_address, scope._server_port, scope._request_model, scope._response_model,
-        environment, application_name, False, 0, scope._end_time - scope._start_time, version)
+    common_span_attributes(
+        scope,
+        operation_type,
+        SemanticConvention.GEN_AI_SYSTEM_AG2,
+        scope._server_address,
+        scope._server_port,
+        scope._request_model,
+        scope._response_model,
+        environment,
+        application_name,
+        False,
+        0,
+        scope._end_time - scope._start_time,
+        version,
+    )
 
     # Span Attributes for Agent-specific parameters
     scope._span.set_attribute(SemanticConvention.GEN_AI_AGENT_NAME, scope._agent_name)
 
     # Span Attributes for Response parameters
     if hasattr(scope, "_input_tokens"):
-        scope._span.set_attribute(SemanticConvention.GEN_AI_USAGE_INPUT_TOKENS, scope._input_tokens)
-        scope._span.set_attribute(SemanticConvention.GEN_AI_USAGE_OUTPUT_TOKENS, scope._output_tokens)
-        scope._span.set_attribute(SemanticConvention.GEN_AI_CLIENT_TOKEN_USAGE, scope._input_tokens + scope._output_tokens)
+        scope._span.set_attribute(
+            SemanticConvention.GEN_AI_USAGE_INPUT_TOKENS, scope._input_tokens
+        )
+        scope._span.set_attribute(
+            SemanticConvention.GEN_AI_USAGE_OUTPUT_TOKENS, scope._output_tokens
+        )
+        scope._span.set_attribute(
+            SemanticConvention.GEN_AI_CLIENT_TOKEN_USAGE,
+            scope._input_tokens + scope._output_tokens,
+        )
         scope._span.set_attribute(SemanticConvention.GEN_AI_USAGE_COST, scope._cost)
 
     # Span Attributes for Content
     if capture_message_content and hasattr(scope, "_chat_history"):
         chat_content = format_content(scope._chat_history)
-        scope._span.set_attribute(SemanticConvention.GEN_AI_CONTENT_COMPLETION, chat_content)
+        scope._span.set_attribute(
+            SemanticConvention.GEN_AI_CONTENT_COMPLETION, chat_content
+        )
 
         # To be removed once the change to span_attributes (from span events) is complete
         scope._span.add_event(
@@ -104,20 +140,51 @@ def common_agent_logic(scope, pricing_info, environment, application_name, metri
 
     # Set agent description for create agent operation
     if hasattr(scope, "_system_message"):
-        scope._span.set_attribute(SemanticConvention.GEN_AI_AGENT_DESCRIPTION, scope._system_message)
+        scope._span.set_attribute(
+            SemanticConvention.GEN_AI_AGENT_DESCRIPTION, scope._system_message
+        )
 
     scope._span.set_status(Status(StatusCode.OK))
 
     # Metrics
     if not disable_metrics and hasattr(scope, "_input_tokens"):
-        record_completion_metrics(metrics, operation_type, SemanticConvention.GEN_AI_SYSTEM_AG2,
-            scope._server_address, scope._server_port, scope._request_model, scope._response_model, environment,
-            application_name, scope._start_time, scope._end_time, scope._input_tokens, scope._output_tokens,
-            scope._cost, 0, scope._end_time - scope._start_time)
-
-def process_agent_creation(agent_name, llm_config, system_message, pricing_info, server_port, server_address,
-    environment, application_name, metrics, start_time, span, capture_message_content=False,
-    disable_metrics=False, version="1.0.0", **kwargs):
+        record_completion_metrics(
+            metrics,
+            operation_type,
+            SemanticConvention.GEN_AI_SYSTEM_AG2,
+            scope._server_address,
+            scope._server_port,
+            scope._request_model,
+            scope._response_model,
+            environment,
+            application_name,
+            scope._start_time,
+            scope._end_time,
+            scope._input_tokens,
+            scope._output_tokens,
+            scope._cost,
+            0,
+            scope._end_time - scope._start_time,
+        )
+
+
+def process_agent_creation(
+    agent_name,
+    llm_config,
+    system_message,
+    pricing_info,
+    server_port,
+    server_address,
+    environment,
+    application_name,
+    metrics,
+    start_time,
+    span,
+    capture_message_content=False,
+    disable_metrics=False,
+    version="1.0.0",
+    **kwargs,
+):
     """
     Process agent creation and generate Telemetry
     """
@@ -134,12 +201,36 @@ def process_agent_creation(agent_name, llm_config, system_message, pricing_info,
     scope._system_message = system_message
     scope._server_address, scope._server_port = server_address, server_port
 
-    common_agent_logic(scope, pricing_info, environment, application_name, metrics,
-        capture_message_content, disable_metrics, version, SemanticConvention.GEN_AI_OPERATION_TYPE_CREATE_AGENT)
-
-def process_agent_run(response, agent_name, request_model, pricing_info, server_port, server_address,
-    environment, application_name, metrics, start_time, span, capture_message_content=False,
-    disable_metrics=False, version="1.0.0", **kwargs):
+    common_agent_logic(
+        scope,
+        pricing_info,
+        environment,
+        application_name,
+        metrics,
+        capture_message_content,
+        disable_metrics,
+        version,
+        SemanticConvention.GEN_AI_OPERATION_TYPE_CREATE_AGENT,
+    )
+
+
+def process_agent_run(
+    response,
+    agent_name,
+    request_model,
+    pricing_info,
+    server_port,
+    server_address,
+    environment,
+    application_name,
+    metrics,
+    start_time,
+    span,
+    capture_message_content=False,
+    disable_metrics=False,
+    version="1.0.0",
+    **kwargs,
+):
     """
     Process agent run and generate Telemetry
     """
@@ -157,19 +248,31 @@ def process_agent_run(response, agent_name, request_model, pricing_info, server_
 
     # Calculate tokens and cost
     scope._input_tokens, scope._output_tokens, scope._cost = calculate_tokens_and_cost(
-        response, request_model, pricing_info)
+        response, request_model, pricing_info
+    )
 
     # Extract response model from cost data
     try:
         if hasattr(response, "cost") and response.cost is not None:
             cost_data = response.cost.get("usage_including_cached_inference", {})
-            scope._response_model = list(cost_data.keys())[1] if len(cost_data) > 1 else request_model
+            scope._response_model = (
+                list(cost_data.keys())[1] if len(cost_data) > 1 else request_model
+            )
         else:
             scope._response_model = request_model
     except (AttributeError, IndexError, KeyError, TypeError):
         scope._response_model = request_model
 
-    common_agent_logic(scope, pricing_info, environment, application_name, metrics,
-        capture_message_content, disable_metrics, version, SemanticConvention.GEN_AI_OPERATION_TYPE_EXECUTE_AGENT_TASK)
+    common_agent_logic(
+        scope,
+        pricing_info,
+        environment,
+        application_name,
+        metrics,
+        capture_message_content,
+        disable_metrics,
+        version,
+        SemanticConvention.GEN_AI_OPERATION_TYPE_EXECUTE_AGENT_TASK,
+    )
 
     return response