Update app.py

app.py

@@ -24,14 +24,9 @@ from langchain_openai import OpenAIEmbeddings, ChatOpenAI
 from langchain_community.vectorstores import Qdrant
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.embeddings import SentenceTransformerEmbeddings
-
+import nest_asyncio
 
 torch.cuda.empty_cache()
-
-
-
-
-import nest_asyncio
 nest_asyncio.apply()
 co = cohere.Client(st.secrets["COHERE_API_KEY"])
 
@@ -102,17 +97,40 @@ retriever = vector_store.as_retriever()
 
 
 
-# Dynamically initialize LLM based on selection
-OPENAI_API_KEY = st.secrets["OPENAI_API_KEY"]
-selected_model = st.session_state["selected_model"]
+#selected_model = st.session_state["selected_model"]
+
 if "OpenAI" in selected_model:
-    llm = ChatOpenAI(model="gpt-4o", temperature=0.2, api_key=OPENAI_API_KEY)
+    from langchain_openai import ChatOpenAI
+    llm = ChatOpenAI(model="gpt-4o", temperature=0.2, api_key=st.secrets["OPENAI_API_KEY"])
+
 elif "LLaMA" in selected_model:
-    st.warning("LLaMA integration is not implemented yet.")
-    st.stop()
+    from groq import Groq
+
+    client = Groq(api_key=st.secrets["GROQ_API_KEY"])  # Store in `.streamlit/secrets.toml`
+    def get_llama_response(prompt):
+        completion = client.chat.completions.create(
+            model="meta-llama/llama-4-maverick-17b-128e-instruct",
+            messages=[{"role": "user", "content": prompt}],
+            temperature=1,
+            max_completion_tokens=1024,
+            top_p=1,
+            stream=False
+        )
+        return completion.choices[0].message.content
+
+    llm = get_llama_response  # use this in place of llm.invoke()
+
 elif "Gemini" in selected_model:
-    st.warning("Gemini integration is not implemented yet.")
-    st.stop()
+    import google.generativeai as genai
+    genai.configure(api_key=st.secrets["GEMINI_API_KEY"])  # Store in `.streamlit/secrets.toml`
+
+    gemini_model = genai.GenerativeModel("gemini-2.5-pro-exp-03-25")
+    def get_gemini_response(prompt):
+        response = gemini_model.generate_content(prompt)
+        return response.text
+
+    llm = get_gemini_response
+
 else:
     st.error("Unsupported model selected.")
     st.stop()
@@ -290,7 +308,14 @@ def get_reranked_response(query: str):
     reranked_docs = rerank_with_cohere(query, docs)
     context = "\n\n".join([doc.page_content for doc in reranked_docs])
     prompt = AI_PROMPT_TEMPLATE.format(question=query, context=context)
-    return llm.invoke([{"role": "system", "content": prompt}])
+
+    if callable(llm):
+        # Gemini or LLaMA
+        return type("Obj", (), {"content": llm(prompt)})
+    else:
+        # OpenAI LangChain interface
+        return llm.invoke([{"role": "system", "content": prompt}])
+
 
 # === App UI ===