참고소스 수정본

참고/guardrails-main/guardrails/telemetry/open_inference.py

@@ -0,0 +1,163 @@
+import json
+from typing import Any, Dict, List, Optional
+
+from guardrails.telemetry.common import (
+    get_span,
+    to_dict,
+    serialize,
+    recursive_key_operation,
+    redact,
+)
+
+try:
+    from openinference.semconv.trace import SpanAttributes  # type: ignore
+except ImportError:
+    SpanAttributes = None
+
+
+def trace_operation(
+    *,
+    input_mime_type: Optional[str] = None,
+    input_value: Optional[Any] = None,
+    output_mime_type: Optional[str] = None,
+    output_value: Optional[Any] = None,
+):
+    """Traces an operation (any function call) using OpenInference semantic
+    conventions."""
+    current_span = get_span()
+
+    if current_span is None:
+        return
+
+    ser_input_mime_type = serialize(input_mime_type)
+    if ser_input_mime_type:
+        current_span.set_attribute("input.mime_type", ser_input_mime_type)
+
+    ser_input_value = serialize(input_value)
+    if ser_input_value:
+        current_span.set_attribute("input.value", ser_input_value)
+
+    ser_output_mime_type = serialize(output_mime_type)
+    if ser_output_mime_type:
+        current_span.set_attribute("output.mime_type", ser_output_mime_type)
+
+    ser_output_value = serialize(output_value)
+    if ser_output_value:
+        current_span.set_attribute("output.value", ser_output_value)
+
+
+def trace_llm_call(
+    *,
+    function_call: Optional[
+        Dict[str, Any]
+    ] = None,  # JSON String	"{function_name: 'add', args: [1, 2]}"	Object recording details of a function call in models or APIs  # noqa
+    input_messages: Optional[
+        List[Dict[str, Any]]
+    ] = None,  # List of objects†	[{"message.role": "user", "message.content": "hello"}]	List of messages sent to the LLM in a chat API request  # noqa
+    invocation_parameters: Optional[
+        Dict[str, Any]
+    ] = None,  # JSON string	"{model_name: 'gpt-3', temperature: 0.7}"	Parameters used during the invocation of an LLM or API  # noqa
+    model_name: Optional[
+        str
+    ] = None,  # String	"gpt-3.5-turbo"	The name of the language model being utilized  # noqa
+    output_messages: Optional[
+        List[Dict[str, Any]]
+    ] = None,  # List of objects	[{"message.role": "user", "message.content": "hello"}]	List of messages received from the LLM in a chat API request  # noqa
+    prompt_template_template: Optional[
+        str
+    ] = None,  # String	"Weather forecast for {city} on {date}"	Template used to generate prompts as Python f-strings  # noqa
+    prompt_template_variables: Optional[
+        Dict[str, Any]
+    ] = None,  # JSON String	{ context: "<context from retrieval>", subject: "math" }	JSON of key value pairs applied to the prompt template  # noqa
+    prompt_template_version: Optional[
+        str
+    ] = None,  # String	"v1.0"	The version of the prompt template  # noqa
+    token_count_completion: Optional[
+        int
+    ] = None,  # Integer	15	The number of tokens in the completion  # noqa
+    token_count_prompt: Optional[
+        int
+    ] = None,  # Integer	5	The number of tokens in the prompt  # noqa
+    token_count_total: Optional[
+        int
+    ] = None,  # Integer	20	Total number of tokens, including prompt and completion  # noqa
+):
+    """Traces an LLM call using OpenInference semantic conventions."""
+    current_span = get_span()
+
+    if current_span is None:
+        return
+    if SpanAttributes is not None:
+        current_span.set_attribute(SpanAttributes.OPENINFERENCE_SPAN_KIND, "GUARDRAIL")
+    ser_function_call = serialize(function_call)
+    if ser_function_call:
+        current_span.set_attribute("llm.function_call", ser_function_call)
+
+    if input_messages and isinstance(input_messages, list):
+        for i, message in enumerate(input_messages):
+            msg_obj = to_dict(message)
+            for key, value in msg_obj.items():
+                if value is not None:
+                    standardized_key = f"message.{key}" if "message" not in key else key
+                    current_span.set_attribute(
+                        f"llm.input_messages.{i}.{standardized_key}",
+                        serialize(value),  # type: ignore
+                    )
+
+    ser_invocation_parameters = serialize(invocation_parameters)
+    redacted_ser_invocation_parameters = recursive_key_operation(
+        ser_invocation_parameters, redact
+    )
+    reser_invocation_parameters = (
+        json.dumps(redacted_ser_invocation_parameters)
+        if isinstance(redacted_ser_invocation_parameters, dict)
+        or isinstance(redacted_ser_invocation_parameters, list)
+        else redacted_ser_invocation_parameters
+    )
+    if reser_invocation_parameters:
+        current_span.set_attribute(
+            "llm.invocation_parameters", reser_invocation_parameters
+        )
+
+    ser_model_name = serialize(model_name)
+    if ser_model_name:
+        current_span.set_attribute("llm.model_name", ser_model_name)
+
+    if output_messages and isinstance(output_messages, list):
+        for i, message in enumerate(output_messages):
+            # Most responses are either dictionaries or Pydantic models
+            msg_obj = to_dict(message)
+            for key, value in msg_obj.items():
+                if value is not None:
+                    standardized_key = f"message.{key}" if "message" not in key else key
+                    current_span.set_attribute(
+                        f"llm.output_messages.{i}.{standardized_key}",
+                        serialize(value),  # type: ignore
+                    )
+
+    ser_prompt_template_template = serialize(prompt_template_template)
+    if ser_prompt_template_template:
+        current_span.set_attribute(
+            "llm.prompt_template.template", ser_prompt_template_template
+        )
+
+    ser_prompt_template_variables = serialize(prompt_template_variables)
+    if ser_prompt_template_variables:
+        current_span.set_attribute(
+            "llm.prompt_template.variables", ser_prompt_template_variables
+        )
+
+    ser_prompt_template_version = serialize(prompt_template_version)
+    if ser_prompt_template_version:
+        current_span.set_attribute(
+            "llm.prompt_template.version", ser_prompt_template_version
+        )
+
+    if token_count_completion:
+        current_span.set_attribute("llm.token_count.completion", token_count_completion)
+
+    if token_count_prompt:
+        current_span.set_attribute("llm.token_count.prompt", token_count_prompt)
+
+    if token_count_total:
+        current_span.set_attribute("llm.token_count.total", token_count_total)