enable chat for expert, custom inference parameters

backend/src/inference.py

@@ -1,7 +1,7 @@
 import threading
 from stream import Stream
 
-VALID_ROLES = ['user', 'agent', 'system']
+VALID_ROLES = ['user', 'assistant', 'system']
 
 class LlamaMessage:
     def __init__(self, role, content):
@@ -15,15 +15,14 @@ class Models:
     llama_8b = "meta/meta-llama-3-8b-instruct"
     llama_70b = "meta/meta-llama-3-70b-instruct"
 
-async def replicate_llama_completion(app, messages, model=None):
+async def replicate_llama_completion(app, messages, model=None, temperature=0.75, max_tokens=500):
     if model is None:
         model = Models.llama_8b
 
     # Llama instruct
     full_prompt = "<|begin_of_text|>"
     for message in messages:
-        header_name = message.role if message.role != "agent" else "assistant"
-        full_prompt += f"<|start_header_id|>{header_name}<|end_header_id|>\n\n{message.content}<|eot_id|>"
+        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|>"
     
     content_stream = Stream(app)
@@ -34,6 +33,8 @@ async def replicate_llama_completion(app, messages, model=None):
             'model': model,
             'prompt': full_prompt,
             'content_stream': content_stream,
+            'temperature': temperature,
+            'max_tokens': max_tokens,
         }
     )
     thread.start()
@@ -41,14 +42,14 @@ async def replicate_llama_completion(app, messages, model=None):
     return response.lstrip()
 
 
-def run_replicate_inference_loop(replicate_client, model, prompt, content_stream):
+def run_replicate_inference_loop(replicate_client, model, prompt, content_stream, temperature, max_tokens):
     print("START INFERENCE")
     completion = replicate_client.run(
         model,
         input={
             "prompt": prompt,
-            "temperature": 0.75,
-            "max_new_tokens": 500,
+            "temperature": temperature,
+            "max_new_tokens": max_tokens,
             "system_prompt": "" # We handle system prompt in the formatted prompt
         }
     )