Update agent.py

agent.py

@@ -10,7 +10,7 @@ from langchain.agents import AgentExecutor
 from langchain.agents.mrkl.base import ZeroShotAgent
 from prompts import FORMAT_INSTRUCTIONS, QUESTION_PROMPT, QUESTION_PROMPT1, SUFFIX
 from tools import make_tools
- 
+import os
 from rmrkl import ChatZeroShotAgent, RetryAgentExecutor
 from langchain_ollama import OllamaLLM
 import base64
@@ -35,17 +35,28 @@ def _make_llm(model, temp, api_key, streaming: bool = False):
        model_name=model,
 max_tokens = 5000,
        openai_api_key=api_key,
-       base_url="https://www.dmxapi.com/v1"
+       base_url=os.getenv("OPENAI_API_BASE")
    )
     elif model.startswith("gpt") or model.startswith("deepseek"):
-        llm = ChatOpenAI(model=model,
-            temperature = 0.1,
-            
-            timeout=1000,
-            
-            callbacks=[StreamingStdOutCallbackHandler()],
-            openai_api_key=api_key,base_url="https://www.dmxapi.com/v1"
-            )
+        if os.getenv("OPENAI_API_BASE"):
+            llm = ChatOpenAI(model=model,
+                temperature = 0.1,
+                
+                timeout=1000,
+                
+                callbacks=[StreamingStdOutCallbackHandler()],
+                openai_api_key=api_key,base_url = os.getenv("OPENAI_API_BASE")
+                )
+        else: 
+            llm = ChatOpenAI(model=model,
+                temperature = 0.1,
+                
+                timeout=1000,
+                
+                callbacks=[StreamingStdOutCallbackHandler()],
+                openai_api_key=api_key 
+                )
+        
     elif model.startswith("llama") :
             llm = OllamaLLM(model=model,
                 temperature = 0.1,    
@@ -82,7 +93,7 @@ class TeLLAgent:
             raise ValueError("Invalid OpenAI API key")
 
         if tools is None:
-            api_keys["OPENAI_API_KEY"] = 'sk-itPrztYm9F6XZZpsBMJB9O7Vq0pYUABVVBSoThuBxEGTnDik'
+            api_keys["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")
             tools_llm = _make_llm(tools_model, temp, openai_api_key, streaming)
             tools = make_tools(tools_llm, api_keys=api_keys, verbose=verbose, image_path = image_path, file_path = file_path)