Update app.py

app.py

@@ -18,25 +18,30 @@ from sentence_transformers import CrossEncoder
 import google.generativeai as genai
 from typing import List
 from langchain_core.language_models import BaseLanguageModel
+from langchain_core.runnables import Runnable
 
 import google.generativeai as genai
 
 
-class GeminiLLM(BaseLanguageModel):
+class GeminiLLM(Runnable):
     def __init__(self, model_name="models/gemini-1.5-pro-latest", api_key=None):
         self.api_key = api_key or st.secrets["GOOGLE_API_KEY"]
         if not self.api_key:
-            raise ValueError("GOOGLE_API_KEY not found in Streamlit secrets.")
+            raise ValueError("GOOGLE_API_KEY not found.")
         genai.configure(api_key=self.api_key)
         self.model = genai.GenerativeModel(model_name)
-
-    def _call(self, prompt, stop=None):
+        
+    def _call(self, prompt: str, stop=None) -> str:
         response = self.model.generate_content(prompt)
         return response.text
 
     @property
-    def _llm_type(self):
+    def _llm_type(self) -> str:
         return "custom_gemini"
+    
+    def invoke(self, input, config=None):
+        response = self.model.generate_content(input)
+        return response.text.strip()
 
 class GeminiEmbeddings(Embeddings):
     def __init__(self, model_name="models/embedding-001", api_key=None):
@@ -64,19 +69,6 @@ class GeminiEmbeddings(Embeddings):
             task_type="retrieval_query"
         )["embedding"]
 
-
-class GeminiLLM:
-    def __init__(self, model_name="models/gemini-1.5-pro-latest", api_key=None):
-        api_key = api_key or os.getenv("GOOGLE_API_KEY")
-        if not api_key:
-            raise ValueError("Missing GOOGLE_API_KEY")
-        genai.configure(api_key=api_key)
-        self.model = genai.GenerativeModel(model_name)
-
-    def predict(self, prompt: str) -> str:
-        response = self.model.generate_content(prompt)
-        return response.text.strip()
-    
 reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
 
 vectorstore_global = None
@@ -88,8 +80,8 @@ def load_environment():
 def preload_modtran_document():
     global vectorstore_global
     embeddings = GeminiEmbeddings()
-    vectorstore = FAISS.load_local("modtran_vectorstore", embeddings, allow_dangerous_deserialization=True)
-    set_global_vectorstore(vectorstore)
+    st.session_state.vectorstore = FAISS.load_local("modtran_vectorstore", embeddings, allow_dangerous_deserialization=True)
+    set_global_vectorstore(st.session_state.vectorstore)
     st.session_state.chat_ready = True
 
 def convert_pdf_to_xml(pdf_file, xml_path):
@@ -183,7 +175,7 @@ def self_reasoning(query, context):
 
     **Answer:**
     """
-    return llm.predict(reasoning_prompt)
+    return llm._call(reasoning_prompt)
 
 def faiss_search_with_keywords(query):
     global vectorstore_global
@@ -222,7 +214,7 @@ faiss_reasoning_tool = Tool(
 )
 
 def initialize_chatbot_agent():
-    llm = GeminiLLM()  # <-- Gemini instead of OpenAI
+    llm = GeminiLLM() 
     memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
     tools = [faiss_keyword_tool, faiss_reasoning_tool]
     agent = initialize_agent(
@@ -250,12 +242,12 @@ def handle_user_query(query):
 def main():
     load_environment()
 
-    if "agent" not in st.session_state:
-        st.session_state.agent = None
     if "chat_ready" not in st.session_state:
         st.session_state.chat_ready = False
     if "chat_history" not in st.session_state:
         st.session_state.chat_history = []
+    if "vectorstore" not in st.session_state:
+        st.session_state.vectorstore = None
 
     st.header("Chat with MODTRAN Documents 📄")