black formatting

ddtrace/llmobs/_evaluators/ragas/base.py

@@ -28,9 +28,7 @@ def __init__(self):
 
         self.ragas_version = parse_version(ragas.__version__)
         if self.ragas_version >= (0, 2, 0) or self.ragas_version < (0, 1, 10):
-            raise NotImplementedError(
-                "Ragas version: {} is not supported".format(self.ragas_version),
-            )
+            raise NotImplementedError("Ragas version: {} is not supported".format(self.ragas_version))
 
         from ragas.llms import llm_factory
 

ddtrace/llmobs/_evaluators/ragas/faithfulness.py

@@ -84,14 +84,7 @@ def evaluate(self, span_event: dict) -> Tuple[Union[float, str], Optional[dict]]
         evaluation_metadata = {EVALUATION_KIND_METADATA: "faithfulness"}  # type: dict[str, Union[str, dict, list]]
 
         # initialize data we annotate for tracing ragas
-        score, question, answer, context, statements, faithfulness_list = (
-            math.nan,
-            None,
-            None,
-            None,
-            None,
-            None,
-        )
+        score, question, answer, context, statements, faithfulness_list = (math.nan, None, None, None, None, None)
 
         with self.llmobs_service.workflow(
             "dd-ragas.faithfulness", ml_app=_get_ml_app_for_ragas_trace(span_event)
@@ -145,9 +138,7 @@ def evaluate(self, span_event: dict) -> Tuple[Union[float, str], Optional[dict]]
 
     def _create_statements(self, question: str, answer: str) -> Optional[List[str]]:
         with self.llmobs_service.workflow("dd-ragas.create_statements"):
-            self.llmobs_service.annotate(
-                input_data={"question": question, "answer": answer},
-            )
+            self.llmobs_service.annotate(input_data={"question": question, "answer": answer})
             statements_prompt = self._create_statements_prompt(answer=answer, question=question)
 
             """LLM step to break down the answer into simpler statements"""
@@ -160,9 +151,7 @@ def _create_statements(self, question: str, answer: str) -> Optional[List[str]]:
             statements = [item["simpler_statements"] for item in statements.dicts()]
             statements = [item for sublist in statements for item in sublist]
 
-            self.llmobs_service.annotate(
-                output_data=statements,
-            )
+            self.llmobs_service.annotate(output_data=statements)
             if not isinstance(statements, List):
                 return None
             return statements
@@ -172,10 +161,7 @@ def _create_verdicts(self, context: str, statements: List[str]):
         Returns: `StatementFaithfulnessAnswers` model detailing which statements are faithful to the context
         """
         with self.llmobs_service.workflow("dd-ragas.create_verdicts") as create_verdicts_workflow:
-            self.llmobs_service.annotate(
-                span=create_verdicts_workflow,
-                input_data=statements,
-            )
+            self.llmobs_service.annotate(span=create_verdicts_workflow, input_data=statements)
             """Check which statements contradict the conntext"""
             raw_nli_results = self.ragas_faithfulness_instance.llm.generate_text(
                 self._create_natural_language_inference_prompt(context, statements)
@@ -206,10 +192,7 @@ def _create_verdicts(self, context: str, statements: List[str]):
                 logger.debug("Failed to parse faithfulness_list", exc_info=e)
                 return None
             finally:
-                self.llmobs_service.annotate(
-                    span=create_verdicts_workflow,
-                    output_data=faithfulness_list,
-                )
+                self.llmobs_service.annotate(span=create_verdicts_workflow, output_data=faithfulness_list)
 
     def _create_statements_prompt(self, answer, question):
         # Returns: `ragas.llms.PromptValue` object
@@ -242,10 +225,7 @@ def _compute_score(self, faithfulness_list) -> float:
             else:
                 score = math.nan
             self.llmobs_service.annotate(
-                metadata={
-                    "faithful_statements": faithful_statements,
-                    "num_statements": num_statements,
-                },
+                metadata={"faithful_statements": faithful_statements, "num_statements": num_statements},
                 output_data=score,
             )
             return score

ddtrace/llmobs/_evaluators/runner.py

@@ -15,9 +15,7 @@
 logger = get_logger(__name__)
 
 
-SUPPORTED_EVALUATORS = {
-    RagasFaithfulnessEvaluator.LABEL: RagasFaithfulnessEvaluator,
-}
+SUPPORTED_EVALUATORS = {RagasFaithfulnessEvaluator.LABEL: RagasFaithfulnessEvaluator}
 
 
 class EvaluatorRunner(PeriodicService):
@@ -59,10 +57,7 @@ def __init__(self, interval: float, llmobs_service=None, evaluators=None):
                         namespace=TELEMETRY_APM_PRODUCT.LLMOBS,
                         name="evaluators.init",
                         value=1,
-                        tags=(
-                            ("evaluator_label", evaluator),
-                            ("state", evaluator_init_state),
-                        ),
+                        tags=(("evaluator_label", evaluator), ("state", evaluator_init_state)),
                     )
             else:
                 raise ValueError("Parsed unsupported evaluator: {}".format(evaluator))
@@ -83,11 +78,7 @@ def _stop_service(self) -> None:
         self.executor.shutdown(wait=True)
 
     def recreate(self) -> "EvaluatorRunner":
-        return self.__class__(
-            interval=self._interval,
-            llmobs_service=self.llmobs_service,
-            evaluators=self.evaluators,
-        )
+        return self.__class__(interval=self._interval, llmobs_service=self.llmobs_service, evaluators=self.evaluators)
 
     def enqueue(self, span_event: Dict, span: Span) -> None:
         with self._lock:

ddtrace/llmobs/_integrations/anthropic.py

@@ -32,11 +32,7 @@ class AnthropicIntegration(BaseLLMIntegration):
     _integration_name = "anthropic"
 
     def _set_base_span_tags(
-        self,
-        span: Span,
-        model: Optional[str] = None,
-        api_key: Optional[str] = None,
-        **kwargs: Dict[str, Any],
+        self, span: Span, model: Optional[str] = None, api_key: Optional[str] = None, **kwargs: Dict[str, Any]
     ) -> None:
         """Set base level tags that should be present on all Anthropic spans (if they are not None)."""
         if model is not None:
@@ -48,12 +44,7 @@ def _set_base_span_tags(
                 span.set_tag_str(API_KEY, api_key)
 
     def _llmobs_set_tags(
-        self,
-        span: Span,
-        args: List[Any],
-        kwargs: Dict[str, Any],
-        response: Optional[Any] = None,
-        operation: str = "",
+        self, span: Span, args: List[Any], kwargs: Dict[str, Any], response: Optional[Any] = None, operation: str = ""
     ) -> None:
         """Extract prompt/response tags from a completion and set them as temporary "_ml_obs.*" tags."""
         parameters = {}

ddtrace/llmobs/_integrations/base.py

@@ -198,12 +198,7 @@ def trunc(self, text: str) -> str:
         return text
 
     def llmobs_set_tags(
-        self,
-        span: Span,
-        args: List[Any],
-        kwargs: Dict[str, Any],
-        response: Optional[Any] = None,
-        operation: str = "",
+        self, span: Span, args: List[Any], kwargs: Dict[str, Any], response: Optional[Any] = None, operation: str = ""
     ) -> None:
         """Extract input/output information from the request and response to be submitted to LLMObs."""
         if not self.llmobs_enabled:
@@ -216,11 +211,6 @@ def llmobs_set_tags(
 
     @abc.abstractmethod
     def _llmobs_set_tags(
-        self,
-        span: Span,
-        args: List[Any],
-        kwargs: Dict[str, Any],
-        response: Optional[Any] = None,
-        operation: str = "",
+        self, span: Span, args: List[Any], kwargs: Dict[str, Any], response: Optional[Any] = None, operation: str = ""
     ) -> None:
         raise NotImplementedError()

ddtrace/llmobs/_integrations/bedrock.py

@@ -28,12 +28,7 @@ class BedrockIntegration(BaseLLMIntegration):
     _integration_name = "bedrock"
 
     def _llmobs_set_tags(
-        self,
-        span: Span,
-        args: List[Any],
-        kwargs: Dict[str, Any],
-        response: Optional[Any] = None,
-        operation: str = "",
+        self, span: Span, args: List[Any], kwargs: Dict[str, Any], response: Optional[Any] = None, operation: str = ""
     ) -> None:
         """Extract prompt/response tags from a completion and set them as temporary "_ml_obs.*" tags."""
         if span.get_tag(PROPAGATED_PARENT_ID_KEY) is None:

ddtrace/llmobs/_integrations/gemini.py

@@ -33,12 +33,7 @@ def _set_base_span_tags(
             span.set_tag_str("google_generativeai.request.model", str(model))
 
     def _llmobs_set_tags(
-        self,
-        span: Span,
-        args: List[Any],
-        kwargs: Dict[str, Any],
-        response: Optional[Any] = None,
-        operation: str = "",
+        self, span: Span, args: List[Any], kwargs: Dict[str, Any], response: Optional[Any] = None, operation: str = ""
     ) -> None:
         instance = kwargs.get("instance", None)
         metadata = llmobs_get_metadata_google(kwargs, instance)

ddtrace/llmobs/_integrations/langchain.py

@@ -46,11 +46,7 @@
 VERTEXAI_PROVIDER_NAME = "vertexai"
 GEMINI_PROVIDER_NAME = "google_palm"
 
-ROLE_MAPPING = {
-    "human": "user",
-    "ai": "assistant",
-    "system": "system",
-}
+ROLE_MAPPING = {"human": "user", "ai": "assistant", "system": "system"}
 
 SUPPORTED_OPERATIONS = ["llm", "chat", "chain", "embedding", "retrieval", "tool"]
 
@@ -175,12 +171,7 @@ def _llmobs_set_tags_from_llm(
         span._set_ctx_item(output_tag_key, message_content)
 
     def _llmobs_set_tags_from_chat_model(
-        self,
-        span: Span,
-        args: List[Any],
-        kwargs: Dict[str, Any],
-        chat_completions: Any,
-        is_workflow: bool = False,
+        self, span: Span, args: List[Any], kwargs: Dict[str, Any], chat_completions: Any, is_workflow: bool = False
     ) -> None:
         span._set_ctx_items(
             {
@@ -369,8 +360,7 @@ def _llmobs_set_meta_tags_from_embedding(
                 embeddings_count = len(output_embedding)
             embedding_dim = len(output_values[0])
             span._set_ctx_item(
-                output_tag_key,
-                "[{} embedding(s) returned with size {}]".format(embeddings_count, embedding_dim),
+                output_tag_key, "[{} embedding(s) returned with size {}]".format(embeddings_count, embedding_dim)
             )
         except (TypeError, IndexError):
             log.warning("Failed to write output vectors", output_embedding)
@@ -425,12 +415,7 @@ def _llmobs_set_meta_tags_from_tool(self, span: Span, tool_inputs: Dict[str, Any
         if not span.error and tool_output is not None:
             formatted_outputs = self.format_io(tool_output)
         span._set_ctx_items(
-            {
-                SPAN_KIND: "tool",
-                METADATA: metadata,
-                INPUT_VALUE: formatted_input,
-                OUTPUT_VALUE: formatted_outputs,
-            }
+            {SPAN_KIND: "tool", METADATA: metadata, INPUT_VALUE: formatted_input, OUTPUT_VALUE: formatted_outputs}
         )
 
     def _set_base_span_tags(  # type: ignore[override]
@@ -537,10 +522,7 @@ def check_token_usage_ai_message(self, ai_message):
 
         return (input_tokens, output_tokens, total_tokens), run_id_base
 
-    def format_io(
-        self,
-        messages,
-    ):
+    def format_io(self, messages):
         """
         Formats input and output messages for serialization to JSON.
         Specifically, makes sure that any schema messages are converted to strings appropriately.

ddtrace/llmobs/_integrations/vertexai.py

@@ -33,12 +33,7 @@ def _set_base_span_tags(
             span.set_tag_str("vertexai.request.model", model)
 
     def _llmobs_set_tags(
-        self,
-        span: Span,
-        args: List[Any],
-        kwargs: Dict[str, Any],
-        response: Optional[Any] = None,
-        operation: str = "",
+        self, span: Span, args: List[Any], kwargs: Dict[str, Any], response: Optional[Any] = None, operation: str = ""
     ) -> None:
         instance = kwargs.get("instance", None)
         history = kwargs.get("history", [])

ddtrace/llmobs/_llmobs.py

@@ -105,8 +105,7 @@ def __init__(self, tracer=None):
             timeout=float(os.getenv("_DD_LLMOBS_WRITER_TIMEOUT", 5.0)),
         )
         self._evaluator_runner = EvaluatorRunner(
-            interval=float(os.getenv("_DD_LLMOBS_EVALUATOR_INTERVAL", 1.0)),
-            llmobs_service=self,
+            interval=float(os.getenv("_DD_LLMOBS_EVALUATOR_INTERVAL", 1.0)), llmobs_service=self
         )
 
         forksafe.register(self._child_after_fork)
@@ -638,9 +637,11 @@ def agent(cls, name: Optional[str] = None, session_id: Optional[str] = None, ml_
         if cls.enabled is False:
             log.warning(SPAN_START_WHILE_DISABLED_WARNING)
         return cls._instance._start_span("agent", name=name, session_id=session_id, ml_app=ml_app)
-    
+
     @classmethod
-    def _experiment(cls, name: Optional[str] = None, session_id: Optional[str] = None, ml_app: Optional[str] = None) -> Span:
+    def _experiment(
+        cls, name: Optional[str] = None, session_id: Optional[str] = None, ml_app: Optional[str] = None
+    ) -> Span:
         """
         Trace a dynamic workflow in which an embedded language model (agent) decides what sequence of actions to take.
 
@@ -1026,10 +1027,7 @@ def submit_evaluation_for(
                 raise TypeError(
                     "`span_with_tag_value` must be a dict with keys 'tag_key' and 'tag_value' containing string values"
                 )
-            join_on["tag"] = {
-                "key": span_with_tag_value.get("tag_key"),
-                "value": span_with_tag_value.get("tag_value"),
-            }
+            join_on["tag"] = {"key": span_with_tag_value.get("tag_key"), "value": span_with_tag_value.get("tag_value")}
 
         timestamp_ms = timestamp_ms if timestamp_ms else int(time.time() * 1000)
 
@@ -1052,10 +1050,7 @@ def submit_evaluation_for(
             log.warning("tags must be a dictionary of string key-value pairs.")
             tags = {}
 
-        evaluation_tags = {
-            "ddtrace.version": ddtrace.__version__,
-            "ml_app": ml_app,
-        }
+        evaluation_tags = {"ddtrace.version": ddtrace.__version__, "ml_app": ml_app}
 
         if tags:
             for k, v in tags.items():
@@ -1181,10 +1176,7 @@ def submit_evaluation(
             return
 
         # initialize tags with default values that will be overridden by user-provided tags
-        evaluation_tags = {
-            "ddtrace.version": ddtrace.__version__,
-            "ml_app": ml_app,
-        }
+        evaluation_tags = {"ddtrace.version": ddtrace.__version__, "ml_app": ml_app}
 
         if tags:
             for k, v in tags.items():

ddtrace/llmobs/_log_writer.py

@@ -59,10 +59,7 @@ def __init__(self, site, api_key, interval, timeout):
         self._endpoint = "/api/v2/logs"  # type: str
         self._site = site  # type: str
         self._intake = "http-intake.logs.%s" % self._site  # type: str
-        self._headers = {
-            "DD-API-KEY": self._api_key,
-            "Content-Type": "application/json",
-        }
+        self._headers = {"DD-API-KEY": self._api_key, "Content-Type": "application/json"}
         logger.debug("started log writer to %r", self._url)
 
     def start(self, *args, **kwargs):

ddtrace/llmobs/_utils.py

@@ -167,7 +167,7 @@ def _get_session_id(span: Span) -> Optional[str]:
 def _inject_llmobs_parent_id(span_context):
     """Inject the LLMObs parent ID into the span context for reconnecting distributed LLMObs traces."""
     span = ddtrace.tracer.current_span()
-    
+
     if span is None:
         log.warning("No active span to inject LLMObs parent ID info.")
         return
@@ -207,11 +207,11 @@ def __init__(self, resp) -> None:
 
     @property
     def status_code(self) -> int:
-        if hasattr(self._resp, 'status'):
+        if hasattr(self._resp, "status"):
             return self._resp.status
-        elif hasattr(self._resp, 'code'):
+        elif hasattr(self._resp, "code"):
             return self._resp.code
-        elif hasattr(self._resp, 'getcode'):
+        elif hasattr(self._resp, "getcode"):
             return self._resp.getcode()
         else:
             raise AttributeError(f"Could not find status code in response object of type {type(self._resp)}")
@@ -222,7 +222,7 @@ def read(self) -> bytes:
         return self._content
 
     def text(self) -> str:
-        return self.read().decode('utf-8')
+        return self.read().decode("utf-8")
 
     def json(self) -> dict:
         return json.loads(self.text())