Update app.py

app.py

@@ -29,6 +29,9 @@ 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"
+
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -69,22 +72,13 @@ 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")
     
-    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"] = "sk-proj-iQ8piK0xt54XBKEI4nDzCg9CZE7a13xZqCaN1B78zqZTyhBwrXOCjfMjNWG0w1gprhdj6_"
-            
+    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
+    
     st.divider()
     
     # Reset button
@@ -93,7 +87,6 @@ with st.sidebar:
             del st.session_state[key]
         st.rerun()
 
-
 #---------------------------- Utility Functions ----------------------------#
 
 def extract_text_from_pdf(file):
@@ -398,6 +391,8 @@ class CaseBreakdownCrew:
         self.api_key = api_key
         
     def create_metadata_agent(self):
+        self.api_key = api_key
+        self.llm = LLM(model='gemini/gemini-2.0-flash', api_key=self.api_key)  # Create a Gemini LLM instance
         return Agent(
             role="Metadata Analyzer",
             goal="Extract title and author information from document content",
@@ -409,6 +404,7 @@ 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",
@@ -420,6 +416,7 @@ 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",
@@ -505,6 +502,7 @@ 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()
@@ -607,16 +605,9 @@ def create_teaching_plan_crew(file_paths, llm_provider="gemini"):
     tracker.set_placeholder(st.empty())
     
     # Initialize LLM based on provider
-    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")
-        )
+    my_llm = LLM(model='gemini/gemini-2.0-flash',
+        api_key=os.environ.get("GEMINI_API_KEY")
+    )
     
     # Define agents with callbacks for UI updates
     pdf_analyzer = Agent(
@@ -708,24 +699,15 @@ def create_teaching_plan_crew(file_paths, llm_provider="gemini"):
 #---------------------------- Board Plan Generator ----------------------------#
 
 class BoardPlanAnalyzer:
-    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"
-            
+    def __init__(self):
+        api_key = os.environ.get('GEMINI_API_KEY')
         if not api_key:
-            raise ValueError(f"{llm_provider.capitalize()} API key not found")
-            
-        if llm_provider == "gemini":
-            os.environ['GEMINI_API_KEY'] = api_key
-        else:
-            os.environ['OPENAI_API_KEY'] = api_key
+            raise ValueError("Gemini API key not found")
+        # Create an LLM instance configured for Gemini
+        self.llm = LLM(model='gemini/gemini-2.0-flash', api_key=api_key)
             
         litellm.set_verbose = True
-        
+
         # Create agents
         self.create_agents()
 
@@ -745,6 +727,7 @@ class BoardPlanAnalyzer:
                 description="Extracts text content from PDF files"
             )],
             allow_delegation=False,
+            llm=self.llm,
             verbose=True
         )
 
@@ -761,6 +744,7 @@ class BoardPlanAnalyzer:
                 description="Analyzes case study and creates structured board plan"
             )],
             allow_delegation=False,
+            llm=self.llm,
             verbose=True
         )
 
@@ -823,7 +807,7 @@ class BoardPlanAnalyzer:
         
         try:
             response = litellm.completion(
-                model=self.model,
+                model=self.llm,
                 messages=messages,
                 response_format={"type": "json_object"}
             )
@@ -1098,7 +1082,7 @@ if st.session_state.uploaded_files:
                         progress_bar = progress_placeholder.progress(0)
                         
                         # Select LLM provider
-                        llm_provider = "gemini" if api_key_source == "Google (Gemini)" else "openai"
+                        llm_provider = "gemini"
                         
                         # Update progress
                         progress_bar.progress(10)
@@ -1189,7 +1173,7 @@ if st.session_state.uploaded_files:
                 if st.button("Generate Board Plan", key="board_plan_button"):
                     try:
                         # Select LLM provider
-                        llm_provider = "gemini" if api_key_source == "Google (Gemini)" else "openai"
+                        llm_provider = "gemini"
                         
                         # Initialize the board plan analyzer
                         analyzer = BoardPlanAnalyzer(llm_provider=llm_provider)