Skip to content

Commit

Permalink
chore: Use serialize_callable instead of `serialize_callback_handle…
Browse files Browse the repository at this point in the history
…r` in Ollama (#461)

* change function

* ruff
  • Loading branch information
ZanSara authored Feb 23, 2024
1 parent 7a6556f commit 2d2d0ba
Showing 1 changed file with 3 additions and 3 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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,
Expand All @@ -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]:
Expand Down

0 comments on commit 2d2d0ba

Please sign in to comment.