Ollama

backend/src/inference.py

@@ -1,6 +1,7 @@
 import threading
 from stream import Stream
 
+
 VALID_ROLES = ['user', 'assistant', 'system']
 
 class LlamaMessage:
@@ -17,6 +18,28 @@ class Models:
     qwen3_8b = "lucataco/qwen3-vl-8b-instruct:39e893666996acf464cff75688ad49ac95ef54e9f1c688fbc677330acc478e11"
     qwen3_235b = "qwen/qwen3-235b-a22b-instruct-2507"
 
+async def ollama_chat_completion(app, messages, model, opts={}):
+    from ollama import AsyncClient
+
+    content_stream = Stream(app)
+
+    print("START OLLAMA INFERENCE")
+    stream_response = await AsyncClient().chat(
+        model=model,
+        messages=[{'role': msg.role, 'content': msg.content} for msg in messages],
+        stream=True,
+    )
+
+    async for part in stream_response:
+        content_stream.feed(part.message.content)
+        print(part['message']['content'], end='', flush=True)
+
+    content_stream.close()
+
+    print("END OLLAMA INFERENCE")
+    return await content_stream.content_wait()
+
+
 async def replicate_llama_completion(app, messages, model=None, opts={}):
     if model is None:
         model = Models.llama_8b