restored to original

added the support for the gpt-4o-mini

app.py

@@ -1,10 +1,13 @@
 import os
 import warnings
 from typing import *
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-from langchain.llms import HuggingFacePipeline
+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
 
 from interface import create_demo
 from medrax.agent import *
@@ -12,6 +15,8 @@ from medrax.tools import *
 from medrax.utils import *
 
 warnings.filterwarnings("ignore")
+logging.set_verbosity_error()
+_ = load_dotenv()
 
 def initialize_agent(
     prompt_file,
@@ -19,11 +24,27 @@ def initialize_agent(
     model_dir="./model-weights",
     temp_dir="temp",
     device="cuda",
+    model="gpt-4o-mini",
     temperature=0.7,
-    top_p=0.95
+    top_p=0.95,
+    openai_kwargs={}
 ):
-    """Initialize the MedRAX agent with specified tools and configuration."""
-
+    """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"]
 
@@ -45,6 +66,7 @@ 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:
@@ -52,24 +74,7 @@ def initialize_agent(
             tools_dict[tool_name] = all_tools[tool_name]()
 
     checkpointer = MemorySaver()
-
-    # Load local Hugging Face model
-    hf_model_id = model_id = "HuggingFaceH4/zephyr-7b-alpha"
-    tokenizer = AutoTokenizer.from_pretrained(hf_model_id)
-    raw_model = AutoModelForCausalLM.from_pretrained(hf_model_id, device_map="auto")
-
-    pipe = pipeline(
-        "text-generation",
-        model=raw_model,
-        tokenizer=tokenizer,
-        max_new_tokens=512,
-        temperature=temperature,
-        top_p=top_p,
-        return_full_text=False,
-    )
-
-    model = HuggingFacePipeline(pipeline=pipe)
-
+    model = ChatOpenAI(model=model, temperature=temperature, top_p=top_p, **openai_kwargs)
     agent = Agent(
         model,
         tools=list(tools_dict.values()),
@@ -82,10 +87,15 @@ def initialize_agent(
     print("Agent initialized")
     return agent, tools_dict
 
-
 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",
@@ -98,15 +108,26 @@ if __name__ == "__main__":
         # "ChestXRayGeneratorTool",
     ]
 
+    # Collect the ENV variables
+    openai_kwargs = {}
+    if api_key := os.getenv("OPENAI_API_KEY"):
+        openai_kwargs["api_key"] = api_key
+
+    if base_url := os.getenv("OPENAI_BASE_URL"):
+        openai_kwargs["base_url"] = base_url
+
     agent, tools_dict = initialize_agent(
         "medrax/docs/system_prompts.txt",
         tools_to_use=selected_tools,
-        model_dir="./model-weights",
-        temp_dir="temp",
-        device="cuda",
+        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-mini",  # Change this to the model you want to use, e.g. gpt-4o-mini
         temperature=0.7,
-        top_p=0.95
+        top_p=0.95,
+        openai_kwargs=openai_kwargs
     )
-
     demo = create_demo(agent, tools_dict)
+
+    # demo.launch(server_name="0.0.0.0", server_port=8585, share=True)
     demo.launch(debug=True, ssr_mode=False)