From 2d2d0ba9904c27be06074f650a1d5109de21994f Mon Sep 17 00:00:00 2001 From: ZanSara Date: Fri, 23 Feb 2024 10:41:38 +0100 Subject: [PATCH] chore: Use `serialize_callable` instead of `serialize_callback_handler` in Ollama (#461) * change function * ruff --- .../components/generators/ollama/generator.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/integrations/ollama/src/haystack_integrations/components/generators/ollama/generator.py b/integrations/ollama/src/haystack_integrations/components/generators/ollama/generator.py index 321eab9f3..f3ab86282 100644 --- a/integrations/ollama/src/haystack_integrations/components/generators/ollama/generator.py +++ b/integrations/ollama/src/haystack_integrations/components/generators/ollama/generator.py @@ -3,8 +3,8 @@ import requests from haystack import component, default_from_dict, default_to_dict -from haystack.components.generators.utils import deserialize_callback_handler, serialize_callback_handler from haystack.dataclasses import StreamingChunk +from haystack.utils.callable_serialization import deserialize_callable, serialize_callable from requests import Response @@ -57,7 +57,7 @@ def to_dict(self) -> Dict[str, Any]: Serialize this component to a dictionary. :return: The serialized component as a dictionary. """ - callback_name = serialize_callback_handler(self.streaming_callback) if self.streaming_callback else None + callback_name = serialize_callable(self.streaming_callback) if self.streaming_callback else None return default_to_dict( self, timeout=self.timeout, @@ -80,7 +80,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "OllamaGenerator": init_params = data.get("init_parameters", {}) serialized_callback_handler = init_params.get("streaming_callback") if serialized_callback_handler: - data["init_parameters"]["streaming_callback"] = deserialize_callback_handler(serialized_callback_handler) + data["init_parameters"]["streaming_callback"] = deserialize_callable(serialized_callback_handler) return default_from_dict(cls, data) def _create_json_payload(self, prompt: str, stream: bool, generation_kwargs=None) -> Dict[str, Any]: