refactor: Switch RAG source from static path to dynamic download

.gradio/certificate.pem

@@ -0,0 +1,31 @@
+-----BEGIN CERTIFICATE-----
+MIIFazCCA1OgAwIBAgIRAIIQz7DSQONZRGPgu2OCiwAwDQYJKoZIhvcNAQELBQAw
+TzELMAkGA1UEBhMCVVMxKTAnBgNVBAoTIEludGVybmV0IFNlY3VyaXR5IFJlc2Vh
+cmNoIEdyb3VwMRUwEwYDVQQDEwxJU1JHIFJvb3QgWDEwHhcNMTUwNjA0MTEwNDM4
+WhcNMzUwNjA0MTEwNDM4WjBPMQswCQYDVQQGEwJVUzEpMCcGA1UEChMgSW50ZXJu
+ZXQgU2VjdXJpdHkgUmVzZWFyY2ggR3JvdXAxFTATBgNVBAMTDElTUkcgUm9vdCBY
+MTCCAiIwDQYJKoZIhvcNAQEBBQADggIPADCCAgoCggIBAK3oJHP0FDfzm54rVygc
+h77ct984kIxuPOZXoHj3dcKi/vVqbvYATyjb3miGbESTtrFj/RQSa78f0uoxmyF+
+0TM8ukj13Xnfs7j/EvEhmkvBioZxaUpmZmyPfjxwv60pIgbz5MDmgK7iS4+3mX6U
+A5/TR5d8mUgjU+g4rk8Kb4Mu0UlXjIB0ttov0DiNewNwIRt18jA8+o+u3dpjq+sW
+T8KOEUt+zwvo/7V3LvSye0rgTBIlDHCNAymg4VMk7BPZ7hm/ELNKjD+Jo2FR3qyH
+B5T0Y3HsLuJvW5iB4YlcNHlsdu87kGJ55tukmi8mxdAQ4Q7e2RCOFvu396j3x+UC
+B5iPNgiV5+I3lg02dZ77DnKxHZu8A/lJBdiB3QW0KtZB6awBdpUKD9jf1b0SHzUv
+KBds0pjBqAlkd25HN7rOrFleaJ1/ctaJxQZBKT5ZPt0m9STJEadao0xAH0ahmbWn
+OlFuhjuefXKnEgV4We0+UXgVCwOPjdAvBbI+e0ocS3MFEvzG6uBQE3xDk3SzynTn
+jh8BCNAw1FtxNrQHusEwMFxIt4I7mKZ9YIqioymCzLq9gwQbooMDQaHWBfEbwrbw
+qHyGO0aoSCqI3Haadr8faqU9GY/rOPNk3sgrDQoo//fb4hVC1CLQJ13hef4Y53CI
+rU7m2Ys6xt0nUW7/vGT1M0NPAgMBAAGjQjBAMA4GA1UdDwEB/wQEAwIBBjAPBgNV
+HRMBAf8EBTADAQH/MB0GA1UdDgQWBBR5tFnme7bl5AFzgAiIyBpY9umbbjANBgkq
+hkiG9w0BAQsFAAOCAgEAVR9YqbyyqFDQDLHYGmkgJykIrGF1XIpu+ILlaS/V9lZL
+ubhzEFnTIZd+50xx+7LSYK05qAvqFyFWhfFQDlnrzuBZ6brJFe+GnY+EgPbk6ZGQ
+3BebYhtF8GaV0nxvwuo77x/Py9auJ/GpsMiu/X1+mvoiBOv/2X/qkSsisRcOj/KK
+NFtY2PwByVS5uCbMiogziUwthDyC3+6WVwW6LLv3xLfHTjuCvjHIInNzktHCgKQ5
+ORAzI4JMPJ+GslWYHb4phowim57iaztXOoJwTdwJx4nLCgdNbOhdjsnvzqvHu7Ur
+TkXWStAmzOVyyghqpZXjFaH3pO3JLF+l+/+sKAIuvtd7u+Nxe5AW0wdeRlN8NwdC
+jNPElpzVmbUq4JUagEiuTDkHzsxHpFKVK7q4+63SM1N95R1NbdWhscdCb+ZAJzVc
+oyi3B43njTOQ5yOf+1CceWxG1bQVs5ZufpsMljq4Ui0/1lvh+wjChP4kqKOJ2qxq
+4RgqsahDYVvTH9w7jXbyLeiNdd8XM2w9U/t7y0Ff/9yi0GE44Za4rF2LN9d11TPA
+mRGunUHBcnWEvgJBQl9nJEiU0Zsnvgc/ubhPgXRR4Xq37Z0j4r7g1SgEEzwxA57d
+emyPxgcYxn/eR44/KJ4EBs+lVDR3veyJm+kXQ99b21/+jh5Xos1AnX5iItreGCc=
+-----END CERTIFICATE-----

