openrouter fix

app.py

@@ -5,7 +5,7 @@ from dotenv import load_dotenv
 from transformers import logging
 
 from langgraph.checkpoint.memory import MemorySaver
-from langchain_openai import ChatOpenAI
+from langchain_community.chat_models import ChatOpenRouter
 
 from interface import create_demo
 from medrax.agent import *
@@ -14,7 +14,7 @@ from medrax.utils import *
 
 warnings.filterwarnings("ignore")
 logging.set_verbosity_error()
-_ = load_dotenv()  # Loads your .env file (OPENAI_API_KEY and OPENAI_BASE_URL)
+_ = load_dotenv()
 
 def initialize_agent(
     prompt_file,
@@ -22,11 +22,12 @@ def initialize_agent(
     model_dir="./model-weights",
     temp_dir="temp",
     device="cuda",
-    model="google/gemini-1.5-flash-latest",  # ✅ updated model name for OpenRouter
+    model="google/gemini-2.5-pro-exp-03-25:free",
     temperature=0.7,
-    top_p=0.95, 
-    openai_kwargs=None
+    top_p=0.95
 ):
+    """Initialize the MedRAX agent with specified tools and configuration."""
+
     prompts = load_prompts_from_file(prompt_file)
     prompt = prompts["MEDICAL_ASSISTANT"]
 
@@ -55,8 +56,14 @@ def initialize_agent(
             tools_dict[tool_name] = all_tools[tool_name]()
 
     checkpointer = MemorySaver()
+    model = ChatOpenRouter(
+        model_name=model,
+        api_key=os.getenv("OPENAI_API_KEY"),
+        base_url=os.getenv("OPENAI_BASE_URL"),
+        temperature=temperature,
+        top_p=top_p,
+    )
 
-    model = ChatOpenAI(model=model, temperature=temperature, top_p=top_p, **openai_kwargs)
     agent = Agent(
         model,
         tools=list(tools_dict.values()),
@@ -85,23 +92,15 @@ if __name__ == "__main__":
         # "ChestXRayGeneratorTool",
     ]
 
-    # ✅ 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",
         temp_dir="temp",
         device="cuda",
-        model="google/gemini-1.5-flash-latest",  # ✅ Updated OpenRouter model
+        model="google/gemini-2.5-pro-exp-03-25:free",
         temperature=0.7,
-        top_p=0.95,
-        openai_kwargs=openai_kwargs
+        top_p=0.95
     )
 
     demo = create_demo(agent, tools_dict)