updates

app.py

@@ -4,8 +4,6 @@ from typing import *
 from dotenv import load_dotenv
 from transformers import logging
 
-from langgraph.checkpoint.memory import MemorySaver
-from langchain_openai import ChatOpenAI
 from langgraph.checkpoint.memory import MemorySaver
 from langchain_openai import ChatOpenAI
 
@@ -16,12 +14,7 @@ from medrax.utils import *
 
 warnings.filterwarnings("ignore")
 logging.set_verbosity_error()
-_ = load_dotenv()
-
-openai_kwargs = {
-"openai_api_key": os.environ.get("OPENAI_API_KEY"),
-"openai_api_base": os.environ.get("OPENAI_BASE_URL"),
-}
+_ = load_dotenv()  # Loads your .env file (OPENAI_API_KEY and OPENAI_BASE_URL)
 
 def initialize_agent(
     prompt_file,
@@ -29,28 +22,11 @@ def initialize_agent(
     model_dir="./model-weights",
     temp_dir="temp",
     device="cuda",
-    model="google/gemini-2.5-pro-exp-03-25:free",
+    model="google/gemini-1.5-flash-latest",  # ✅ updated model name for OpenRouter
     temperature=0.7,
     top_p=0.95, 
-    openai_kwargs=openai_kwargs
-
+    openai_kwargs=None
 ):
-    """Initialize the MedRAX agent with specified tools and configuration.
-
-    Args:
-        prompt_file (str): Path to file containing system prompts
-        tools_to_use (List[str], optional): List of tool names to initialize. If None, all tools are initialized.
-        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.
-        openai_kwargs (dict, optional): Additional keyword arguments for OpenAI API, such as API key and base URL.
-
-    Returns:
-        Tuple[Agent, Dict[str, BaseTool]]: Initialized agent and dictionary of tool instances
-    """
     prompts = load_prompts_from_file(prompt_file)
     prompt = prompts["MEDICAL_ASSISTANT"]
 
@@ -72,7 +48,6 @@ def initialize_agent(
         "DicomProcessorTool": lambda: DicomProcessorTool(temp_dir=temp_dir),
     }
 
-    # Initialize only selected tools or all if none specified
     tools_dict = {}
     tools_to_use = tools_to_use or all_tools.keys()
     for tool_name in tools_to_use:
@@ -80,6 +55,7 @@ def initialize_agent(
             tools_dict[tool_name] = all_tools[tool_name]()
 
     checkpointer = MemorySaver()
+
     model = ChatOpenAI(model=model, temperature=temperature, top_p=top_p, **openai_kwargs)
     agent = Agent(
         model,
@@ -95,14 +71,8 @@ def initialize_agent(
 
 
 if __name__ == "__main__":
-    """
-    This is the main entry point for the MedRAX application.
-    It initializes the agent with the selected tools and creates the demo.
-    """
     print("Starting server...")
 
-    # Example: initialize with only specific tools
-    # Here three tools are commented out, you can uncomment them to use them
     selected_tools = [
         "ImageVisualizerTool",
         "DicomProcessorTool",
@@ -115,26 +85,24 @@ if __name__ == "__main__":
         # "ChestXRayGeneratorTool",
     ]
 
-    # Collect the ENV variables
+    # ✅ Collect environment variables and pass to model
     openai_kwargs = {}
     if api_key := os.getenv("OPENAI_API_KEY"):
         openai_kwargs["openai_api_key"] = api_key
-    
     if base_url := os.getenv("OPENAI_BASE_URL"):
         openai_kwargs["openai_api_base"] = base_url
-    
 
     agent, tools_dict = initialize_agent(
         "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
+        model_dir="./model-weights",
+        temp_dir="temp",
+        device="cuda",
+        model="google/gemini-1.5-flash-latest",  # ✅ Updated OpenRouter model
         temperature=0.7,
         top_p=0.95,
         openai_kwargs=openai_kwargs
     )
-    demo = create_demo(agent, tools_dict)
 
+    demo = create_demo(agent, tools_dict)
     demo.launch(server_name="0.0.0.0", server_port=8585, share=True)