minor refactoring

integrations/ollama/src/haystack_integrations/components/generators/ollama/generator.py

@@ -103,7 +103,7 @@ def _create_json_payload(self, prompt: str, stream: bool, generation_kwargs=None
             "options": generation_kwargs,
         }
 
-    def _convert_to_haystack_response(self, ollama_response: Response) -> Dict[str, List[Any]]:
+    def _convert_to_response(self, ollama_response: Response) -> Dict[str, List[Any]]:
         """
         Convert a response from the Ollama API to the required Haystack format.
         :param ollama_response: A response (requests library) from the Ollama API.
@@ -117,7 +117,7 @@ def _convert_to_haystack_response(self, ollama_response: Response) -> Dict[str,
 
         return {"replies": replies, "meta": [meta]}
 
-    def _convert_to_response(self, chunks: List[StreamingChunk]) -> Dict[str, List[Any]]:
+    def _convert_to_streaming_response(self, chunks: List[StreamingChunk]) -> Dict[str, List[Any]]:
         """
         Convert a list of chunks response required Haystack format.
         :param chunks: List of StreamingChunks
@@ -184,6 +184,6 @@ def run(
 
         if stream:
             chunks: List[StreamingChunk] = self._handle_streaming_response(response)
-            return self._convert_to_response(chunks)
+            return self._convert_to_streaming_response(chunks)
 
-        return self._convert_to_haystack_response(response)
+        return self._convert_to_response(response)

integrations/ollama/tests/test_generator.py

@@ -165,8 +165,6 @@ def test_create_json_payload(self, configuration: dict[str, Any], stream: bool):
 
         assert observed == expected
 
-
-class TestOllamaStreamingGenerator:
     @pytest.mark.integration
     def test_ollama_generator_run_streaming(self):
         class Callback:
@@ -187,4 +185,4 @@ def __call__(self, chunk):
         assert "Amsterdam" in results["replies"][0]
         assert len(results["meta"]) == 1
         assert callback.responses == results["replies"][0]
-        assert callback.count_calls > 1
+        assert callback.count_calls > 1