Qwen support for expert agent

backend/src/agents/expert.py

@@ -3,12 +3,15 @@ import json
 import inspect
 from string import Template
 
-from inference import LlamaMessage, replicate_llama_completion, Models
+import inference
+from inference import LlamaMessage
 from agent import Agent, AgentField
 
 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]
 """.strip()
 
 class ExpertAgent(Agent):
@@ -16,6 +19,7 @@ class ExpertAgent(Agent):
     description = CLASS_DESCRIPTION
     schema_fields = [
         AgentField('system_prompt', 'string', default="You are an expert AI assistant."),
+        AgentField('model', 'string', default="qwen-big"),
     ]
 
     def __init__(self, app, data):
@@ -45,7 +49,20 @@ class ExpertAgent(Agent):
         if prompt:
             messages.append(LlamaMessage("user", prompt))
 
-        response = await replicate_llama_completion(self.app, messages, model=Models.llama_70b, max_tokens=2048)
+        model = self.config['model']
+
+        if model == "qwen-big":
+            response = await inference.replicate_qwen_completion(self.app, messages, model=inference.Models.quen3_235b)
+        elif model == "qwen-small":
+            response = await inference.replicate_qwen_completion(self.app, messages, model=inference.Models.quen3_8b)
+        elif model == "llama-big":
+            response = await inference.replicate_llama_completion(self.app, messages, model=inference.Models.llama_70b)
+        elif model == "llama-small":
+            response = await inference.replicate_llama_completion(self.app, messages, model=inference.Models.llama_8b)
+        else:
+            # An error here is questionable, model is user configurable
+            raise Exception("Unkonwn agent model: " + model)
+
         return response
 
     async def received_message(self, message):

backend/src/database.py

@@ -119,6 +119,9 @@ class Thread(Base):
     name: Mapped[Optional[str]]
     description: Mapped[Optional[str]]
 
+    # Threads can form trees where threads continue each other
+    #previous_thread_id: Mapped[Optional[int]] = mapped_column(ForeignKey("threads.id"))
+
 # A persistent object or value associated with an agent
 class AgentResource(Base):
     __modelname__ = 'agent_resource'
@@ -134,3 +137,11 @@ class AgentResource(Base):
     key: Mapped[str]
     value = mapped_column(JSON, default=None)
 
+class MessageDatum(Base):
+    __modelname__ = 'message_datum'
+    __tablename__ = 'message_data'
+
+    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
+    message_id: Mapped[int] = mapped_column(ForeignKey("messages.id"))
+    user_id: Mapped[int] = mapped_column(ForeignKey("users.id"))
+    value = mapped_column(JSON, default=None)

backend/src/inference.py

@@ -14,27 +14,51 @@ class LlamaMessage:
 class Models:
     llama_8b = "meta/meta-llama-3-8b-instruct"
     llama_70b = "meta/meta-llama-3-70b-instruct"
+    qwen3_8b = "lucataco/qwen3-vl-8b-instruct:39e893666996acf464cff75688ad49ac95ef54e9f1c688fbc677330acc478e11"
+    quen3_235b = "qwen/qwen3-235b-a22b-instruct-2507"
 
-async def replicate_llama_completion(app, messages, model=None, temperature=0.75, max_tokens=500):
+async def replicate_llama_completion(app, messages, model=None, opts={}):
     if model is None:
         model = Models.llama_8b
 
-    # Llama instruct
-    full_prompt = "<|begin_of_text|>"
-    for message in messages:
-        full_prompt += f"<|start_header_id|>{message.role}<|end_header_id|>\n\n{message.content}<|eot_id|>"
-    full_prompt += f"<|start_header_id|>assistant<|end_header_id|>"
+    return await replicate_base_completion(app, model, {
+        "prompt": format_instruct(messages),
+        "temperature": opts.get('temperature', 0.75),
+        "max_new_tokens": opts.get('max_tokens', 2048),
+        "system_prompt": "" # System prompt is formatted in the prompt
+    })
+
+
+async def replicate_qwen_completion(app, messages, model=None, opts={}):
+    if model is None:
+        model = Models.qwen3_8b
+
+    if model == Models.qwen3_8b:
+        input = {
+            "prompt": format_instruct(messages),
+            "temperature": opts.get('temperature', 0.75),
+            "max_new_tokens": opts.get('max_tokens', 2048),
+        }
+    elif model == Models.quen3_235b:
+        input = {
+            "prompt": format_instruct(messages),
+            "temperature": opts.get('temperature', 0.75),
+            "max_tokens": opts.get('max_tokens', 2048),
+        }
+    else:
+        raise Exception("Invalid qwen model: " + model)
+
+    return await replicate_base_completion(app, model, input)
     
+async def replicate_base_completion(app, model, input):
     content_stream = Stream(app)
     thread = threading.Thread(
         target=run_replicate_inference_loop,
         kwargs={
             'replicate_client': app.replicate_client,
             'model': model,
-            'prompt': full_prompt,
             'content_stream': content_stream,
-            'temperature': temperature,
-            'max_tokens': max_tokens,
+            'input': input
         }
     )
     thread.start()
@@ -42,19 +66,18 @@ async def replicate_llama_completion(app, messages, model=None, temperature=0.75
     return response.lstrip()
 
 
-def run_replicate_inference_loop(replicate_client, model, prompt, content_stream, temperature, max_tokens):
+def run_replicate_inference_loop(replicate_client, model, input, content_stream):
     print("START INFERENCE")
-    completion = replicate_client.run(
-        model,
-        input={
-            "prompt": prompt,
-            "temperature": temperature,
-            "max_new_tokens": max_tokens,
-            "system_prompt": "" # We handle system prompt in the formatted prompt
-        }
-    )
+    completion = replicate_client.run(model, input=input)
     for event in completion:
         if event:
             content_stream.feed(str(event))
     print("END INFERENCE")
-    content_stream.close()
+    content_stream.close()
+
+def format_instruct(messages):
+    prompt = "<|begin_of_text|>"
+    for message in messages:
+        prompt += f"<|start_header_id|>{message.role}<|end_header_id|>\n\n{message.content}<|eot_id|>"
+    prompt += f"<|start_header_id|>assistant<|end_header_id|>"
+    return prompt