Fix ollama response parsing (#1406)

griptape-ai · Dec 9, 2024 · b9d09e0 · b9d09e0
1 parent 3606156
commit b9d09e0
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 19 deletions.
diff --git a/griptape/drivers/prompt/ollama_prompt_driver.py b/griptape/drivers/prompt/ollama_prompt_driver.py
@@ -29,7 +29,7 @@
 logger = logging.getLogger(Defaults.logging_config.logger_name)
 
 if TYPE_CHECKING:
-    from ollama import Client
+    from ollama import ChatResponse, Client
 
     from griptape.tokenizers.base_tokenizer import BaseTokenizer
     from griptape.tools import BaseTool
@@ -81,13 +81,10 @@ def try_run(self, prompt_stack: PromptStack) -> Message:
         response = self.client.chat(**params)
         logger.debug(response)
 
-        if isinstance(response, dict):
-            return Message(
-                content=self.__to_prompt_stack_message_content(response),
-                role=Message.ASSISTANT_ROLE,
-            )
-        else:
-            raise Exception("invalid model response")
+        return Message(
+            content=self.__to_prompt_stack_message_content(response),
+            role=Message.ASSISTANT_ROLE,
+        )
 
     @observable
     def try_stream(self, prompt_stack: PromptStack) -> Iterator[DeltaMessage]:
@@ -213,7 +210,7 @@ def __to_ollama_role(self, message: Message, message_content: Optional[BaseMessa
             else:
                 return "user"
 
-    def __to_prompt_stack_message_content(self, response: dict) -> list[BaseMessageContent]:
+    def __to_prompt_stack_message_content(self, response: ChatResponse) -> list[BaseMessageContent]:
         content = []
         message = response["message"]
 

diff --git a/tests/unit/drivers/prompt/test_ollama_prompt_driver.py b/tests/unit/drivers/prompt/test_ollama_prompt_driver.py
@@ -230,16 +230,6 @@ def test_try_run(self, mock_client, prompt_stack, messages, use_native_tools):
         assert message.value[1].value.path == "test"
         assert message.value[1].value.input == {"foo": "bar"}
 
-    def test_try_run_bad_response(self, mock_client):
-        # Given
-        prompt_stack = PromptStack()
-        driver = OllamaPromptDriver(model="llama")
-        mock_client.return_value.chat.return_value = "bad-response"
-
-        # When/Then
-        with pytest.raises(Exception, match="invalid model response"):
-            driver.try_run(prompt_stack)
-
     def test_try_stream_run(self, mock_stream_client):
         # Given
         prompt_stack = PromptStack()