Use Replicate streaming API (#16597)

run-llama · Oct 18, 2024 · dcf2419 · dcf2419
1 parent 6e677c3
commit dcf2419
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/llama-index-integrations/llms/llama-index-llms-replicate/llama_index/llms/replicate/base.py b/llama-index-integrations/llms/llama-index-llms-replicate/llama_index/llms/replicate/base.py
@@ -144,11 +144,12 @@ def stream_complete(
         if not formatted:
             prompt = self.completion_to_prompt(prompt)
         input_dict = self._get_input_dict(prompt, **kwargs)
-        response_iter = replicate.run(self.model, input=input_dict)
+        response_iter = replicate.stream(self.model, input=input_dict)
 
         def gen() -> CompletionResponseGen:
             text = ""
-            for delta in response_iter:
+            for server_event in response_iter:
+                delta = str(server_event)
                 text += delta
                 yield CompletionResponse(
                     delta=delta,

diff --git a/llama-index-integrations/llms/llama-index-llms-replicate/pyproject.toml b/llama-index-integrations/llms/llama-index-llms-replicate/pyproject.toml
@@ -27,7 +27,7 @@ exclude = ["**/BUILD"]
 license = "MIT"
 name = "llama-index-llms-replicate"
 readme = "README.md"
-version = "0.2.1"
+version = "0.3.0"
 
 [tool.poetry.dependencies]
 python = ">=3.8.1,<4.0"