fix

server/common/agent.py

@@ -13,22 +13,8 @@
 from langchain_openai import ChatOpenAI
 from pydantic import BaseModel
 
+from server.common.openai_utils import get_cached_openai_api_key
 from server.common.tools import create_tools
-from server.common.utils.aws_secret_utils import get_secret_key
-
-
-def get_openai_api_key():
-    deployment_stage = os.getenv("DEPLOYMENT_STAGE", "test")
-    if deployment_stage == "test":
-        return os.getenv("OPENAI_API_KEY")
-    return get_secret_key(f"{deployment_stage}/explorer/agent", region_name="us-west-2")["OPENAI_API_KEY"]
-
-
-def get_cached_openai_api_key():
-    """Get OpenAI API key from cache or fetch it if not cached"""
-    if not hasattr(current_app, "openai_api_key"):
-        current_app.openai_api_key = get_openai_api_key()
-    return current_app.openai_api_key
 
 
 class AgentMessage(BaseModel):

server/common/tools.py

@@ -5,6 +5,7 @@
 from langchain_core.tools import Tool
 from langchain_openai import ChatOpenAI
 from pydantic import BaseModel, Field
+from server.common.openai_utils import get_cached_openai_api_key
 
 
 class ColorByGeneSchema(BaseModel):
@@ -304,7 +305,8 @@ def no_more_steps(summary: str):
 
 
 def call_llm_with_structured_output(query: str, schema: Type[T]) -> T:
-    llm = ChatOpenAI(temperature=0, model_name="gpt-4o").with_structured_output(schema)
+    api_key = get_cached_openai_api_key()
+    llm = ChatOpenAI(temperature=0, model_name="gpt-4o", api_key=api_key).with_structured_output(schema)
     return llm.invoke(query).model_dump()