[Feature Improvement] Add support for Ollama, Gemini, and Claude, with Web UI

agents.py

@@ -1,3 +1,5 @@
+import json
+
 from utils import *
 from tools import *
 from inference import *

ai_lab_repo.py

@@ -544,6 +544,18 @@ def parse_arguments():
         help='Provide the DeepSeek API key.'
     )
 
+    parser.add_argument(
+        '--google-api-key',
+        type=str,
+        help='Provide the Google API key.'
+    )
+
+    parser.add_argument(
+        '--anthropic-api-key',
+        type=str,
+        help='Provide the Anthropic API key.'
+    )
+
     parser.add_argument(
         '--load-existing',
         type=str,
@@ -638,13 +650,20 @@ def parse_arguments():
 
     api_key = os.getenv('OPENAI_API_KEY') or args.api_key
     deepseek_api_key = os.getenv('DEEPSEEK_API_KEY') or args.deepseek_api_key
+    google_api_key = os.getenv('GOOGLE_API_KEY') or args.google_api_key
+    anthropic_api_key = os.getenv('ANTHROPIC_API_KEY') or args.anthropic_api_key
     if args.api_key is not None and os.getenv('OPENAI_API_KEY') is None:
         os.environ["OPENAI_API_KEY"] = args.api_key
     if args.deepseek_api_key is not None and os.getenv('DEEPSEEK_API_KEY') is None:
         os.environ["DEEPSEEK_API_KEY"] = args.deepseek_api_key
-
-    if not api_key and not deepseek_api_key:
-        raise ValueError("API key must be provided via --api-key / -deepseek-api-key or the OPENAI_API_KEY / DEEPSEEK_API_KEY environment variable.")
+    if args.google_api_key is not None and os.getenv('GOOGLE_API_KEY') is None:
+        os.environ["GOOGLE_API_KEY"] = args.google_api_key
+    if args.anthropic_api_key is not None and os.getenv('ANTHROPIC_API_KEY') is None:
+        os.environ["ANTHROPIC_API_KEY"] = args.anthropic_api_key
+
+    if not api_key and not deepseek_api_key and not google_api_key and not anthropic_api_key:
+        raise ValueError("API key must be provided via --api-key / -deepseek-api-key / --google-api-key / --anthropic-api-key argument "
+                         "or the OPENAI_API_KEY / DEEPSEEK_API_KEY / GOOGLE_API_KEY / ANTHROPIC_API_KEY environment variable.")
 
     ##########################################################
     # Research question that the agents are going to explore #

config.py

@@ -0,0 +1,3 @@
+GOOGLE_GENERATIVE_API_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/"
+DEEPSEEK_API_BASE_URL = "https://api.deepseek.com/v1"
+OLLAMA_API_BASE_URL = "http://localhost:11434/v1/"

config_gradio.py

@@ -0,0 +1,258 @@
+import os
+import subprocess
+import sys
+from typing import Any
+import webbrowser
+
+import gradio as gr
+
+
+def get_existing_saves() -> list:
+    """Retrieve list of existing save files from state_saves directory."""
+    saves_dir = 'state_saves'
+    try:
+        os.makedirs(saves_dir, exist_ok=True)
+        # List all .pkl files in the directory
+        saves = [f for f in os.listdir(saves_dir) if f.endswith('.pkl')]
+        return saves if saves else ["No saved states found"]
+    except Exception as e:
+        print(f"Error retrieving saves: {e}")
+        return ["No saved states found"]
+
+
+def refresh_saves_dropdown():
+    """
+    IMPORTANT PART:
+    Return a *new* gr.Dropdown component populated with fresh choices.
+    This replaces the existing dropdown instead of attempting to update it.
+    """
+    new_saves = get_existing_saves()
+    return gr.Dropdown(
+        choices=new_saves,
+        label="Select Saved Research State",
+        interactive=True
+    )
+
+
+def run_research_process(
+    research_topic: str,
+    api_key: str,
+    llm_backend: str,
+    custom_llm_backend: str,
+    language: str,
+    copilot_mode: bool,
+    compile_latex: bool,
+    num_papers_lit_review: Any,  # Gradio numbers may come as float
+    mlesolver_max_steps: Any,
+    papersolver_max_steps: Any,
+    deepseek_api_key: str = "",
+    google_api_key: str = "",
+    anthropic_api_key: str = "",
+    load_existing: bool = False,
+    load_existing_path: str = ""
+) -> str:
+    try:
+        # Determine which LLM backend to use:
+        # For Ollama, if a custom backend is provided, use it;
+        # otherwise, use the dropdown value.
+        if api_key.strip().lower() == "ollama":
+            chosen_backend = custom_llm_backend.strip() if custom_llm_backend.strip() else llm_backend
+        else:
+            chosen_backend = llm_backend
+
+        # Prepare the command arguments
+        cmd = [
+            sys.executable, 'ai_lab_repo.py',
+            '--research-topic', research_topic,
+            '--api-key', api_key,
+            '--llm-backend', chosen_backend,
+            '--language', language,
+            '--copilot-mode', str(copilot_mode).lower(),
+            '--compile-latex', str(compile_latex).lower(),
+            '--num-papers-lit-review', str(num_papers_lit_review),
+            '--mlesolver-max-steps', str(mlesolver_max_steps),
+            '--papersolver-max-steps', str(papersolver_max_steps)
+        ]
+
+        # Add DeepSeek API key if provided
+        if deepseek_api_key:
+            cmd.extend(['--deepseek-api-key', deepseek_api_key])
+
+        # Add Google API key if provided
+        if google_api_key:
+            cmd.extend(['--google-api-key', google_api_key])
+
+        # Add Anthropic API key if provided
+        if anthropic_api_key:
+            cmd.extend(['--anthropic-api-key', anthropic_api_key])
+
+        # Add load existing flags if selected
+        if load_existing and load_existing_path and load_existing_path != "No saved states found":
+            cmd.extend([
+                '--load-existing', 'True',
+                '--load-existing-path', os.path.join('state_saves', load_existing_path)
+            ])
+
+        # Open a new terminal window with the research process
+        if sys.platform == 'win32':
+            subprocess.Popen(['start', 'cmd', '/k'] + cmd, shell=True)
+        elif sys.platform == 'darwin':
+            subprocess.Popen(['open', '-a', 'Terminal'] + cmd)
+        else:
+            subprocess.Popen(['x-terminal-emulator', '-e'] + cmd)
+
+        return "Research process started in a new terminal window. Please check the terminal for progress."
+    except Exception as e:
+        error_message = f"Error starting research process: {e}"
+        print(error_message)
+        return error_message
+
+
+def create_gradio_config() -> gr.Blocks:
+    # Standard backend options for the dropdown
+    llm_backend_options = [
+        "o1",
+        "o1-preview",
+        "o1-mini",
+        "gpt-4o",
+        "gpt-4o-mini",
+        "deepseek-chat",
+        "claude-3-5-sonnet",
+        "claude-3-5-haiku",
+        "gemini-2.0-flash",
+        "gemini-2.0-flash-lite",
+    ]
+    languages = [
+        "English", "Chinese-Simplified", "Chinese-Traditional",
+        "Japanese", "Korean", "Filipino", "French",
+        "Slovak", "Portuguese", "Spanish", "Turkish", "Hindi", "Bengali",
+        "Vietnamese", "Russian", "Arabic", "Farsi", "Italian"
+    ]
+
+    with gr.Blocks() as demo:
+        gr.Markdown("# Agent Laboratory Configuration")
+
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("## Research Configuration")
+                research_topic = gr.Textbox(
+                    label="Research Topic",
+                    placeholder="Enter your research idea...",
+                    lines=3
+                )
+                api_key = gr.Textbox(
+                    label="OpenAI API Key",
+                    type="password",
+                    placeholder="Enter your OpenAI API key (for Ollama, set API key to 'ollama', must be set)"
+                )
+                deepseek_api_key = gr.Textbox(
+                    label="DeepSeek API Key (Optional)",
+                    type="password",
+                    placeholder="Enter your DeepSeek API key if using DeepSeek model"
+                )
+                google_api_key = gr.Textbox(
+                    label="Google API Key (Optional)",
+                    type="password",
+                    placeholder="Enter your Google API key if using Google models"
+                )
+                anthropic_api_key = gr.Textbox(
+                    label="Anthropic API Key (Optional)",
+                    type="password",
+                    placeholder="Enter your Anthropic API key if using Anthropic models"
+                )
+
+                with gr.Row():
+                    # Dropdown for standard LLM backend options
+                    llm_backend = gr.Dropdown(
+                        choices=llm_backend_options,
+                        label="LLM Backend",
+                        value="o1-mini"
+                    )
+                    language = gr.Dropdown(
+                        choices=languages,
+                        label="Language",
+                        value="English"
+                    )
+
+                with gr.Row():
+                    copilot_mode = gr.Checkbox(label="Enable Human-in-Loop Mode")
+                    compile_latex = gr.Checkbox(label="Compile LaTeX", value=True)
+
+                with gr.Row():
+                    num_papers_lit_review = gr.Number(
+                        label="Papers in Literature Review",
+                        value=5, precision=0, minimum=1, maximum=20
+                    )
+                    mlesolver_max_steps = gr.Number(
+                        label="MLE Solver Max Steps",
+                        value=3, precision=0, minimum=1, maximum=10
+                    )
+                    papersolver_max_steps = gr.Number(
+                        label="Paper Solver Max Steps",
+                        value=5, precision=0, minimum=1, maximum=10
+                    )
+
+                # Saved States Section
+                with gr.Accordion("Resume Previous Research", open=False):
+                    load_existing = gr.Checkbox(label="Load Existing Research State")
+                    existing_saves = gr.Dropdown(
+                        choices=get_existing_saves(),
+                        label="Select Saved Research State",
+                        interactive=True
+                    )
+                    refresh_saves_btn = gr.Button("Refresh Saved States")
+
+                # Custom LLM Backend textbox for Ollama.
+                # This is optional and only used when API key is set to "ollama".
+                custom_llm_backend = gr.Textbox(
+                    label="Custom LLM Backend (For Ollama)",
+                    placeholder="Enter your custom model string (optional)",
+                    value=""
+                )
+
+                submit_btn = gr.Button("Start Research in Terminal", variant="primary")
+
+            with gr.Column():
+                gr.Markdown("## Instructions")
+                gr.Markdown(
+                    """
+                    - Fill in the research configuration.
+                    - Optionally load a previous research state.
+                    - **For standard models:** Select the desired backend from the dropdown.
+                    - **For Ollama:** Set the API key to `ollama` and, if needed, enter your custom model string in the **Custom LLM Backend** field.
+                    - If the custom field is left empty when using Ollama, the dropdown value will be used.
+                    - Click **Start Research in Terminal**.
+                    - A new terminal window will open with the research process.
+                    """
+                )
+
+        # Instead of returning just a list, return a new Dropdown from refresh_saves_dropdown()
+        refresh_saves_btn.click(
+            fn=refresh_saves_dropdown,
+            inputs=None,
+            outputs=existing_saves
+        )
+
+        # Connect submit button to a research process with the new custom_llm_backend input.
+        submit_btn.click(
+            fn=run_research_process,
+            inputs=[
+                research_topic, api_key, llm_backend, custom_llm_backend, language,
+                copilot_mode, compile_latex, num_papers_lit_review,
+                mlesolver_max_steps, papersolver_max_steps,
+                deepseek_api_key, google_api_key, anthropic_api_key,
+                load_existing, existing_saves,
+            ],
+            outputs=gr.Textbox(label="Status")
+        )
+
+    return demo
+
+
+def main():
+    demo = create_gradio_config()
+    demo.launch()
+
+
+if __name__ == "__main__":
+    main()