refactor: replace bare dict with dict[str, Any] in ops trace providers (#35082)

api/core/ops/arize_phoenix_trace/arize_phoenix_trace.py

@@ -778,7 +778,7 @@ class ArizePhoenixDataTrace(BaseTraceInstance):
             logger.info("[Arize/Phoenix] Failed to construct project URL: %s", str(e), exc_info=True)
             raise ValueError(f"[Arize/Phoenix] Failed to construct project URL: {str(e)}")
 
-    def _construct_llm_attributes(self, prompts: dict | list | str | None) -> dict[str, str]:
+    def _construct_llm_attributes(self, prompts: dict[str, Any] | list[Any] | str | None) -> dict[str, str]:
         """Construct LLM attributes with passed prompts for Arize/Phoenix."""
         attributes: dict[str, str] = {}
 
@@ -797,7 +797,9 @@ class ArizePhoenixDataTrace(BaseTraceInstance):
             path = f"{SpanAttributes.LLM_INPUT_MESSAGES}.{message_index}.{key}"
             set_attribute(path, value)
 
-        def set_tool_call_attributes(message_index: int, tool_index: int, tool_call: dict | object | None) -> None:
+        def set_tool_call_attributes(
+            message_index: int, tool_index: int, tool_call: dict[str, Any] | object | None
+        ) -> None:
             """Extract and assign tool call details safely."""
             if not tool_call:
                 return

api/core/ops/mlflow_trace/mlflow_trace.py

@@ -242,7 +242,7 @@ class MLflowDataTrace(BaseTraceInstance):
 
         return inputs, attributes
 
-    def _parse_knowledge_retrieval_outputs(self, outputs: dict):
+    def _parse_knowledge_retrieval_outputs(self, outputs: dict[str, Any]):
         """Parse KR outputs and attributes from KR workflow node"""
         retrieved = outputs.get("result", [])
 
@@ -319,7 +319,7 @@ class MLflowDataTrace(BaseTraceInstance):
             end_time_ns=datetime_to_nanoseconds(trace_info.end_time),
         )
 
-    def _get_message_user_id(self, metadata: dict) -> str | None:
+    def _get_message_user_id(self, metadata: dict[str, Any]) -> str | None:
         if (end_user_id := metadata.get("from_end_user_id")) and (
             end_user_data := db.session.get(EndUser, end_user_id)
         ):
@@ -468,7 +468,7 @@ class MLflowDataTrace(BaseTraceInstance):
         }
         return node_type_mapping.get(node_type, "CHAIN")  # type: ignore[arg-type,call-overload]
 
-    def _set_trace_metadata(self, span: Span, metadata: dict):
+    def _set_trace_metadata(self, span: Span, metadata: dict[str, Any]):
         token = None
         try:
             # NB: Set span in context such that we can use update_current_trace() API
@@ -490,7 +490,7 @@ class MLflowDataTrace(BaseTraceInstance):
             return messages
         return prompts  # Fallback to original format
 
-    def _parse_single_message(self, item: dict):
+    def _parse_single_message(self, item: dict[str, Any]):
         """Postprocess single message format to be standard chat message"""
         role = item.get("role", "user")
         msg = {"role": role, "content": item.get("text", "")}

api/core/ops/opik_trace/opik_trace.py

@@ -3,7 +3,7 @@ import logging
 import os
 import uuid
 from datetime import datetime, timedelta
-from typing import cast
+from typing import Any, cast
 
 from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionMetadataKey
 from opik import Opik, Trace
@@ -436,7 +436,7 @@ class OpikDataTrace(BaseTraceInstance):
 
         self.add_span(span_data)
 
-    def add_trace(self, opik_trace_data: dict) -> Trace:
+    def add_trace(self, opik_trace_data: dict[str, Any]) -> Trace:
         try:
             trace = self.opik_client.trace(**opik_trace_data)
             logger.debug("Opik Trace created successfully")
@@ -444,7 +444,7 @@ class OpikDataTrace(BaseTraceInstance):
         except Exception as e:
             raise ValueError(f"Opik Failed to create trace: {str(e)}")
 
-    def add_span(self, opik_span_data: dict):
+    def add_span(self, opik_span_data: dict[str, Any]):
         try:
             self.opik_client.span(**opik_span_data)
             logger.debug("Opik Span created successfully")