app.py

@@ -1,186 +1,219 @@
-import torch
-import transformers
 import gradio as gr
+import google.generativeai as genai
+from google.generativeai.types import generation_types
 from ragatouille import RAGPretrainedModel
-from huggingface_hub import InferenceClient
+import arxiv
+import os
 import re
 from datetime import datetime
-import json
-import os
-import arxiv
-from utils import get_md_text_abstract, search_cleaner, get_arxiv_live_search
-
-retrieve_results = 10 
-show_examples = False
-llm_models_to_choose = ['mistralai/Mixtral-8x7B-Instruct-v0.1','mistralai/Mistral-7B-Instruct-v0.3', 'google/gemma-2-2b-it', 'None']
-
-token = os.getenv("HF_TOKEN")
-
-generate_kwargs = dict(
-    temperature = None,
-    max_new_tokens = 512,
-    top_p = None,
-    do_sample = False,
-    )
+from utils import get_md_text_abstract
+from huggingface_hub import snapshot_download
+
+# --- Core Configuration ---
+hf_token = os.getenv("HF_TOKEN")
+gemini_api_key = os.getenv("GEMINI_API_KEY")
+RAG_SOURCE = os.getenv("RAG_SOURCE")
+LOCAL_DATA_DIR = './rag_index_data'
+
+LLM_MODELS_TO_CHOOSE = [
+    'google/gemma-3-4b-it',
+    'google/gemma-3-12b-it',
+    'google/gemma-3-27b-it',
+    'None'
+]
+DEFAULT_LLM_MODEL = 'google/gemma-3-4b-it'
+RETRIEVE_RESULTS = 20
+
+# --- Gemini API Configuration ---
+if gemini_api_key:
+    genai.configure(api_key=gemini_api_key)
+else:
+    print("CRITICAL WARNING: GEMINI_API_KEY environment variable not set. The application will not function without it.")
 
-## RAG Model
-RAG = RAGPretrainedModel.from_index("colbert/indexes/arxiv_colbert")
+GEMINI_GENERATION_CONFIG = genai.types.GenerationConfig(
+    temperature=0.2,
+    max_output_tokens=450,
+    top_p=0.95,
+)
 
+# --- RAG & Data Source Setup ---
 try:
-  gr.Info("Setting up retriever, please wait...")
-  rag_initial_output = RAG.search("what is Mistral?", k = 1)
-  gr.Info("Retriever working successfully!")
-    
-except:
-  gr.Warning("Retriever not working!")
-
-## Header
-mark_text = '# 🔍 Search Results\n'
-header_text = "# ArXiv CS RAG \n"
-
+    gr.Info("Setting up the RAG retriever...")
+    # If the local index directory doesn't exist, download it from Hugging Face.
+    if not os.path.exists(LOCAL_DATA_DIR):
+        if not RAG_SOURCE or not hf_token:
+            raise ValueError("RAG index not found locally, and RAG_SOURCE or HF_TOKEN environment variables are not set. Cannot download index.")
+        
+        snapshot_download(
+            repo_id=RAG_SOURCE,
+            repo_type="dataset", # Your index is stored as a dataset repo
+            token=hf_token,
+            local_dir=LOCAL_DATA_DIR
+        )
+        gr.Info("Index downloaded successfully.")
+    else:
+        gr.Info(f"Found existing local index at {LOCAL_DATA_DIR}.")
+
+    # Load the RAG model from the (now existing) local index path.
+    gr.Info(f'''Loading index from {os.path.join(LOCAL_DATA_DIR, "arxiv_colbert")}...''')
+    RAG = RAGPretrainedModel.from_index(os.path.join(LOCAL_DATA_DIR, "arxiv_colbert"))
+    RAG.search("Test query", k=1) # Warm-up query
+    gr.Info("Retriever loaded successfully!")
+
+except Exception as e:
+    gr.Warning(f"Could not initialize the RAG retriever. The app may not function correctly. Error: {e}")
+    RAG = None
+
+
+# --- UI Text and Metadata ---
+MARKDOWN_SEARCH_RESULTS_HEADER = '# 🔍 Search Results\n'
+APP_HEADER_TEXT = "# ArXiv CS RAG\n"
+INDEX_INFO = "Semantic Search"
 try:
