feat: updates to agent and interface

interface.py

@@ -1,4 +1,5 @@
 import re
+import base64
 import gradio as gr
 from pathlib import Path
 import time
@@ -106,10 +107,29 @@ class ChatInterface:
 
         messages = []
         image_path = self.original_file_path or display_image
+
         if image_path is not None:
-            messages.append({"role": "user", "content": f"path: {image_path}"})
+            # Send path for tools
+            messages.append({"role": "user", "content": f"image_path: {image_path}"})
+
+            # Load and encode image for multimodal
+            with open(image_path, "rb") as img_file:
+                img_base64 = base64.b64encode(img_file.read()).decode("utf-8")
+
+            messages.append(
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "image_url",
+                            "image_url": {"url": f"data:image/jpeg;base64,{img_base64}"},
+                        }
+                    ],
+                }
+            )
+
         if message is not None:
-            messages.append({"role": "user", "content": message})
+            messages.append({"role": "user", "content": [{"type": "text", "text": message}]})
 
         try:
             for event in self.agent.workflow.stream(

main.py

@@ -19,7 +19,14 @@ _ = load_dotenv()
 
 
 def initialize_agent(
-    prompt_file, tools_to_use=None, model_dir="/model-weights", temp_dir="temp", device="cuda"
+    prompt_file,
+    tools_to_use=None,
+    model_dir="/model-weights",
+    temp_dir="temp",
+    device="cuda",
+    model="chatgpt-4o-latest",
+    temperature=0.7,
+    top_p=0.95,
 ):
     """Initialize the MedRAX agent with specified tools and configuration.
 
@@ -29,6 +36,9 @@ def initialize_agent(
         model_dir (str, optional): Directory containing model weights. Defaults to "/model-weights".
         temp_dir (str, optional): Directory for temporary files. Defaults to "temp".
         device (str, optional): Device to run models on. Defaults to "cuda".
+        model (str, optional): Model to use. Defaults to "chatgpt-4o-latest".
+        temperature (float, optional): Temperature for the model. Defaults to 0.7.
+        top_p (float, optional): Top P for the model. Defaults to 0.95.
 
     Returns:
         Tuple[Agent, Dict[str, BaseTool]]: Initialized agent and dictionary of tool instances
@@ -62,7 +72,7 @@ def initialize_agent(
             tools_dict[tool_name] = all_tools[tool_name]()
 
     checkpointer = MemorySaver()
-    model = ChatOpenAI(model="gpt-4o", temperature=0.7, top_p=0.95)
+    model = ChatOpenAI(model=model, temperature=temperature, top_p=top_p)
     agent = Agent(
         model,
         tools=list(tools_dict.values()),
@@ -98,8 +108,15 @@ if __name__ == "__main__":
     ]
 
     agent, tools_dict = initialize_agent(
-        "medrax/docs/system_prompts.txt", tools_to_use=selected_tools, model_dir="/model-weights"
+        "medrax/docs/system_prompts.txt",
+        tools_to_use=selected_tools,
+        model_dir="/model-weights",  # Change this to the path of the model weights
+        temp_dir="temp",  # Change this to the path of the temporary directory
+        device="cuda",  # Change this to the device you want to use
+        model="gpt-4o",  # Change this to the model you want to use, e.g. gpt-4o-mini
+        temperature=0.7,
+        top_p=0.95,
     )
     demo = create_demo(agent, tools_dict)
 
-    demo.launch(server_name="0.0.0.0", server_port=8585, share=True)
+    demo.launch(server_name="0.0.0.0", server_port=8989, share=True)

medrax/agent/agent.py

@@ -81,6 +81,7 @@ class Agent:
             log_dir (str, optional): Directory to save logs. Defaults to 'logs'.
         """
         self.system_prompt = system_prompt
+        print(f"System prompt: {self.system_prompt}")
         self.log_tools = log_tools
 
         if self.log_tools:

medrax/docs/system_prompts.txt

@@ -1,8 +1,9 @@
 [MEDICAL_ASSISTANT]
-You are an expert medical AI assistant.
+You are an expert medical AI assistant who can answer any medical questions and analyze medical images similar to a doctor.
+Solve using our own vision and reasoning and use tools to complement your reasoning.
 Make multiple tool calls in parallel or sequence as needed for comprehensive answers.
 Critically think about and criticize the tool outputs.
-If you need to look up some information before asking a follow up question, you are allowed to do that!
+If you need to look up some information before asking a follow up question, you are allowed to do that.
 
 [GENERAL_ASSISTANT]
 You are a helpful AI assistant. Your role is to assist users with a wide range of tasks and questions, providing accurate and useful information on various topics.