fast/thinking toggle

app.py

@@ -382,20 +382,28 @@ def ingest_url_ui(
 def send_chat_ui(
     notebook_id: str | None,
     question: str,
+    rag_mode: str,
     history: list[dict[str, str]] | None,
     current_username: str,
     profile: gr.OAuthProfile | None,
     request: gr.Request,
-) -> tuple[list[dict[str, str]], str, str]:
+) -> tuple[str, list[dict[str, str]]]:
     """Send one chat question and append the grounded answer to the chat history."""
 
     username: str = _resolve_username(profile, request, current_username)
     if not notebook_id:
-        raise gr.Error("Select a notebook before asking a question.")
+        raise gr.Error("Select a notebook before sending a message.")
     if not question or not question.strip():
-        raise gr.Error("Enter a question before sending.")
+        raise gr.Error("Message cannot be empty.")
+
+    chat_history: list[dict[str, str]] = history or []
+    try:
+        response: ChatResponse = answer_question(username, notebook_id, question.strip(), rag_mode)
+    except Exception as e:
+        chat_history.append({"role": "user", "content": question.strip()})
+        chat_history.append({"role": "assistant", "content": f"Error: {e}"})
+        return "", chat_history
 
-    response: ChatResponse = answer_question(username, notebook_id, question.strip())
     updated_history: list[dict[str, str]] = list(history or [])
     updated_history.append({"role": "user", "content": question.strip()})
     updated_history.append(
@@ -404,7 +412,7 @@ def send_chat_ui(
             "content": response["content"] + _render_citations(response["citations"]),
         }
     )
-    return updated_history, "", f"Question answered with {len(response['citations'])} citations."
+    return "", updated_history
 
 
 def _append_artifact_path(current_paths: list[str] | None, artifact: ArtifactRef) -> tuple[list[str], gr.Dropdown]:
@@ -528,9 +536,24 @@ with gr.Blocks(title="NotebookLM Clone") as demo:
 
         with gr.Column():
             gr.Markdown("## Chat")
-            chat_history = gr.Chatbot(label="Grounded Chat")
-            question_input = gr.Textbox(label="Question", placeholder="Ask about this notebook")
-            ask_button = gr.Button("Ask")
+            chat_history = gr.Chatbot(
+                elem_id="chat-history",
+                show_label=False,
+            )
+            with gr.Row():
+                chat_input = gr.Textbox(
+                    show_label=False,
+                    placeholder="Ask a question about your sources...",
+                    scale=4,
+                )
+                rag_mode = gr.Radio(
+                    choices=["Fast", "Reasoning"],
+                    value="Reasoning",
+                    label="RAG Mode",
+                    scale=1,
+                    interactive=True,
+                )
+            chat_submit = gr.Button("Send", variant="primary")
 
         with gr.Column():
             gr.Markdown("## Artifacts")
@@ -580,10 +603,15 @@ with gr.Blocks(title="NotebookLM Clone") as demo:
         outputs=[ingest_status, uploaded_docs_state, uploaded_docs_display],
     )
 
-    ask_button.click(
+    chat_submit.click(
+        send_chat_ui,
+        inputs=[notebook_dropdown, chat_input, rag_mode, chat_history, username_state],
+        outputs=[chat_input, chat_history],
+    )
+    chat_input.submit(
         send_chat_ui,
-        inputs=[notebook_dropdown, question_input, chat_history, username_state],
-        outputs=[chat_history, question_input, activity_status],
+        inputs=[notebook_dropdown, chat_input, rag_mode, chat_history, username_state],
+        outputs=[chat_input, chat_history],
     )
 
     report_button.click(

src/notebooklm_clone/chat.py

@@ -240,7 +240,7 @@ def _generate_answer(question: str, context: str) -> str:
     raise ChatGenerationError("Chat model returned an empty response.")
 
 
-def answer_question(username: str, notebook_id: str, question: str) -> ChatResponse:
+def answer_question(username: str, notebook_id: str, question: str, rag_mode: str = "Reasoning") -> ChatResponse:
     """Answer a notebook question using retrieved chunks and inline citations.
 
     Spec references:
@@ -270,6 +270,7 @@ def answer_question(username: str, notebook_id: str, question: str) -> ChatRespo
             notebook_id=notebook_id,
             query=normalized_question,
             k=_RETRIEVAL_K,
+            rag_mode=rag_mode,
         )
 
         if not retrieved_chunks:

src/notebooklm_clone/retrieval.py

@@ -418,6 +418,7 @@ def retrieve(
     notebook_id: str,
     query: str,
     k: int,
+    rag_mode: str = "Reasoning",
 ) -> list[RetrievalResult]:
     """Retrieve top notebook chunks with hybrid scoring, query expansion, and reranking.
 
@@ -459,7 +460,7 @@ def retrieve(
         }
 
         # Query expansion: generate alt phrasings and merge scores
-        queries: list[str] = _expand_query(query)
+        queries: list[str] = _expand_query(query) if rag_mode == "Reasoning" else [query]
         bm25_raw, vector_raw = _multi_query_scores(
             chunk_documents, collection, queries, len(ids)
         )
@@ -515,10 +516,13 @@ def retrieve(
 
         ranked_results.sort(key=lambda item: (-item["score"], item["chunk_id"]))
 
-        # Rerank only top-N candidates to control latency (default: 10)
-        _rerank_n: int = int(os.getenv("NOTEBOOKLM_RERANK_TOP_N", "10"))
-        rerank_pool: list[RetrievalResult] = ranked_results[:_rerank_n]
-        result: list[RetrievalResult] = _rerank(query, rerank_pool, k)
+        if rag_mode == "Fast":
+            result: list[RetrievalResult] = ranked_results[:k]
+        else:
+            # Rerank only top-N candidates to control latency (default: 10)
+            _rerank_n: int = int(os.getenv("NOTEBOOKLM_RERANK_TOP_N", "10"))
+            rerank_pool: list[RetrievalResult] = ranked_results[:_rerank_n]
+            result: list[RetrievalResult] = _rerank(query, rerank_pool, k)
 
         _log_retrieval(username, notebook_id, "success", started_at)
         return result