-  with open("README.md", "r") as f:
-      mdfile = f.read()
-  date_pattern = r'Index Last Updated : \d{4}-\d{2}-\d{2}'
-  match = re.search(date_pattern, mdfile)
-  date = match.group().split(': ')[1]
-  formatted_date = datetime.strptime(date, '%Y-%m-%d').strftime('%d %b %Y')
-  header_text += f'Index Last Updated: {formatted_date}\n'
-  index_info = f"Semantic Search - up to {formatted_date}"  
-except:
-  index_info = "Semantic Search"
-
-database_choices = [index_info,'Arxiv Search - Latest - (EXPERIMENTAL)']
-
-## Arxiv API
-arx_client = arxiv.Client()
-is_arxiv_available = True
-check_arxiv_result = get_arxiv_live_search("What is Mistral?", arx_client, retrieve_results)
-if len(check_arxiv_result) == 0:
-  is_arxiv_available = False
-  print("Arxiv search not working, switching to default search ...")
-  database_choices = [index_info]
-
-
-
-## Show examples (disabled)
-if show_examples:
-    with open("sample_outputs.json", "r") as f:
-      sample_outputs = json.load(f)
-    output_placeholder = sample_outputs['output_placeholder']
-    md_text_initial = sample_outputs['search_placeholder']
-    
-else:
-    output_placeholder = None 
-    md_text_initial = ''
-
-
-def rag_cleaner(inp):
-    rank = inp['rank']
-    title = inp['document_metadata']['title']
-    content = inp['content']
-    date = inp['document_metadata']['_time']
-    return f"{rank}. <b> {title} </b> \n Date : {date} \n Abstract: {content}"
+    with open("README.md", "r") as f:
+        mdfile = f.read()
+    date_match = re.search(r'Index Last Updated : (\d{4}-\d{2}-\d{2})', mdfile)
+    if date_match:
+        date = date_match.group(1)
+        formatted_date = datetime.strptime(date, '%Y-%m-%d').strftime('%d %b %Y')
+        APP_HEADER_TEXT += f'Index Last Updated: {formatted_date}\n'
+        INDEX_INFO = f"Semantic Search - up to {formatted_date}"
+except Exception:
+    print("README.md not found or is invalid. Using default data source info.")
+
+DATABASE_CHOICES = [INDEX_INFO, 'Arxiv Search - Latest - (EXPERIMENTAL)']
+ARX_CLIENT = arxiv.Client()
+
+# --- Helper Functions ---
+def get_prompt_text(question, context):
+    """Formats the prompt for the Gemma 3 model."""
+    system_instruction = (
+        "Based on the provided scientific paper abstracts, provide a comprehensive answer of 6-7 lines. "
+        "Synthesize information from multiple sources if possible. Your answer must be grounded in the "
+        "details found in the abstracts. Cite the titles of the papers you use as sources in your answer."
+    )
+    message = f"Abstracts:\n{context}\n\nQuestion: {question}"
+    return f"{system_instruction}\n\n{message}"
+
+def update_with_rag_md(message, llm_results_use, database_choice):
+    """Fetches documents, updates the UI, and creates the final prompt for the LLM."""
+    prompt_context = ""
+    rag_out = []
+    source_used = database_choice
+
+    try:
+        if database_choice == INDEX_INFO and RAG:
+            rag_out = RAG.search(message, k=RETRIEVE_RESULTS)
+        else:
+            rag_out = list(ARX_CLIENT.results(arxiv.Search(query=message, max_results=RETRIEVE_RESULTS, sort_by=arxiv.SortCriterion.Relevance)))
+            if not rag_out:
+                gr.Warning("Live Arxiv search returned no results. Falling back to semantic search.")
+                if RAG:
+                    rag_out = RAG.search(message, k=RETRIEVE_RESULTS)
+                    source_used = INDEX_INFO
+    except Exception as e:
+        gr.Warning(f"An error occurred during search: {e}. Falling back to semantic search.")
+        if RAG:
+            rag_out = RAG.search(message, k=RETRIEVE_RESULTS)
+            source_used = INDEX_INFO
+
+    md_text_updated = MARKDOWN_SEARCH_RESULTS_HEADER
+    for i, rag_answer in enumerate(rag_out):
+        md_text_paper, prompt_text = get_md_text_abstract(rag_answer, source=source_used, return_prompt_formatting=True)
+        if i < llm_results_use:
+            prompt_context += f"{i+1}. {prompt_text}\n"
+        md_text_updated += md_text_paper
+
+    final_prompt = get_prompt_text(message, prompt_context)
+    return md_text_updated, final_prompt
+
+def ask_gemma_llm(prompt, llm_model_picked, stream_outputs):
+    """Sends a prompt to the Google Gemini API and streams the response."""
+    if not prompt or not prompt.strip():
+        yield "Error: The generated prompt is empty. Please try a different query."
+        return
+
+    if llm_model_picked == 'None':
+        yield "LLM Model is disabled."
+        return
+
+    if not gemini_api_key:
+        yield "Error: GEMINI_API_KEY is not configured. Cannot contact the LLM."
+        return
+
+    try:
+        safety_settings = [
+            {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_ONLY_HIGH"},
+            {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_ONLY_HIGH"},
+            {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_ONLY_HIGH"},
+            {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_ONLY_HIGH"},
+        ]
+
+        gemini_model_name = llm_model_picked.split('/')[-1]
+        model = genai.GenerativeModel(gemini_model_name)
+        response = model.generate_content(
+            prompt,
+            generation_config=GEMINI_GENERATION_CONFIG,
+            stream=stream_outputs,
+            safety_settings=safety_settings
+        )
+
+        if stream_outputs:
+            output = ""
+            for chunk in response:
+                try:
+                    text = chunk.parts[0].text
+                    output += text
+                    yield output
+                except (IndexError, AttributeError):
+                    # Ignore empty chunks, which can occur at the end of a stream.
+                    pass
+            if not output:
+                 yield "Model returned an empty or blocked stream. This may be due to the safety settings or the nature of the prompt."
+        else:
+            # Handle non-streaming responses.
+            try:
+                yield response.parts[0].text
+            except (IndexError, AttributeError):
+                reason = "UNKNOWN"
+                if response.prompt_feedback.block_reason:
+                    reason = response.prompt_feedback.block_reason.name
+                elif response.candidates and not response.candidates[0].content.parts:
+                     reason = response.candidates[0].finish_reason.name
+                yield f"Model returned an empty or blocked response."
+
+    except Exception as e:
+        error_message = f"An error occurred with the Gemini API: {e}"
+        print(error_message) # Server side log
+        gr.Warning("An error occurred with the Gemini API. Check the server logs for details.")
+        yield error_message
+
+# --- Gradio User Interface ---
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown(APP_HEADER_TEXT)
 
-def get_prompt_text(question, context, formatted = True, llm_model_picked = 'mistralai/Mistral-7B-Instruct-v0.3'):
-    if formatted:
-      sys_instruction = f"Context:\n {context} \n Given the following scientific paper abstracts, take a deep breath and lets think step by step to answer the question. Cite the titles of your sources when answering, do not cite links or dates."
-      message = f"Question: {question}"
-        
-      if 'mistralai' in llm_model_picked:
-          return f"<s>" + f"[INST] {sys_instruction}" +  f" {message}[/INST]"
-          
-      elif 'gemma' in llm_model_picked:
-          return f"<bos><start_of_turn>user\n{sys_instruction}" +  f" {message}<end_of_turn>\n"
-          
-    return f"Context:\n {context} \n Given the following info, take a deep breath and lets think step by step to answer the question: {question}. Cite the titles of your sources when answering.\n\n"
-
-def get_references(question, retriever, k = retrieve_results):
-    rag_out = retriever.search(query=question, k=k)
-    return rag_out
-
-def get_rag(message):
-    return get_references(message, RAG)
-
-with gr.Blocks(theme = gr.themes.Soft()) as demo:
-    header = gr.Markdown(header_text)
-    
     with gr.Group():
-      msg = gr.Textbox(label = 'Search', placeholder = 'What is Mistral?')
-        
-      with gr.Accordion("Advanced Settings", open=False):
-        with gr.Row(equal_height = True):
-          llm_model = gr.Dropdown(choices = llm_models_to_choose, value = 'mistralai/Mistral-7B-Instruct-v0.3', label = 'LLM Model')
-          llm_results = gr.Slider(minimum=4, maximum=10, value=5, step=1, interactive=True, label="Top n results as context")
-          database_src = gr.Dropdown(choices = database_choices, value = index_info, label = 'Search Source')
-          stream_results = gr.Checkbox(value = True, label = "Stream output", visible = False)
-
-    output_text = gr.Textbox(show_label = True, container = True, label = 'LLM Answer', visible = True, placeholder = output_placeholder)
-    input = gr.Textbox(show_label = False, visible = False)
-    gr_md = gr.Markdown(mark_text + md_text_initial)
-
-    def update_with_rag_md(message, llm_results_use = 5, database_choice = index_info, llm_model_picked = 'mistralai/Mistral-7B-Instruct-v0.3'):
-        prompt_text_from_data = ""
-        database_to_use = database_choice
-        if database_choice == index_info:
-          rag_out = get_rag(message)
-        else:
-          arxiv_search_success = True
-          try:
-            rag_out = get_arxiv_live_search(message, arx_client, retrieve_results)
-            if len(rag_out) == 0:
-              arxiv_search_success = False 
-          except:
-            arxiv_search_success = False
- 
-
-          if not arxiv_search_success:
-            gr.Warning("Arxiv Search not working, switching to semantic search ...")
-            rag_out = get_rag(message)
-            database_to_use = index_info 
-
-        md_text_updated = mark_text
-        for i in range(retrieve_results):
-          rag_answer = rag_out[i]
-          if i < llm_results_use:
-            md_text_paper, prompt_text = get_md_text_abstract(rag_answer, source = database_to_use, return_prompt_formatting = True)
-            prompt_text_from_data += f"{i+1}. {prompt_text}"
-          else:
-            md_text_paper = get_md_text_abstract(rag_answer, source = database_to_use)
-          md_text_updated += md_text_paper
-        prompt = get_prompt_text(message, prompt_text_from_data, llm_model_picked = llm_model_picked)
-        return md_text_updated, prompt
-
-    def ask_llm(prompt, llm_model_picked = 'mistralai/Mistral-7B-Instruct-v0.3', stream_outputs = False):
-       model_disabled_text = "LLM Model is disabled"
-       output = ""
-        
-       if llm_model_picked == 'None':
-          if stream_outputs:
-              for out in model_disabled_text:
-                output += out
-                yield output
-              return output 
-          else:
-              return model_disabled_text
-              
-       client = InferenceClient(llm_model_picked, token = token)
-       try:
-           stream = client.text_generation(prompt, **generate_kwargs,  stream=stream_outputs, details=False, return_full_text=False)
-           
-       except:
-           gr.Warning("LLM Inference rate limit reached, try again later!")
-           return ""
-       
-       if stream_outputs:
-           for response in stream:
-              output += response
-              yield output
-           return output
-       else:
-           return stream
-
-
-    msg.submit(update_with_rag_md, [msg, llm_results,  database_src, llm_model], [gr_md, input]).success(ask_llm, [input, llm_model, stream_results], output_text)
-
-demo.queue().launch()
+        msg = gr.Textbox(label='Search', placeholder='e.g., What is Mixtral?')
+        with gr.Accordion("Advanced Settings", open=False):
+            llm_model = gr.Dropdown(choices=LLM_MODELS_TO_CHOOSE, value=DEFAULT_LLM_MODEL, label='LLM Model')
+            llm_results = gr.Slider(5, 20, value=10, step=1, label="Top n results as context")
+            database_src = gr.Dropdown(choices=DATABASE_CHOICES, value=INDEX_INFO, label='Search Source')
+            stream_results = gr.Checkbox(value=True, label="Stream output")
+
+    output_text = gr.Textbox(label='LLM Answer', placeholder="The model's answer will appear here...", interactive=False, lines=8)
+    input_prompt = gr.Textbox(visible=False)
+    gr_md = gr.Markdown(MARKDOWN_SEARCH_RESULTS_HEADER)
+
+    msg.submit(
+        fn=update_with_rag_md,
+        inputs=[msg, llm_results, database_src],
+        outputs=[gr_md, input_prompt]
+    ).then(
+        fn=ask_gemma_llm,
+        inputs=[input_prompt, llm_model, stream_results],
+        outputs=[output_text]
+    )
+
+if __name__ == "__main__":
+    # Launch the app
+    demo.queue().launch()

requirements.txt

@@ -1,4 +1,5 @@
-ragatouille
+ragatouille<0.0.10
+google-generativeai
 huggingface_hub
 arxiv
 transformers[torch]<=4.49.0