fix types [AUTH-1050]

Opentrons · Nov 19, 2024 · 145e925 · 145e925
1 parent f188e93
commit 145e925
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 9 deletions.
diff --git a/opentrons-ai-client/src/assets/localization/en/protocol_generator.json b/opentrons-ai-client/src/assets/localization/en/protocol_generator.json
@@ -30,7 +30,7 @@
   "login": "Login",
   "logout": "Logout",
   "make_sure_your_prompt": "Write a prompt in a natural language for OpentronsAI to generate a protocol using the Opentrons Python Protocol API v2. The better the prompt, the better the quality of the protocol produced by OpentronsAI.",
-  "modify_intro": "Modify the following Python code using the Opentrons Python Protocol API v2. Ensure that the new labware and pipettes are compatible with the robot type specified in the protocol.\n\n",
+  "modify_intro": "Modify the following Python code using the Opentrons Python Protocol API v2. Ensure that the new labware and pipettes are compatible with the Flex robot.\n\n",
   "modify_python_code": "Original Python Code:\n",
   "modify_type_of_update": "Type of update:\n- ",
   "modify_details_of_change": "Detail of changes:\n- ",

diff --git a/opentrons-ai-server/api/domain/anthropic_predict.py b/opentrons-ai-server/api/domain/anthropic_predict.py
@@ -5,7 +5,14 @@
 import requests
 import structlog
 from anthropic import Anthropic
-from anthropic.types import Message, MessageParam
+from anthropic.types import (
+    Message,
+    MessageParam,
+    TextBlock,
+    TextBlockParam,
+    ToolParam,
+    ToolUseBlock,
+)
 from ddtrace import tracer
 
 from api.domain.config_anthropic import DOCUMENTS, PROMPT, SYSTEM_PROMPT
@@ -27,11 +34,11 @@ def __init__(self, settings: Settings) -> None:
             {
                 "role": "user",
                 "content": [
-                    {"type": "text", "text": DOCUMENTS.format(doc_content=self.get_docs()), "cache_control": {"type": "ephemeral"}}
+                    TextBlockParam(type="text", text=DOCUMENTS.format(doc_content=self.get_docs()), cache_control={"type": "ephemeral"})
                 ],
             }
         ]
-        self.tools: List[Dict[str, Any]] = [
+        self.tools: List[ToolParam] = [
             {
                 "name": "simulate_protocol",
                 "description": "Simulates the python protocol on user input. Returned value is text indicating if protocol is successful.",
@@ -104,7 +111,7 @@ def predict(self, prompt: str) -> str | None:
         try:
             self._messages.append({"role": "user", "content": PROMPT.format(USER_PROMPT=prompt)})
             response = self.generate_message()
-            if response.content[-1].type == "tool_use":
+            if isinstance(response.content[-1], ToolUseBlock):
                 tool_use = response.content[-1]
                 self._messages.append({"role": "assistant", "content": response.content})
                 result = self.handle_tool_use(tool_use.name, tool_use.input)
@@ -125,7 +132,7 @@ def predict(self, prompt: str) -> str | None:
                 self._messages.append({"role": "assistant", "content": response_text})
                 return response_text
 
-            elif response.content[0].type == "text":
+            elif isinstance(response.content[0], TextBlock):
                 response_text = response.content[0].text
                 self._messages.append({"role": "assistant", "content": response_text})
                 return response_text
@@ -154,7 +161,7 @@ def reset(self) -> None:
             {
                 "role": "user",
                 "content": [
-                    {"type": "text", "text": DOCUMENTS.format(doc_content=self.get_docs()), "cache_control": {"type": "ephemeral"}}
+                    TextBlockParam(type="text", text=DOCUMENTS.format(doc_content=self.get_docs()), cache_control={"type": "ephemeral"})
                 ],
             }
         ]
@@ -175,9 +182,9 @@ def simulate_protocol(self, protocol: str) -> str:
         response_data = response.json()
         if "error_message" in response_data:
             logger.error("Simulation error", extra={"error": response_data["error_message"]})
-            return response_data["error_message"]
+            return str(response_data["error_message"])
         elif "protocol_name" in response_data:
-            return response_data["run_status"]
+            return str(response_data["run_status"])
         else:
             logger.error("Unexpected response", extra={"response": response_data})
             return "Unexpected response"