Qwen support for expert agent

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