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):