Add Gemini provider selection

README.md

@@ -15,8 +15,8 @@ Upload PDFs, index them, and chat with a multi-agent RAG workflow.
 
 1. Create a new Space and choose **Docker**.
 2. Upload this repository contents.
-3. Add a secret named `GROQ_API_KEY` in **Settings → Secrets**.
-4. The app will start automatically.
+3. Add a secret named `GROQ_API_KEY` or `GEMINI_API_KEY` in **Settings → Secrets**.
+4. Choose the provider and model in the app sidebar.
 
 ## Notes
 

agents/base_agent.py

@@ -1,7 +1,14 @@
 import os
 from pathlib import Path
 from langchain_groq import ChatGroq
-from config import GROQ_API_KEY, LLM_MODEL
+from langchain_google_genai import ChatGoogleGenerativeAI
+from config import (
+    GROQ_API_KEY,
+    GEMINI_API_KEY,
+    LLM_MODEL,
+    DEFAULT_PROVIDER,
+    DEFAULT_MODEL,
+)
 from utils import get_logger
 
 logger = get_logger(__name__)
@@ -16,23 +23,51 @@ class BaseAgent:
     provides a ChatGroq client.
     """
 
-    def __init__(self, prompt_file: str, temperature: float = 0.0, max_tokens: int = 512):
-        if not GROQ_API_KEY:
-            raise EnvironmentError(
-                "GROQ_API_KEY is not set. "
-                "Add it to your .env file or Streamlit secrets."
+    def __init__(
+        self,
+        prompt_file: str,
+        temperature: float = 0.0,
+        max_tokens: int = 512,
+        model_provider: str | None = None,
+        model_name: str | None = None,
+    ):
+        provider = (model_provider or DEFAULT_PROVIDER).lower()
+        model = model_name or DEFAULT_MODEL or LLM_MODEL
+
+        if provider == "groq":
+            if not GROQ_API_KEY:
+                raise EnvironmentError(
+                    "GROQ_API_KEY is not set. "
+                    "Add it to your .env file or Streamlit secrets."
+                )
+
+            self.llm = ChatGroq(
+                model_name=model,
+                temperature=temperature,
+                max_tokens=max_tokens,
+                groq_api_key=GROQ_API_KEY,
             )
+        elif provider == "gemini":
+            if not GEMINI_API_KEY:
+                raise EnvironmentError(
+                    "GEMINI_API_KEY is not set. "
+                    "Add it to your .env file or Streamlit secrets."
+                )
 
-        self.llm = ChatGroq(
-            model_name=LLM_MODEL,
-            temperature=temperature,
-            max_tokens=max_tokens,
-            groq_api_key=GROQ_API_KEY,
-        )
+            self.llm = ChatGoogleGenerativeAI(
+                model=model,
+                temperature=temperature,
+                max_output_tokens=max_tokens,
+                google_api_key=GEMINI_API_KEY,
+            )
+        else:
+            raise ValueError(f"Unknown model provider: {provider}")
 
         prompt_path = PROMPT_DIR / prompt_file
         self.prompt_template = prompt_path.read_text(encoding="utf-8")
-        logger.info(f"{self.__class__.__name__} ready (model={LLM_MODEL})")
+        logger.info(
+            f"{self.__class__.__name__} ready (provider={provider}, model={model})"
+        )
 
     def _call_llm(self, prompt: str) -> str:
         response = self.llm.invoke(prompt)

agents/relevance_agent.py

@@ -13,8 +13,14 @@ class RelevanceAgent(BaseAgent):
     taking conversation history into account.
     """
 
-    def __init__(self):
-        super().__init__(prompt_file="relevance.txt", temperature=0.0, max_tokens=10)
+    def __init__(self, model_provider: str | None = None, model_name: str | None = None):
+        super().__init__(
+            prompt_file="relevance.txt",
+            temperature=0.0,
+            max_tokens=10,
+            model_provider=model_provider,
+            model_name=model_name,
+        )
 
     def check(self, question: str, documents: list[Document], history: str) -> str:
         """

agents/research_agent.py

@@ -13,8 +13,14 @@ class ResearchAgent(BaseAgent):
     Also performs query rewriting when history is present.
     """
 
