use instead of for in Ollama (#1239)

deepset-ai · Dec 10, 2024 · 87dd2cd · 87dd2cd
1 parent cb051b2
commit 87dd2cd
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/...ions/ollama/src/haystack_integrations/components/generators/ollama/chat/chat_generator.py b/...ions/ollama/src/haystack_integrations/components/generators/ollama/chat/chat_generator.py
@@ -109,7 +109,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "OllamaChatGenerator":
         return default_from_dict(cls, data)
 
     def _message_to_dict(self, message: ChatMessage) -> Dict[str, str]:
-        return {"role": message.role.value, "content": message.content}
+        return {"role": message.role.value, "content": message.text}
 
     def _build_message_from_ollama_response(self, ollama_response: ChatResponse) -> ChatMessage:
         """

diff --git a/integrations/ollama/tests/test_chat_generator.py b/integrations/ollama/tests/test_chat_generator.py
@@ -102,7 +102,7 @@ def test_build_message_from_ollama_response(self):
         observed = OllamaChatGenerator(model=model)._build_message_from_ollama_response(ollama_response)
 
         assert observed.role == "assistant"
-        assert observed.content == "Hello! How are you today?"
+        assert observed.text == "Hello! How are you today?"
 
     @pytest.mark.integration
     def test_run(self):
@@ -121,7 +121,7 @@ def test_run(self):
 
             assert isinstance(response, dict)
             assert isinstance(response["replies"], list)
-            assert answer in response["replies"][0].content
+            assert answer in response["replies"][0].text
 
     @pytest.mark.integration
     def test_run_with_chat_history(self):
@@ -137,7 +137,7 @@ def test_run_with_chat_history(self):
 
         assert isinstance(response, dict)
         assert isinstance(response["replies"], list)
-        assert "Manchester" in response["replies"][-1].content or "Glasgow" in response["replies"][-1].content
+        assert "Manchester" in response["replies"][-1].text or "Glasgow" in response["replies"][-1].text
 
     @pytest.mark.integration
     def test_run_model_unavailable(self):
@@ -166,4 +166,4 @@ def test_run_with_streaming(self):
 
         assert isinstance(response, dict)
         assert isinstance(response["replies"], list)
-        assert "Manchester" in response["replies"][-1].content or "Glasgow" in response["replies"][-1].content
+        assert "Manchester" in response["replies"][-1].text or "Glasgow" in response["replies"][-1].text