From 69e8614a98af2cf4a54b936aa6b37e2e32b7e7af Mon Sep 17 00:00:00 2001
From: Shivering Ocelot <shiveringocelot@pm.me>
Date: Thu, 17 Jul 2025 21:10:28 -0500
Subject: [PATCH] enable chat for expert, custom inference parameters

---
 backend/src/agents/expert.py | 56 ++++++++++++++++++++++++++----------
 backend/src/inference.py     | 15 +++++-----
 2 files changed, 49 insertions(+), 22 deletions(-)

diff --git a/backend/src/agents/expert.py b/backend/src/agents/expert.py
index 3aba9f0..ab46222 100644
--- a/backend/src/agents/expert.py
+++ b/backend/src/agents/expert.py
@@ -22,29 +22,55 @@ class ExpertAgent(Agent):
         Agent.__init__(self, app, data)
 
     async def prompt(self, instruction, input=None, **kwargs):
-        # if instruction and input are passed, instruction extends the system prompt
-        # otherwise, use the one passed as the user prompt, with the regular system prompt
-
-        # TODO: if 'history' in kwargs: interpret it as user-agent turns assuming list[str]
+        # Instruction extends system if there's input or messages
+        # Otherwise instruction is the user prompt
+        # Input is always added as a user message at the end, if present
 
+        # Use 'messages' as starting point when passed
+        messages = []
+        if 'messages' in kwargs:
+            messages = [LlamaMessage(msg['role'], msg['content']) for msg in kwargs['messages']]
+        
+        # Construct the default messages
         system = self.config['system_prompt']
-        if instruction != None and input != None:
+        prompt = input
+
+        if input == None and not messages:
+            prompt = instruction
+        elif instruction:
             system += "\n" + instruction
+        
+        # Add default messages to list
+        messages.insert(0, LlamaMessage("system", system))
+        if prompt:
+            messages.append(LlamaMessage("user", prompt))
 
-        prompt = input if input != None else instruction
-        messages = [
-            LlamaMessage("system", system),
-            LlamaMessage("user", prompt)
-        ]
-
-        response = await replicate_llama_completion(self.app, messages, model=Models.llama_70b)
+        response = await replicate_llama_completion(self.app, messages, model=Models.llama_70b, max_tokens=2048)
         return response
 
     async def received_message(self, message):
-        raise Exception("TODO: expert chat support")
+        chat_view = self.app.get_chat_view(message.chat_id)
+        thread = chat_view.get_thread(message.thread_id)
+
+        def msg_role(message):
+            sender = message.sender
+            is_digit = isinstance(sender, (int, str)) and str(sender).isdigit()
+
+            if is_digit and str(sender) == str(self.id):
+                return 'assistant'
+            elif sender == 'user':
+                return 'user'
+            else:
+                raise Exception("Unsupported sender: " + sender)
+        
+        messages = [{'role': msg_role(msg), 'content': msg.content} for msg in thread.messages]
+
+        response = await self.prompt(
+            instruction=None,
+            input=None,
+            messages=messages
+        )
 
-        messages = []
-        response = await replicate_llama_completion(self.app, messages, model=Models.llama_70b)
         self.app.send_chat_message(
             chat_id=message.chat_id,
             sender=self.id,
diff --git a/backend/src/inference.py b/backend/src/inference.py
index ab93f7d..bd4584f 100644
--- a/backend/src/inference.py
+++ b/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
         }
     )