Add multi-LLM provider support and position agent system for cross-modal geophysics

PyHydroGeophysX/agents/__init__.py

@@ -1,9 +1,13 @@
 """
 Multi-Agent System for Automated Geophysical Workflows
 
-This module provides a GPT API-based multi-agent system for automating
-the workflow: "load ERT → invert → convert to water content → report"
-with optional seismic data integration.
+This module provides an automatic cross-modal geophysics agent system for 
+subsurface hydrology, supporting multiple LLM APIs (GPT, Gemini, Claude) to 
+automate workflows that process geophysical data (ERT, seismic, and more) 
+into hydrologic information.
+
+The workflow example: "load geophysical data → process → invert → convert to 
+hydrologic parameters → report" with optional cross-modal integration.
 
 Each agent is specialized for a specific task and communicates through
 a coordinator to execute the complete workflow.

PyHydroGeophysX/agents/agent_coordinator.py

@@ -2,7 +2,9 @@
 Agent Coordinator for Multi-Agent Workflow
 
 Coordinates the execution of multiple specialized agents to complete
-the full geophysical processing workflow.
+the full geophysical processing workflow. Supports cross-modal geophysical
+data processing (ERT, seismic, and more) with multiple LLM API providers
+(GPT, Gemini, Claude).
 """
 
 from typing import Dict, Any, List, Optional
@@ -15,21 +17,24 @@ class AgentCoordinator:
     """
     Coordinates multiple agents to execute a complete workflow.
 
-    The coordinator manages the workflow: 
-    "load ERT → invert → convert to water content → report"
-    with optional seismic data integration.
+    The coordinator manages cross-modal geophysical workflows such as: 
+    "load geophysical data → process → invert → convert to hydrologic parameters → report"
+    with support for multiple data types (ERT, seismic, etc.) and LLM providers (GPT, Gemini, Claude).
     """
 
-    def __init__(self, api_key: Optional[str] = None, output_dir: str = "results/agents"):
+    def __init__(self, api_key: Optional[str] = None, output_dir: str = "results/agents",
+                 llm_provider: str = "openai"):
         """
         Initialize the agent coordinator.
 
         Args:
-            api_key: OpenAI API key for agents
+            api_key: LLM API key for agents
             output_dir: Directory for saving results
+            llm_provider: LLM provider to use ('openai', 'gemini', or 'claude')
         """
-        self.api_key = api_key or os.getenv('OPENAI_API_KEY')
+        self.api_key = api_key or self._get_default_api_key(llm_provider)
         self.output_dir = output_dir