-    def __init__(self):
-        super().__init__(prompt_file="research.txt", temperature=0.1, max_tokens=600)
+    def __init__(self, model_provider: str | None = None, model_name: str | None = None):
+        super().__init__(
+            prompt_file="research.txt",
+            temperature=0.1,
+            max_tokens=600,
+            model_provider=model_provider,
+            model_name=model_name,
+        )
 
         # Load query-rewrite prompt from same prompts/ directory
         from pathlib import Path

agents/verification_agent.py

@@ -10,8 +10,14 @@ class VerificationAgent(BaseAgent):
     Checks whether the draft answer is grounded in the retrieved documents.
     """
 
-    def __init__(self):
-        super().__init__(prompt_file="verification.txt", temperature=0.0, max_tokens=220)
+    def __init__(self, model_provider: str | None = None, model_name: str | None = None):
+        super().__init__(
+            prompt_file="verification.txt",
+            temperature=0.0,
+            max_tokens=220,
+            model_provider=model_provider,
+            model_name=model_name,
+        )
 
     def check(self, answer: str, documents: list[Document]) -> dict:
         """

app.py

@@ -4,7 +4,16 @@ import streamlit as st
 from ingestion import ingest_pdfs
 from retriever import HybridRetriever
 from graph import AgentWorkflow, Turn
-from config import UPLOAD_DIR, INDEX_DIR
+from config import (
+    UPLOAD_DIR,
+    INDEX_DIR,
+    GROQ_FREE_MODELS,
+    GEMINI_FREE_MODELS,
+    DEFAULT_PROVIDER,
+    DEFAULT_MODEL,
+    GROQ_API_KEY,
+    GEMINI_API_KEY,
+)
 
 # ─────────────────────────────────────────────────────────────────────────────
 # Page config
@@ -27,6 +36,26 @@ with st.sidebar:
         "🔍 **Verification Mode**: ~6–10 s — checks answer quality"
     )
     st.divider()
+    st.subheader("Model")
+
+    provider_labels = ["Groq", "Gemini"]
+    provider_index = 0 if st.session_state.model_provider == "groq" else 1
+    provider_label = st.selectbox("Provider", provider_labels, index=provider_index)
+    model_provider = provider_label.lower()
+
+    model_options = GROQ_FREE_MODELS if model_provider == "groq" else GEMINI_FREE_MODELS
+    if st.session_state.model_name not in model_options:
+        st.session_state.model_name = model_options[0]
+
+    model_name = st.selectbox(
+        "Model",
+        model_options,
+        index=model_options.index(st.session_state.model_name),
+    )
+
+    st.session_state.model_provider = model_provider
+    st.session_state.model_name = model_name
+    st.divider()
     st.caption("Conversation memory: last **4** Q&A pairs")
 
 # ─────────────────────────────────────────────────────────────────────────────
@@ -40,6 +69,8 @@ defaults = {
     "retriever":             None,
     "files_indexed":         False,
     "uploaded_file_names":   set(),
+    "model_provider":        DEFAULT_PROVIDER,
+    "model_name":            DEFAULT_MODEL,
 }
 for key, val in defaults.items():
     if key not in st.session_state:
@@ -127,6 +158,13 @@ for msg in st.session_state.chat_history:
 question = st.chat_input("Ask a question about your uploaded PDFs…")
 
 if question:
+    if st.session_state.model_provider == "groq" and not GROQ_API_KEY:
+        st.error("GROQ_API_KEY is not set. Add it to your secrets or .env file.")
+        st.stop()
+    if st.session_state.model_provider == "gemini" and not GEMINI_API_KEY:
+        st.error("GEMINI_API_KEY is not set. Add it to your secrets or .env file.")
+        st.stop()
+
     if not os.path.exists(INDEX_DIR) or not os.listdir(INDEX_DIR):
         st.warning("⚠️ Please upload and index PDFs first.")
         st.stop()
