Update app.py

app.py

@@ -15,7 +15,11 @@ from langchain_openai import ChatOpenAI
 from langchain.prompts import ChatPromptTemplate
 
 DB_PATH = "json_vector.db"
+
+# Read API keys
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")  # NEW
+
 EMBEDDING_MODEL = "text-embedding-ada-002"
 
 if "ingested_batches" not in st.session_state:
@@ -226,7 +230,31 @@ prompt = ChatPromptTemplate.from_messages([
     ("human", "Here are the most relevant records:\n{context}\n\nQuestion: {question}")
 ])
 
-llm = ChatOpenAI(model="gpt-4.1", openai_api_key=OPENAI_API_KEY, temperature=0)
+# --- LLM PROVIDER SELECTION --- # NEW/MODIFIED FOR LLM SELECTION
+llm_provider = st.selectbox(
+    "Select LLM Provider",
+    options=["OpenAI GPT-4", "Mistral (OpenRouter)"],
+    index=0,
+    help="Choose which LLM to use for answering your questions."
+)
+
+def get_llm(llm_provider):
+    if llm_provider == "OpenAI GPT-4":
+        return ChatOpenAI(
+            model="gpt-4.1",
+            openai_api_key=OPENAI_API_KEY,
+            temperature=0,
+        )
+    else:  # "Mistral (OpenRouter)"
+        return ChatOpenAI(
+            model="mistralai/mistral-small-3.1-24b-instruct:free",  # Or another Mistral model if desired
+            openai_api_key=OPENROUTER_API_KEY,
+            openai_api_base="https://openrouter.ai/api/v1",
+            temperature=0,
+        )
+
+llm = get_llm(llm_provider)
+
 retriever = HybridRetriever(top_k=5)
 qa_chain = RetrievalQA.from_chain_type(
     llm=llm,