From 40944dea1c03a1b594fcdee988b46449d75ba18e Mon Sep 17 00:00:00 2001
From: Patrick Frey <patrick-frey@gmx.de>
Date: Mon, 30 Dec 2024 11:27:28 +0100
Subject: [PATCH] Added environment variable to select different OpenAI models.

---
 .env.example                    | 1 +
 src/agents/market_data.py       | 9 ++++++++-
 src/agents/portfolio_manager.py | 8 +++++++-
 3 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/.env.example b/.env.example
index 1e321576..f7025698 100644
--- a/.env.example
+++ b/.env.example
@@ -1,2 +1,3 @@
 OPENAI_API_KEY=your_openai_api_key_here
 FINANCIAL_DATASETS_API_KEY=your_financial_datasets_api_key_here
+OPENAI_MODEL="gpt-4o" # Specify which OpenAI model to use. Examples: gpt-4, gpt-4o, gpt-3.5-turbo-0125 https://openai.com/api/pricing/
diff --git a/src/agents/market_data.py b/src/agents/market_data.py
index c8592e15..886a797b 100644
--- a/src/agents/market_data.py
+++ b/src/agents/market_data.py
@@ -6,7 +6,14 @@
 
 from datetime import datetime
 
-llm = ChatOpenAI(model="gpt-4o")
+from dotenv import load_dotenv
+import os
+
+# Load environment variables from the .env file
+load_dotenv()
+
+# Invoke the LLM
+llm = ChatOpenAI(model=os.getenv("OPENAI_MODEL"), api_key=os.getenv("OPENAI_API_KEY"))
 
 def market_data_agent(state: AgentState):
     """Responsible for gathering and preprocessing market data"""
diff --git a/src/agents/portfolio_manager.py b/src/agents/portfolio_manager.py
index e8a1af97..c4bee15e 100644
--- a/src/agents/portfolio_manager.py
+++ b/src/agents/portfolio_manager.py
@@ -4,6 +4,8 @@
 
 from agents.state import AgentState, show_agent_reasoning
 
+from dotenv import load_dotenv
+import os
 
 ##### Portfolio Management Agent #####
 def portfolio_management_agent(state: AgentState):
@@ -107,8 +109,12 @@ def portfolio_management_agent(state: AgentState):
             "portfolio_stock": portfolio["stock"]
         }
     )
+
+    # Load environment variables from the .env file
+    load_dotenv()
+
     # Invoke the LLM
-    llm = ChatOpenAI(model="gpt-4o")
+    llm = ChatOpenAI(model=os.getenv("OPENAI_MODEL"), api_key=os.getenv("OPENAI_API_KEY"))
     result = llm.invoke(prompt)
 
     # Create the portfolio management message