@@ -144,6 +182,8 @@ if question:
             question=question,
             retriever=st.session_state.retriever,
             conversation_history=st.session_state.conversation_history,
+            model_provider=st.session_state.model_provider,
+            model_name=st.session_state.model_name,
         )
 
     # ── Persist updated history window back to session ────────────────────

config.py

@@ -28,7 +28,22 @@ BATCH_SIZE    = 1000   # nodes per indexing batch for large PDFs
 
 # ── LLM ──────────────────────────────────────────────────────────────────────
 GROQ_API_KEY  = os.getenv("GROQ_API_KEY")
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
 LLM_MODEL     = "llama-3.1-8b-instant"
 
+GROQ_FREE_MODELS = [
+	"llama-3.1-8b-instant",
+	"llama-3.1-70b-versatile",
+	"mixtral-8x7b-32768",
+]
+
+GEMINI_FREE_MODELS = [
+	"gemini-1.5-flash",
+	"gemini-1.5-flash-8b",
+]
+
+DEFAULT_PROVIDER = "groq"
+DEFAULT_MODEL = GROQ_FREE_MODELS[0]
+
 # ── Workflow ──────────────────────────────────────────────────────────────────
 MAX_ITERATIONS = 2     # max research→verify loops before forcing end

graph/nodes.py

@@ -39,7 +39,10 @@ def rewrite_query_node(state: AgentState) -> dict:
 
     logger.info("Node: rewrite_query")
     history = _format_history(state)
-    agent = ResearchAgent()
+    agent = ResearchAgent(
+        model_provider=state.get("model_provider"),
+        model_name=state.get("model_name"),
+    )
     rewritten = agent.rewrite_query(state["question"], history)
     return {"rewritten_query": rewritten}
 
@@ -55,7 +58,10 @@ def check_relevance_node(state: AgentState) -> dict:
 
     logger.info("Node: check_relevance")
     history = _format_history(state)
-    agent = RelevanceAgent()
+    agent = RelevanceAgent(
+        model_provider=state.get("model_provider"),
+        model_name=state.get("model_name"),
+    )
 
     label = agent.check(
         question=state["rewritten_query"],
@@ -85,7 +91,10 @@ def research_node(state: AgentState) -> dict:
 
     logger.info("Node: research")
     history = _format_history(state)
-    agent = ResearchAgent()
+    agent = ResearchAgent(
+        model_provider=state.get("model_provider"),
+        model_name=state.get("model_name"),
+    )
 
     result = agent.generate(
         question=state["rewritten_query"],
@@ -110,7 +119,10 @@ def verify_node(state: AgentState) -> dict:
     from agents.verification_agent import VerificationAgent
 
     logger.info("Node: verify")
-    agent = VerificationAgent()
+    agent = VerificationAgent(
+        model_provider=state.get("model_provider"),
+        model_name=state.get("model_name"),
+    )
     result = agent.check(
         answer=state["draft_answer"],
         documents=state["documents"],

graph/state.py

@@ -32,3 +32,5 @@ class AgentState(TypedDict):
     retriever: Any                   # HybridRetriever instance (passed through)
     iteration_count: int             # tracks research→verify loops
     enable_verification: bool        # toggle slower verification path
+    model_provider: str              # "groq" | "gemini"
+    model_name: str                  # selected model name

graph/workflow.py

@@ -132,6 +132,8 @@ class AgentWorkflow:
         question: str,
         retriever: Any,
         conversation_history: list[Turn] | None = None,
+        model_provider: str | None = None,
+        model_name: str | None = None,
     ) -> dict:
         """
         Run the full pipeline for one user turn.
@@ -183,6 +185,8 @@ class AgentWorkflow:
             "retriever": retriever,
             "iteration_count": 0,
             "enable_verification": self.enable_verification,
+            "model_provider": model_provider or "groq",
+            "model_name": model_name or "",
         }
 
         try:

requirements.txt

@@ -14,3 +14,4 @@ pypdf>=4.2.0
 langchain>=0.1.20
 langgraph>=0.0.40
 langchain-groq>=0.1.4
+langchain-google-genai>=1.0.7