+        self.llm_provider = llm_provider.lower()
         self.agents = {}
         self.workflow_state = {
             'status': 'initialized',
@@ -42,6 +47,16 @@ def __init__(self, api_key: Optional[str] = None, output_dir: str = "results/age
         # Create output directory
         os.makedirs(output_dir, exist_ok=True)
 
+    def _get_default_api_key(self, provider: str) -> Optional[str]:
+        """Get default API key based on provider."""
+        provider_env_map = {
+            'openai': 'OPENAI_API_KEY',
+            'gemini': 'GEMINI_API_KEY',
+            'claude': 'ANTHROPIC_API_KEY'
+        }
+        env_var = provider_env_map.get(provider.lower())
+        return os.getenv(env_var) if env_var else None
+
     def register_agent(self, agent_name: str, agent_instance):
         """
         Register an agent with the coordinator.

PyHydroGeophysX/agents/base_agent.py

@@ -18,18 +18,35 @@ class BaseAgent(ABC):
     with other agents through the coordinator.
     """
 
-    def __init__(self, name: str, api_key: Optional[str] = None, model: str = "gpt-4"):
+    def __init__(self, name: str, api_key: Optional[str] = None, model: Optional[str] = None, 
+                 llm_provider: str = "openai"):
         """
         Initialize the base agent.
 
         Args:
             name: Name identifier for this agent
-            api_key: OpenAI API key (uses OPENAI_API_KEY env var if not provided)
-            model: OpenAI model to use (default: gpt-4, alternatives: gpt-3.5-turbo, gpt-4-turbo)
+            api_key: LLM API key (uses provider-specific env var if not provided)
+            model: LLM model to use (default: gpt-4 for OpenAI, gemini-pro for Gemini, 
+                   claude-3-opus-20240229 for Claude)
+            llm_provider: LLM provider to use ('openai', 'gemini', or 'claude')
         """
         self.name = name
-        self.api_key = api_key or os.getenv('OPENAI_API_KEY')
-        self.model = model or os.getenv('OPENAI_MODEL', 'gpt-4')
+        self.llm_provider = llm_provider.lower()
+
+        # Set API key and default model based on provider
+        if self.llm_provider == "openai":
+            self.api_key = api_key or os.getenv('OPENAI_API_KEY')
+            self.model = model or os.getenv('OPENAI_MODEL', 'gpt-4')
+        elif self.llm_provider == "gemini":
+            self.api_key = api_key or os.getenv('GEMINI_API_KEY')
+            self.model = model or os.getenv('GEMINI_MODEL', 'gemini-pro')
+        elif self.llm_provider == "claude":
+            self.api_key = api_key or os.getenv('ANTHROPIC_API_KEY')
+            self.model = model or os.getenv('CLAUDE_MODEL', 'claude-3-opus-20240229')
+        else:
+            raise ValueError(f"Unsupported LLM provider: {llm_provider}. "
+                           f"Supported providers: 'openai', 'gemini', 'claude'")
+
         self.context = {}
         self.results = {}
 
@@ -49,7 +66,8 @@ def execute(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
     def query_llm(self, prompt: str, system_message: str = None, 
                   temperature: float = 0.7, max_tokens: int = 1000) -> str:
         """
-        Query the LLM (GPT API) for assistance.
+        Query the LLM API for assistance. Supports multiple LLM providers:
+        OpenAI (GPT), Google (Gemini), and Anthropic (Claude).
 
         Args:
             prompt: User prompt for the LLM
@@ -62,34 +80,97 @@ def query_llm(self, prompt: str, system_message: str = None,
         """
         if not self.api_key:
             raise ValueError(
-                "OpenAI API key not found. Set OPENAI_API_KEY environment variable "
-                "or pass api_key during initialization."
+                f"{self.llm_provider.upper()} API key not found. Set the appropriate "
+                f"environment variable or pass api_key during initialization."
             )
 
         try:
-            import openai
-            client = openai.OpenAI(api_key=self.api_key)
-
-            messages = []
-            if system_message:
-                messages.append({"role": "system", "content": system_message})
-            messages.append({"role": "user", "content": prompt})
-
-            response = client.chat.completions.create(
-                model=self.model,
-                messages=messages,
-                temperature=temperature,
-                max_tokens=max_tokens
-            )
-
-            return response.choices[0].message.content
-
-        except ImportError:
+            if self.llm_provider == "openai":
+                return self._query_openai(prompt, system_message, temperature, max_tokens)
+            elif self.llm_provider == "gemini":
+                return self._query_gemini(prompt, system_message, temperature, max_tokens)
+            elif self.llm_provider == "claude":
+                return self._query_claude(prompt, system_message, temperature, max_tokens)
+            else:
+                # This should never happen due to __init__ validation, but handle it anyway
+                raise ValueError(f"Unsupported LLM provider: {self.llm_provider}")
+        except ImportError as e:
             raise ImportError(
-                "OpenAI package not installed. Install with: pip install openai"
+                f"Required package for {self.llm_provider} not installed. "
+                f"Install with: pip install {self._get_package_name()}"
             )
         except Exception as e:
-            raise RuntimeError(f"Error querying LLM: {str(e)}")
+            raise RuntimeError(f"Error querying {self.llm_provider} LLM: {str(e)}")
+
+    def _query_openai(self, prompt: str, system_message: str, 
+                      temperature: float, max_tokens: int) -> str:
+        """Query OpenAI GPT API."""
+        import openai
+        client = openai.OpenAI(api_key=self.api_key)
+
+        messages = []
+        if system_message:
+            messages.append({"role": "system", "content": system_message})
+        messages.append({"role": "user", "content": prompt})
+
+        response = client.chat.completions.create(
+            model=self.model,
+            messages=messages,
+            temperature=temperature,
+            max_tokens=max_tokens
+        )
+
+        return response.choices[0].message.content
+
+    def _query_gemini(self, prompt: str, system_message: str,
+                      temperature: float, max_tokens: int) -> str:
+        """Query Google Gemini API."""
+        import google.generativeai as genai
+        genai.configure(api_key=self.api_key)
+
+        model = genai.GenerativeModel(self.model)
+
+        # Combine system message with prompt for Gemini
+        full_prompt = prompt
+        if system_message:
+            full_prompt = f"{system_message}\n\n{prompt}"
+
+        response = model.generate_content(
+            full_prompt,
+            generation_config=genai.types.GenerationConfig(
+                temperature=temperature,
+                max_output_tokens=max_tokens
+            )
+        )
+
+        return response.text
+
+    def _query_claude(self, prompt: str, system_message: str,
+                      temperature: float, max_tokens: int) -> str:
+        """Query Anthropic Claude API."""
+        import anthropic
+        client = anthropic.Anthropic(api_key=self.api_key)
+
+        message = client.messages.create(
+            model=self.model,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            system=system_message if system_message else "",
+            messages=[
+                {"role": "user", "content": prompt}
+            ]
+        )
+
+        return message.content[0].text
+
+    def _get_package_name(self) -> str:
+        """Get the package name for the current LLM provider."""
+        packages = {
+            "openai": "openai",
+            "gemini": "google-generativeai",
+            "claude": "anthropic"
+        }
+        return packages.get(self.llm_provider, "unknown")
 
     def update_context(self, key: str, value: Any):
         """Update agent's context with new information."""

README.md

@@ -6,7 +6,7 @@ A comprehensive Python package for integrating hydrological model outputs with g
 
 - 🌊 **Hydrological Model Integration:** Seamless loading and processing of MODFLOW and ParFlow outputs  
 - 📊 **ERT Data Processing:** Standardized loading, quality control, and export of ERT field data with RESIPY integration  
-- 🤖 **Multi-Agent AI System:** GPT API-based automated workflow for "load ERT → invert → convert to water content → report" **NEW**
+- 🤖 **Multi-Agent AI System:** Automatic cross-modal geophysics agent supporting multiple LLM APIs (GPT, Gemini, Claude) for automated workflows processing ERT, seismic, and other geophysical data into hydrologic information **NEW**
 - 🪨 **Petrophysical Relationships:** Advanced models for converting between water content, saturation, resistivity, and seismic velocity  
 - ⚡ **Forward Modeling:** Complete ERT and SRT forward modeling capabilities with synthetic data generation  
 - 🔄 **Time-Lapse Inversion:** Sophisticated algorithms for time-lapse ERT inversion with temporal regularization  
@@ -117,16 +117,16 @@ The examples folder provides paired Jupyter notebooks (.ipynb) and Python script
 
 ## 0. Multi-Agent AI Workflow (NEW)
 
-Automate the complete ERT processing workflow using AI agents:
+Automatic cross-modal geophysics agent for subsurface hydrology. Automate geophysical data processing workflows (ERT, seismic, and more) using AI agents with support for multiple LLM APIs (GPT, Gemini, Claude):
 
 ```python
 from PyHydroGeophysX.agents import (
     AgentCoordinator, ERTLoaderAgent, ERTInversionAgent,
     WaterContentAgent, ReportAgent, SeismicAgent
 )
 
-# Initialize coordinator with your OpenAI API key
-coordinator = AgentCoordinator(api_key='your-api-key')
+# Initialize coordinator with your LLM API key (supports OpenAI GPT, Google Gemini, Anthropic Claude)
+coordinator = AgentCoordinator(api_key='your-api-key', llm_provider='openai')  # or 'gemini', 'claude'
 
 # Register specialized agents
 coordinator.register_agent('ert_loader', ERTLoaderAgent())
@@ -160,11 +160,11 @@ if results['status'] == 'success':
 ```
 
 **Key Features:**
-- 🤖 AI-powered parameter selection and interpretation
-- 🔄 Fully automated workflow execution
+- 🤖 AI-powered parameter selection and interpretation with multiple LLM API support (GPT, Gemini, Claude)
+- 🔄 Fully automated workflow execution for cross-modal geophysical data
 - 📊 Automatic quality control and uncertainty quantification
 - 📝 Comprehensive report generation with visualizations
-- 🌊 Optional seismic integration for structural constraints
+- 🌊 Cross-modal integration (ERT, seismic, and other geophysical methods)
 
 See `examples/Ex_multi_agent_workflow.py` for complete examples.