Update app.py

app.py

@@ -29,11 +29,6 @@ from pydantic import BaseModel, Field
 import litellm
 from langchain.tools import Tool
 
-LLM._get_litellm_model_name = lambda self, model_name: f"gemini/{model_name}" if not "/" in model_name else model_name
-os.environ["LITELLM_MODEL_DEFAULT_PROVIDER"] = "gemini"
-
-os.environ["OPENAI_API_KEY"] = "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"
-
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -74,13 +69,22 @@ st.write("---")
 # Sidebar for API key configuration
 with st.sidebar:
     st.title("⚙️ Configuration")
+    api_key_source = st.radio("Select API Key Provider:", 
+                            ["Google (Gemini)", "OpenAI"],
+                            help="Choose which AI provider to use")
     
-    api_key = st.text_input("Enter your Gemini API Key", type="password", 
-                          help="Required for the AI model to function")
-    if api_key:
-        os.environ["GEMINI_API_KEY"] = api_key
-        os.environ["GOOGLE_API_KEY"] = api_key
-    
+    if api_key_source == "Google (Gemini)":
+        api_key = st.text_input("Enter your Gemini API Key", type="password", 
+                                help="Required for the AI model to function")
+        if api_key:
+            os.environ["GEMINI_API_KEY"] = api_key
+            os.environ["GOOGLE_API_KEY"] = api_key
+    else:
+        api_key = st.text_input("Enter your OpenAI API Key", type="password", 
+                                help="Required for the AI model to function")
+        if api_key:
+            os.environ["OPENAI_API_KEY"] = api_key
+            
     st.divider()
     
     # Reset button
@@ -89,6 +93,7 @@ with st.sidebar:
             del st.session_state[key]
         st.rerun()
 
+
 #---------------------------- Utility Functions ----------------------------#
 
 def extract_text_from_pdf(file):
@@ -393,7 +398,6 @@ class CaseBreakdownCrew:
         self.api_key = api_key
         
     def create_metadata_agent(self):
-        llm = LLM(model='gemini/gemini-2.0-flash', api_key=self.api_key)
         return Agent(
             role="Metadata Analyzer",
             goal="Extract title and author information from document content",
@@ -405,7 +409,6 @@ class CaseBreakdownCrew:
         )
     
     def create_content_generator_agent(self):
-        llm = LLM(model='gemini/gemini-2.0-flash', api_key=self.api_key)
         return Agent(
             role="Case Study Content Generator",
             goal="Generate comprehensive case analysis content based on section requirements",
@@ -417,7 +420,6 @@ class CaseBreakdownCrew:
         )
     
     def create_content_reviewer_agent(self):
-        llm = LLM(model='gemini/gemini-2.0-flash', api_key=self.api_key)
         return Agent(
             role="Content Quality Reviewer",
             goal="Evaluate and score content for quality, relevance, and depth",
@@ -503,7 +505,6 @@ class CaseBreakdownCrew:
             agents=[self.create_metadata_agent()],
             tasks=[metadata_task],
             process=Process.sequential,
-            llm=LLM(model='gemini/gemini-2.0-flash', api_key=self.api_key),
             verbose=False
         )
         result = crew.kickoff()
@@ -606,9 +607,16 @@ def create_teaching_plan_crew(file_paths, llm_provider="gemini"):
     tracker.set_placeholder(st.empty())
     
     # Initialize LLM based on provider
-    my_llm = LLM(model='gemini/gemini-2.0-flash',
-        api_key=os.environ.get("GEMINI_API_KEY")
-    )
+    if llm_provider == "gemini":
+        my_llm = LLM(
+            model='gemini/gemini-2.0-flash',
+            api_key=os.environ.get("GEMINI_API_KEY")
+        )
+    else:
+        my_llm = LLM(
+            model='gpt-4-turbo',
+            api_key=os.environ.get("OPENAI_API_KEY")
+        )
     
     # Define agents with callbacks for UI updates
     pdf_analyzer = Agent(
@@ -700,20 +708,24 @@ def create_teaching_plan_crew(file_paths, llm_provider="gemini"):
 #---------------------------- Board Plan Generator ----------------------------#
 
 class BoardPlanAnalyzer:
-    def __init__(self):
-        api_key = os.environ.get('GEMINI_API_KEY')
-        self.model = "gemini/gemini-2.0-flash"
+    def __init__(self, llm_provider="gemini"):
+        if llm_provider == "gemini":
+            api_key = os.environ.get('GEMINI_API_KEY')
+            self.model = "gemini/gemini-2.0-flash"
+        else:
+            api_key = os.environ.get('OPENAI_API_KEY')
+            self.model = "gpt-4-turbo"
             
         if not api_key:
-            raise ValueError("Gemini API key not found")
+            raise ValueError(f"{llm_provider.capitalize()} API key not found")
             
-        os.environ['GEMINI_API_KEY'] = api_key
-        self.api_key = api_key
+        if llm_provider == "gemini":
+            os.environ['GEMINI_API_KEY'] = api_key
+        else:
+            os.environ['OPENAI_API_KEY'] = api_key
             
         litellm.set_verbose = True
-
-        os.environ["LITELLM_MODEL_DEFAULT_PROVIDER"] = "gemini"
-
+        
         # Create agents
         self.create_agents()
 
@@ -733,7 +745,6 @@ class BoardPlanAnalyzer:
                 description="Extracts text content from PDF files"
             )],
             allow_delegation=False,
-            llm=self.model,
             verbose=True
         )
 
@@ -750,7 +761,6 @@ class BoardPlanAnalyzer:
                 description="Analyzes case study and creates structured board plan"
             )],
             allow_delegation=False,
-            llm=self.model,
             verbose=True
         )
 
@@ -1088,7 +1098,7 @@ if st.session_state.uploaded_files:
                         progress_bar = progress_placeholder.progress(0)
                         
                         # Select LLM provider
-                        llm_provider = "gemini"
+                        llm_provider = "gemini" if api_key_source == "Google (Gemini)" else "openai"
                         
                         # Update progress
                         progress_bar.progress(10)
@@ -1179,7 +1189,7 @@ if st.session_state.uploaded_files:
                 if st.button("Generate Board Plan", key="board_plan_button"):
                     try:
                         # Select LLM provider
-                        llm_provider = "gemini"
+                        llm_provider = "gemini" if api_key_source == "Google (Gemini)" else "openai"
                         
                         # Initialize the board plan analyzer
                         analyzer = BoardPlanAnalyzer(llm_provider=llm_provider)