Add multi-provider LLM support and two-tier model configuration

.gitignore

@@ -232,3 +232,14 @@ data/ddinter_raw/
 # Sphinx build
 /docs/build/
 /docs/source/api/
+biomni_env/setup_path.sh
+CLAUDE.md
+
+# Analysis data/outputs generated during test runs
+geo_data/
+outputs/
+outputs_dn_meta/
+dn_analysis/
+dn_glom_meta_probelevel.tsv
+test_dn_meta_analysis.py
+workspace/

biomni/agent/env_collection.py

@@ -194,15 +194,20 @@ def _consolidate_tasks(self, chunk_results: list[str]) -> dict[str, list[dict[st
         prompt = self.consolidation_prompt.format(task_lists=all_tasks)
         response = self.llm.invoke(prompt)
 
+        # Normalize content (handles string, list of blocks, etc.)
+        from biomni.utils import normalize_llm_content
+
+        response_text = normalize_llm_content(response.content)
+
         # Extract the JSON from the response
         try:
             # Try to parse the entire response as JSON
-            result = json.loads(response.content)
+            result = json.loads(response_text)
         except json.JSONDecodeError:
             # If that fails, try to extract JSON from the text
             try:
                 # Look for JSON-like content between triple backticks
-                json_match = re.search(r"```(?:json)?\s*([\s\S]*?)\s*```", response.content)
+                json_match = re.search(r"```(?:json)?\s*([\s\S]*?)\s*```", response_text)
                 if json_match:
                     result = json.loads(json_match.group(1))
                 else:
@@ -211,7 +216,7 @@ def _consolidate_tasks(self, chunk_results: list[str]) -> dict[str, list[dict[st
                         "tasks": [
                             {
                                 "error": "Could not parse JSON",
-                                "raw_response": response.content,
+                                "raw_response": response_text,
                             }
                         ],
                         "databases": [],
@@ -222,7 +227,7 @@ def _consolidate_tasks(self, chunk_results: list[str]) -> dict[str, list[dict[st
                     "tasks": [
                         {
                             "error": "Could not parse JSON",
-                            "raw_response": response.content,
+                            "raw_response": response_text,
                         }
                     ],
                     "databases": [],

biomni/config.py

@@ -29,11 +29,13 @@ class BiomniConfig:
     """
 
     # Data and execution settings
-    path: str = "./data"
+    path: str = "./data"  # Data lake path (input data)
+    workspace: str = "./workspace"  # Output directory for generated files
     timeout_seconds: int = 600
 
     # LLM settings (API keys still from environment)
-    llm: str = "claude-sonnet-4-5"
+    llm: str = "claude-sonnet-4-5"  # Primary model for agent reasoning
+    llm_lite: str | None = None  # Lightweight model for simple tasks; auto-inferred from llm provider if not set
     temperature: float = 0.7
 
     # Tool settings
@@ -58,10 +60,14 @@ def __post_init__(self):
         # Support both old and new names for backwards compatibility
         if os.getenv("BIOMNI_PATH") or os.getenv("BIOMNI_DATA_PATH"):
             self.path = os.getenv("BIOMNI_PATH") or os.getenv("BIOMNI_DATA_PATH")
+        if os.getenv("BIOMNI_WORKSPACE"):
+            self.workspace = os.getenv("BIOMNI_WORKSPACE")
         if os.getenv("BIOMNI_TIMEOUT_SECONDS"):
             self.timeout_seconds = int(os.getenv("BIOMNI_TIMEOUT_SECONDS"))
         if os.getenv("BIOMNI_LLM") or os.getenv("BIOMNI_LLM_MODEL"):
             self.llm = os.getenv("BIOMNI_LLM") or os.getenv("BIOMNI_LLM_MODEL")
+        if os.getenv("BIOMNI_LLM_LITE"):
+            self.llm_lite = os.getenv("BIOMNI_LLM_LITE")
         if os.getenv("BIOMNI_USE_TOOL_RETRIEVER"):
             self.use_tool_retriever = os.getenv("BIOMNI_USE_TOOL_RETRIEVER").lower() == "true"
         if os.getenv("BIOMNI_COMMERCIAL_MODE"):
@@ -80,12 +86,34 @@ def __post_init__(self):
         if env_token:
             self.protocols_io_access_token = env_token
 
+        # Auto-infer llm_lite from the main llm's provider if not explicitly set
+        if self.llm_lite is None:
+            self.llm_lite = self._infer_lite_model(self.llm)
+
+    @staticmethod
+    def _infer_lite_model(model: str) -> str:
+        """Return a lightweight model that matches the provider of the given model."""
+        _LITE_MODELS = {
+            "claude-": "claude-haiku-4-5",
+            "gpt-": "gpt-5-mini",
+            "gemini-": "gemini-1.5-flash",
+            "groq": "llama-3.1-8b-instant",
+        }
+        model_lower = model.lower()
+        for prefix, lite in _LITE_MODELS.items():
+            if model_lower.startswith(prefix) or prefix in model_lower:
+                return lite
+        # Fallback: same provider detection as llm.py
+        return "claude-haiku-4-5"
+
     def to_dict(self) -> dict:
         """Convert config to dictionary for easy access."""
         return {
             "path": self.path,
+            "workspace": self.workspace,
             "timeout_seconds": self.timeout_seconds,
             "llm": self.llm,
+            "llm_lite": self.llm_lite,
             "temperature": self.temperature,
             "use_tool_retriever": self.use_tool_retriever,
             "commercial_mode": self.commercial_mode,

biomni/env_desc.py

@@ -93,6 +93,7 @@
     "biotite": "[Python Package] A comprehensive library for computational molecular biology, providing tools for sequence analysis, structure analysis, and more.",
     "lazyslide": "[Python Package] A Python framework that brings interoperable, reproducible whole slide image analysis, enabling seamless histopathology workflows from preprocessing to deep learning.",
     # Genomics & Variant Analysis (Python)
+    "GEOparse": "[Python Package] A library for downloading and parsing data from NCBI GEO (Gene Expression Omnibus). Use GEOparse.get_GEO(geo='GSE12345', destdir='./output') to download series or sample data.",
     "gget": "[Python Package] A toolkit for accessing genomic databases and retrieving sequences, annotations, and other genomic data.",
     "lifelines": "[Python Package] A complete survival analysis library for fitting models, plotting, and statistical tests.",
     # "scvi-tools": "[Python Package] A package for probabilistic modeling of single-cell omics data, including deep generative models.",

biomni/env_desc_cm.py

@@ -93,6 +93,7 @@
     "biotite": "[Python Package] A comprehensive library for computational molecular biology, providing tools for sequence analysis, structure analysis, and more.",
     "lazyslide": "[Python Package] A Python framework that brings interoperable, reproducible whole slide image analysis, enabling seamless histopathology workflows from preprocessing to deep learning.",
     # Genomics & Variant Analysis (Python)
+    "GEOparse": "[Python Package] A library for downloading and parsing data from NCBI GEO (Gene Expression Omnibus). Use GEOparse.get_GEO(geo='GSE12345', destdir='./output') to download series or sample data.",
     "gget": "[Python Package] A toolkit for accessing genomic databases and retrieving sequences, annotations, and other genomic data.",
     "lifelines": "[Python Package] A complete survival analysis library for fitting models, plotting, and statistical tests.",
     # "scvi-tools": "[Python Package] A package for probabilistic modeling of single-cell omics data, including deep generative models.",

biomni/llm.py

@@ -9,6 +9,74 @@
 SourceType = Literal["OpenAI", "AzureOpenAI", "Anthropic", "Ollama", "Gemini", "Bedrock", "Groq", "Custom"]
 ALLOWED_SOURCES: set[str] = set(SourceType.__args__)
 
+# Default models for each provider (used when no model is specified)
+DEFAULT_MODELS = {
+    "Anthropic": "claude-sonnet-4-5",
+    "OpenAI": "gpt-4o",
+    "Gemini": "gemini-1.5-pro",
+    "Groq": "llama-3.1-70b-versatile",
+    "Bedrock": "anthropic.claude-3-5-sonnet-20241022-v2:0",
+}
+
+# Lightweight models for simple tasks (parsing, classification, etc.)
+DEFAULT_MODELS_LITE = {
+    "Anthropic": "claude-haiku-4-5",
+    "OpenAI": "gpt-5-mini",
+    "Gemini": "gemini-1.5-flash",
+    "Groq": "llama-3.1-8b-instant",
+    "Bedrock": "anthropic.claude-3-5-haiku-20241022-v1:0",
+}
+
+
+def _get_default_model() -> str:
+    """Select the default model based on available API keys.
+
+    Checks for API keys in order of preference and returns the appropriate
+    default model for the first available provider.
+
+    Returns:
+        str: The default model name for the available provider.
+    """
+    # Check providers in order of preference
+    if os.getenv("ANTHROPIC_API_KEY"):
+        return DEFAULT_MODELS["Anthropic"]
+    if os.getenv("OPENAI_API_KEY"):
+        return DEFAULT_MODELS["OpenAI"]
+    if os.getenv("GOOGLE_API_KEY") or os.getenv("GEMINI_API_KEY"):
+        return DEFAULT_MODELS["Gemini"]
+    if os.getenv("GROQ_API_KEY"):
+        return DEFAULT_MODELS["Groq"]
+    if os.getenv("AWS_ACCESS_KEY_ID") or os.getenv("AWS_PROFILE"):
+        return DEFAULT_MODELS["Bedrock"]
+
+    # Fallback to Anthropic model (will fail if no key, but provides clear error)
+    return DEFAULT_MODELS["Anthropic"]
+
+
+def _get_default_model_lite() -> str:
+    """Select the default lightweight model based on available API keys.
+
+    Similar to _get_default_model() but returns cheaper/faster models
+    suitable for simple tasks like parsing, classification, etc.
+
+    Returns:
+        str: The default lite model name for the available provider.
+    """
+    # Check providers in order of preference
+    if os.getenv("ANTHROPIC_API_KEY"):
+        return DEFAULT_MODELS_LITE["Anthropic"]
+    if os.getenv("OPENAI_API_KEY"):
+        return DEFAULT_MODELS_LITE["OpenAI"]
+    if os.getenv("GOOGLE_API_KEY") or os.getenv("GEMINI_API_KEY"):
+        return DEFAULT_MODELS_LITE["Gemini"]
+    if os.getenv("GROQ_API_KEY"):
+        return DEFAULT_MODELS_LITE["Groq"]
+    if os.getenv("AWS_ACCESS_KEY_ID") or os.getenv("AWS_PROFILE"):
+        return DEFAULT_MODELS_LITE["Bedrock"]
+
+    # Fallback to Anthropic model (will fail if no key, but provides clear error)
+    return DEFAULT_MODELS_LITE["Anthropic"]
+
 
 def get_llm(
     model: str | None = None,
@@ -35,7 +103,7 @@ def get_llm(
     # Use config values for any unspecified parameters
     if config is not None:
         if model is None:
-            model = config.llm_model
+            model = config.llm  # Use config's LLM setting
         if temperature is None:
             temperature = config.temperature
         if source is None:
@@ -45,9 +113,9 @@ def get_llm(
         if api_key is None:
             api_key = config.api_key or "EMPTY"
 
-    # Use defaults if still not specified
+    # Use defaults if still not specified - select based on available API keys
     if model is None:
-        model = "claude-3-5-sonnet-20241022"
+        model = _get_default_model()
     if temperature is None:
         temperature = 0.7
     if api_key is None:
@@ -131,12 +199,14 @@ def _get_request_payload(self, input_, *, stop=None, **kwargs):  # type: ignore[
                 stop_sequences=stop_sequences,
                 use_responses_api=True,
                 output_version="v0",
+                request_timeout=300,  # 5 minute timeout for API requests
             )
         else:
             return ChatOpenAI(
                 model=model,
                 temperature=temperature,
                 stop_sequences=stop_sequences,
+                request_timeout=300,  # 5 minute timeout for API requests
             )
 
     elif source == "AzureOpenAI":
@@ -186,6 +256,7 @@ def _get_request_payload(self, input_, *, stop=None, **kwargs):  # type: ignore[
             temperature=temperature,
             max_tokens=8192,
             stop_sequences=stop_sequences,
+            timeout=300,  # 5 minute timeout for API requests
         )
 
     elif source == "Gemini":

biomni/model/retriever.py

@@ -96,7 +96,10 @@ def prompt_based_retrieval(self, query: str, resources: dict, llm=None) -> dict:
         if hasattr(llm, "invoke"):
             # For LangChain-style LLMs
             response = llm.invoke([HumanMessage(content=prompt)])
-            response_content = response.content
+            # Normalize content (handles string, list of blocks, etc.)
+            from biomni.utils import normalize_llm_content
+
+            response_content = normalize_llm_content(response.content)
         else:
             # For other LLM interfaces
             response_content = str(llm(prompt))