Ollama
This commit is contained in:
1 parent
56deeed52f
commit
2c7b3d61d1
3 files changed
+28
-2
No files matched your search
@@ -5,3 +5,4 @@ sqlalchemy
|
||||
nest_asyncio
|
||||
replicate
|
||||
qdrant-client[fastembed]
|
||||
ollama
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in new issue
Block a user