streamed messages

backend/src/inference.py

@@ -26,20 +26,37 @@ class ModelCache:
     last_used = None
     last_model = None
 
-# llama.cpp
+# llama.cpp 
 async def local_chat_completion(app, messages, model, opts={}):
-    llm = Llama(    
-        model_path=f"{app.data_directory}/models/{model}",
-        # n_gpu_layers=-1, # Uncomment to use GPU acceleration
-        # seed=1337, # Uncomment to set a specific seed
-        # n_ctx=2048, # Uncomment to increase the context window
-    )
-    response = llm.create_chat_completion(
-        messages = [{'role': msg.role, 'content': msg.content} for msg in messages]
-    )
+    content_stream = Stream(app)
 
-    # TODO: KeyError: 'text'
-    return response["choices"][0]["text"]
+    def worker():
+        llm = Llama(    
+            model_path=f"{app.data_directory}/models/{model}",
+            # n_gpu_layers=-1, # Uncomment to use GPU acceleration
+            # seed=1337, # Uncomment to set a specific seed
+            # n_ctx=2048, # Uncomment to increase the context window
+        )
+
+        # With stream=True, returns an iterator that infers tokens on demand (openai compatible)
+        llm_output = llm.create_chat_completion(
+            messages = [{'role': msg.role, 'content': msg.content} for msg in messages],
+            stream = True,
+        )
+
+        for chunk in llm_output:
+            delta = chunk['choices'][0]['delta']
+            if 'content' in delta:
+                chunk_text = delta['content']
+                content_stream.feed(chunk_text)
+            #chunk_text = chunk['choices'][0]['text']
+
+        content_stream.close()
+
+    thread = threading.Thread(target=worker)
+    thread.start()
+
+    return content_stream
 
 async def ollama_list(app, messages, model, opts={}):
     # TODO: test the output, and use this
@@ -112,8 +129,9 @@ async def replicate_base_completion(app, model, input):
         }
     )
     thread.start()
-    response = await content_stream.content_wait()
-    return response.lstrip()
+    #response = await content_stream.content_wait()
+    #return response.lstrip()
+    return content_stream
 
 
 def run_replicate_inference_loop(replicate_client, model, input, content_stream):