From 1ae920a443033872f947528e5376d5af91ecaecf Mon Sep 17 00:00:00 2001 From: anakin87 Date: Mon, 9 Dec 2024 18:07:05 +0100 Subject: [PATCH] use instead of for in Ollama --- .../components/generators/ollama/chat/chat_generator.py | 2 +- integrations/ollama/tests/test_chat_generator.py | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/integrations/ollama/src/haystack_integrations/components/generators/ollama/chat/chat_generator.py b/integrations/ollama/src/haystack_integrations/components/generators/ollama/chat/chat_generator.py index b1be7a2db..f598a6e42 100644 --- a/integrations/ollama/src/haystack_integrations/components/generators/ollama/chat/chat_generator.py +++ b/integrations/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 index 0308f42ec..b3df0fbf1 100644 --- 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