diff --git a/backend/requirements.txt b/backend/requirements.txt index d9edbb5..d47cc7b 100644 --- a/backend/requirements.txt +++ b/backend/requirements.txt @@ -4,4 +4,5 @@ quart-cors sqlalchemy nest_asyncio replicate -qdrant-client[fastembed] \ No newline at end of file +qdrant-client[fastembed] +ollama \ No newline at end of file diff --git a/backend/src/agents/expert.py b/backend/src/agents/expert.py index f01271e..18cb2eb 100644 --- a/backend/src/agents/expert.py +++ b/backend/src/agents/expert.py @@ -11,7 +11,7 @@ CLASS_DESCRIPTION = """ Expert AI in any field, powered by Llama 3 70b. Will get guidelines, documents, and other advanced features in the future. -Available models are: [llama/qwen]-[big/small] +Available models are: [llama/qwen]-[big/small] (replicate), ollama:MODEL_NAME """.strip() class ExpertAgent(Agent): @@ -51,7 +51,9 @@ class ExpertAgent(Agent): model = self.config['model'] - if model == "qwen-big": + if model.startswith("ollama:"): + response = await inference.ollama_chat_completion(self.app, messages, model=model[7:]) + elif model == "qwen-big": response = await inference.replicate_qwen_completion(self.app, messages, model=inference.Models.qwen3_235b) elif model == "qwen-small": response = await inference.replicate_qwen_completion(self.app, messages, model=inference.Models.qwen3_8b) diff --git a/backend/src/inference.py b/backend/src/inference.py index 512fcf0..08df92b 100644 --- a/backend/src/inference.py +++ b/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