update anthropic and bedrock to use new typed dicts

Author: ncybul

ddtrace/contrib/internal/anthropic/utils.py

@@ -24,7 +24,7 @@
 from ddtrace.llmobs._integrations.base import BaseLLMIntegration
 from ddtrace.llmobs._integrations.utils import update_proxy_workflow_input_output_value
 from ddtrace.llmobs._utils import _get_attr
-from ddtrace.llmobs.utils import ToolCall
+from ddtrace.llmobs.utils import Message, ToolCall
 from ddtrace.llmobs.utils import ToolDefinition
 from ddtrace.llmobs.utils import ToolResult
 from ddtrace.trace import Span
@@ -71,7 +71,7 @@ def _llmobs_set_tags(
         system_prompt = kwargs.get("system")
         input_messages = self._extract_input_message(messages, system_prompt)
 
-        output_messages = [{"content": ""}]
+        output_messages: List[Message] = [Message(content="")]
         if not span.error and response is not None:
             output_messages = self._extract_output_message(response)
         span_kind = "workflow" if span._get_ctx_item(PROXY_REQUEST) else "llm"
@@ -92,14 +92,14 @@ def _llmobs_set_tags(
         )
         update_proxy_workflow_input_output_value(span, span_kind)
 
-    def _extract_input_message(self, messages, system_prompt: Optional[Union[str, List[Dict[str, Any]]]] = None):
+    def _extract_input_message(self, messages, system_prompt: Optional[Union[str, List[Dict[str, Any]]]] = None) -> List[Message]:
         """Extract input messages from the stored prompt.
         Anthropic allows for messages and multiple texts in a message, which requires some special casing.
         """
         if not isinstance(messages, Iterable):
             log.warning("Anthropic input must be a list of messages.")
 
-        input_messages = []
+        input_messages: List[Message] = []
         if system_prompt is not None:
             messages = [{"content": system_prompt, "role": "system"}] + messages
 
@@ -115,43 +115,43 @@ def _extract_input_message(self, messages, system_prompt: Optional[Union[str, Li
                 log.warning("Anthropic input message must have content and role.")
 
             if isinstance(content, str):
-                input_messages.append({"content": content, "role": role})
+                input_messages.append(Message(content=content, role=str(role)))
 
             elif isinstance(content, list):
                 for block in content:
                     if _get_attr(block, "type", None) == "text":
-                        input_messages.append({"content": _get_attr(block, "text", ""), "role": role})
+                        input_messages.append(Message(content=_get_attr(block, "text", ""), role=str(role)))
 
                     elif _get_attr(block, "type", None) == "image":
                         # Store a placeholder for potentially enormous binary image data.
-                        input_messages.append({"content": "([IMAGE DETECTED])", "role": role})
+                        input_messages.append(Message(content="([IMAGE DETECTED])", role=str(role)))
 
                     elif _get_attr(block, "type", None) == "tool_use":
                         text = _get_attr(block, "text", None)
                         input_data = _get_attr(block, "input", "")
                         if isinstance(input_data, str):
                             input_data = json.loads(input_data)
                         tool_call_info = ToolCall(
-                            name=_get_attr(block, "name", ""),
+                            name=str(_get_attr(block, "name", "")),
                             arguments=input_data,
-                            tool_id=_get_attr(block, "id", ""),
-                            type=_get_attr(block, "type", ""),
+                            tool_id=str(_get_attr(block, "id", "")),
+                            type=str(_get_attr(block, "type", "")),
                         )
                         if text is None:
                             text = ""
-                        input_messages.append({"content": text, "role": role, "tool_calls": [tool_call_info]})
+                        input_messages.append(Message(content=str(text), role=str(role), tool_calls=[tool_call_info]))
 
                     elif _get_attr(block, "type", None) == "tool_result":
                         content = _get_attr(block, "content", None)
                         formatted_content = self._format_tool_result_content(content)
                         tool_result_info = ToolResult(
                             result=formatted_content,
-                            tool_id=_get_attr(block, "tool_use_id", ""),
+                            tool_id=str(_get_attr(block, "tool_use_id", "")),
                             type="tool_result",
                         )
-                        input_messages.append({"content": "", "role": role, "tool_results": [tool_result_info]})
+                        input_messages.append(Message(content="", role=str(role), tool_results=[tool_result_info]))
                     else:
-                        input_messages.append({"content": str(block), "role": role})
+                        input_messages.append(Message(content=str(block), role=str(role)))
 
         return input_messages
 
@@ -169,34 +169,34 @@ def _format_tool_result_content(self, content) -> str:
             return ",".join(formatted_content)
         return str(content)
 
-    def _extract_output_message(self, response):
+    def _extract_output_message(self, response) -> List[Message]:
         """Extract output messages from the stored response."""
-        output_messages = []
+        output_messages: List[Message] = []
         content = _get_attr(response, "content", "")
         role = _get_attr(response, "role", "")
 
         if isinstance(content, str):
-            return [{"content": content, "role": role}]
+            return [Message(content=content, role=str(role))]
 
         elif isinstance(content, list):
             for completion in content:
                 text = _get_attr(completion, "text", None)
                 if isinstance(text, str):
-                    output_messages.append({"content": text, "role": role})
+                    output_messages.append(Message(content=text, role=str(role)))
                 else:
                     if _get_attr(completion, "type", None) == "tool_use":
                         input_data = _get_attr(completion, "input", "")
                         if isinstance(input_data, str):
                             input_data = json.loads(input_data)
                         tool_call_info = ToolCall(
-                            name=_get_attr(completion, "name", ""),
+                            name=str(_get_attr(completion, "name", "")),
                             arguments=input_data,
-                            tool_id=_get_attr(completion, "id", ""),
-                            type=_get_attr(completion, "type", ""),
+                            tool_id=str(_get_attr(completion, "id", "")),
+                            type=str(_get_attr(completion, "type", "")),
                         )
                         if text is None:
                             text = ""
-                        output_messages.append({"content": text, "role": role, "tool_calls": [tool_call_info]})
+                        output_messages.append(Message(content=str(text), role=str(role), tool_calls=[tool_call_info]))
         return output_messages
 
     def _extract_usage(self, span: Span, usage: Dict[str, Any]):

ddtrace/llmobs/_integrations/anthropic.py

@@ -35,7 +35,7 @@
 from ddtrace.llmobs._telemetry import record_bedrock_agent_span_event_created
 from ddtrace.llmobs._utils import _get_attr
 from ddtrace.llmobs._writer import LLMObsSpanEvent
-from ddtrace.llmobs.utils import ToolDefinition
+from ddtrace.llmobs.utils import Message, ToolDefinition
 from ddtrace.trace import Span
 
 
@@ -191,7 +191,7 @@ def translate_bedrock_traces(self, traces, root_span) -> None:
         self._active_span_by_step_id.clear()
 
     @staticmethod
-    def _extract_input_message_for_converse(prompt: List[Dict[str, Any]]):
+    def _extract_input_message_for_converse(prompt: List[Dict[str, Any]]) -> List[Message]:
         """Extract input messages from the stored prompt for converse
 
         `prompt` is an array of `message` objects. Each `message` has a role and content field.
@@ -203,7 +203,7 @@ def _extract_input_message_for_converse(prompt: List[Dict[str, Any]]):
         """
         if not isinstance(prompt, list):
             log.warning("Bedrock input is not a list of messages or a string.")
-            return [{"content": ""}]
+            return [Message(content="")]
         input_messages = []
         for message in prompt:
             if not isinstance(message, dict):
@@ -342,41 +342,42 @@ def _converse_output_stream_processor() -> (
         return messages, metadata, usage_metrics
 
     @staticmethod
-    def _extract_input_message(prompt):
+    def _extract_input_message(prompt) -> List[Message]:
         """Extract input messages from the stored prompt.
         Anthropic allows for messages and multiple texts in a message, which requires some special casing.
         """
         if isinstance(prompt, str):
-            return [{"content": prompt}]
+            return [Message(content=prompt)]
         if not isinstance(prompt, list):
             log.warning("Bedrock input is not a list of messages or a string.")
-            return [{"content": ""}]
-        input_messages = []
+            return [Message(content="")]
+        input_messages: List[Message] = []
         for p in prompt:
             content = p.get("content", "")
             if isinstance(content, list) and isinstance(content[0], dict):
                 for entry in content:
                     if entry.get("type") == "text":
-                        input_messages.append({"content": entry.get("text", ""), "role": str(p.get("role", ""))})
+                        input_messages.append(Message(content=entry.get("text", ""), role=str(p.get("role", ""))))
                     elif entry.get("type") == "image":
                         # Store a placeholder for potentially enormous binary image data.
-                        input_messages.append({"content": "([IMAGE DETECTED])", "role": str(p.get("role", ""))})
+                        input_messages.append(Message(content="([IMAGE DETECTED])", role=str(p.get("role", ""))))
             else:
-                input_messages.append({"content": content, "role": str(p.get("role", ""))})
+                input_messages.append(Message(content=str(content), role=str(p.get("role", ""))))
         return input_messages
 
     @staticmethod
-    def _extract_output_message(response):
+    def _extract_output_message(response) -> List[Message]:
         """Extract output messages from the stored response.
         Anthropic allows for chat messages, which requires some special casing.
         """
         if isinstance(response["text"], str):
-            return [{"content": response["text"]}]
+            return [Message(content=response["text"])]
         if isinstance(response["text"], list):
             if isinstance(response["text"][0], str):
-                return [{"content": str(content)} for content in response["text"]]
+                return [Message(content=str(content)) for content in response["text"]]
             if isinstance(response["text"][0], dict):
-                return [{"content": response["text"][0].get("text", "")}]
+                return [Message(content=response["text"][0].get("text", ""))]
+        return []
 
     def _get_base_url(self, **kwargs: Dict[str, Any]) -> Optional[str]:
         instance = kwargs.get("instance")
@@ -396,8 +397,8 @@ def _extract_tool_definitions(self, tool_config: Dict[str, Any]) -> List[ToolDef
         for tool in tools:
             tool_spec = _get_attr(tool, "toolSpec", {})
             tool_definition_info = ToolDefinition(
-                name=_get_attr(tool_spec, "name", ""),
-                description=_get_attr(tool_spec, "description", ""),
+                name=str(_get_attr(tool_spec, "name", "")),
+                description=str(_get_attr(tool_spec, "description", "")),
                 schema=_get_attr(tool_spec, "inputSchema", {}),
             )
             tool_definitions.append(tool_definition_info)