Feat/extend llm apis

pyproject.toml

@@ -45,6 +45,8 @@ numpy = "*"
 matplotlib = "*"
 python-dotenv = "^1.0"
 google-genai = "*"
+openai = "*"
+anthropic = "*"
 pydantic = "^2.0"
 
 
@@ -64,6 +66,7 @@ pytest-mock = "^3.14.0"
 detect-secrets = "^1.5.0"
 pytest-env = "^1.1.5"
 pylint-per-file-ignores = "^3.2.0"
+ipykernel = "^7.2.0"
 
 [tool.black]
 line-length = 79

src/gee_mcp/server/analysis.py

@@ -8,15 +8,15 @@
 from loguru import logger
 
 from .coderun import _execute_gee_python
-from .genai import init_genai_client
 from .helpers import extract_tag, extract_xml_tag
+from .llm import init_llm_client
 
 
 def _get_datasets_locations_and_periods(
     question: str,
     gee_datasets: list[dict] = None,
 ) -> dict:
-    genai_client = init_genai_client()
+    genai_client = init_llm_client()
 
     dataset_instructions = (
         f"""
@@ -86,7 +86,7 @@ def _get_datasets_locations_and_periods(
 
 
 def _extract_factuality_issues(question: str, python_code: str) -> str:
-    genai_client = init_genai_client()
+    genai_client = init_llm_client()
 
     prompt = f"""
     You are a helpful assistant for Earth Observation data analysis with Google Earth Engine.
@@ -263,7 +263,7 @@ def __init__(
         python_code_result,
         n_samples_per_code_variable=3,
     ):
-        self.genai_client = init_genai_client()
+        self.genai_client = init_llm_client()
 
         self.question = question
         self.python_code = python_code
@@ -684,7 +684,7 @@ async def _assess_factuality_issue(
     </CODE_RECOMMENDATIONS>
     """
 
-    genai_client = init_genai_client()
+    genai_client = init_llm_client()
     r = genai_client.call(prompt)
 
     code_recommendations = extract_xml_tag(r["answer"], "CODE_RECOMMENDATIONS")

src/gee_mcp/server/codegen.py

@@ -6,13 +6,13 @@
 from loguru import logger
 
 from .coderun import GEEPythonExecution
-from .genai import init_genai_client
 from .helpers import (
     NoTagFoundError,
     extract_tag,
     extract_xml_tag,
     remove_leading_spaces,
 )
+from .llm import init_llm_client
 
 
 class QuestionRecord:
@@ -179,7 +179,7 @@ def __init__(
 
         self.question_record = question_record
         self.qr = self.question_record
-        self.genai_client = init_genai_client()
+        self.genai_client = init_llm_client()
         self.remarks_for_prompts = ""
         self.gee_dataset_list = gee_dataset_list
         self.number_of_fix_iterations = number_of_fix_iterations

src/gee_mcp/server/coderun.py

@@ -2,13 +2,13 @@
 
 from loguru import logger
 
-from .genai import init_genai_client
 from .helpers import extract_xml_tag
+from .llm import init_llm_client
 
 
 class GEEPythonExecution:
     def __init__(self, genai_client=None):
-        self.genai_client = genai_client or init_genai_client()
+        self.genai_client = genai_client or init_llm_client()
 
     def exec(self, code):
         namespace: dict = {}

src/gee_mcp/server/llm/__init__.py

@@ -0,0 +1,24 @@
+# Import provider modules for their import-time side effect: each one's
+# @register_llm decorator populates base._LLM_REGISTRY. Without these imports
+# the registry is empty and init_llm_client() can't resolve a provider.
+from .anthropic_provider import AnthropicLLM
+from .base import BaseLLM, LLMProvider, register_llm
+from .cache import JSONFileCache, NullCache, ResponseCache
+from .factory import init_llm_client
+from .google_provider import GoogleLLM
+from .openai_provider import OpenAILLM
+from .types import LLMCallReturn
+
+__all__ = [
+    "BaseLLM",
+    "LLMProvider",
+    "LLMCallReturn",
+    "ResponseCache",
+    "NullCache",
+    "JSONFileCache",
+    "register_llm",
+    "init_llm_client",
+    "AnthropicLLM",
+    "GoogleLLM",
+    "OpenAILLM",
+]

src/gee_mcp/server/llm/anthropic_provider.py

@@ -0,0 +1,62 @@
+import os
+
+import anthropic
+
+from .base import BaseLLM, LLMProvider, register_llm
+from .cache import ResponseCache
+from .types import LLMCallReturn
+
+
+@register_llm(LLMProvider.ANTHROPIC)
+class AnthropicLLM(BaseLLM):
+    # Anthropic requires an explicit max output token count; adaptive thinking
+    # plus tool-free prose answers fit comfortably under this.
+    MAX_TOKENS = 16000
+
+    def __init__(
+        self,
+        api_key,
+        model,
+        cache: ResponseCache | None = None,
+    ):
+        super().__init__(api_key=api_key, model=model, cache=cache)
+        self.client = anthropic.Anthropic(api_key=api_key)
+
+    @classmethod
+    def from_env(
+        cls, model: str, cache: ResponseCache | None = None
+    ) -> "AnthropicLLM":
+        api_key = os.getenv("ANTHROPIC_API_KEY")
+        if api_key is None:
+            raise ValueError("API key for Anthropic not found.")
+        return cls(api_key=api_key, model=model, cache=cache)
+
+    def _call(self, text: str, include_thinking: bool = True) -> LLMCallReturn:
+        thinking_conf: anthropic.types.ThinkingConfigParam
+        if include_thinking:
+            # Adaptive thinking is the only supported "on" mode on Opus 4.7;
+            # "summarized" surfaces the reasoning text instead of omitting it.
+            thinking_conf = {"type": "adaptive", "display": "summarized"}
+        else:
+            thinking_conf = {"type": "disabled"}
+
+        messages: list[anthropic.types.MessageParam] = [
+            {"role": "user", "content": text}
+        ]
+        response = self.client.messages.create(
+            model=self.model,
+            max_tokens=self.MAX_TOKENS,
+            thinking=thinking_conf,
+            messages=messages,
+        )
+
+        thought = None
+        answer = None
+
+        for block in response.content:
+            if block.type == "thinking":
+                thought = block.thinking
+            elif block.type == "text":
+                answer = block.text
+
+        return {"answer": answer, "thought": thought, "response": response}

src/gee_mcp/server/llm/base.py

@@ -0,0 +1,68 @@
+from abc import ABC, abstractmethod
+from enum import Enum
+
+from .cache import NullCache, ResponseCache
+from .types import LLMCallReturn
+
+
+class LLMProvider(Enum):
+    GOOGLE = "google"
+    ANTHROPIC = "anthropic"
+    OPENAI = "openai"
+
+
+# provider -> implementing class, populated by the @register_llm decorator
+_LLM_REGISTRY: dict[LLMProvider, type["BaseLLM"]] = {}
+
+
+def register_llm(provider: LLMProvider):
+    """Class decorator: tag a `BaseLLM` subclass with its provider and register it."""
+
+    def deco(cls: type["BaseLLM"]) -> type["BaseLLM"]:
+        cls._provider = provider
+        _LLM_REGISTRY[provider] = cls
+        return cls
+
+    return deco
+
+
+class BaseLLM(ABC):
+    _provider: LLMProvider | None = None
+
+    def __init__(
+        self,
+        api_key: str | None,
+        model: str,
+        cache: ResponseCache | None = None,
+    ):
+        if self._provider is None:
+            raise RuntimeError(
+                f"{type(self).__name__} must set a class-level `_provider`"
+            )
+
+        self.api_key = api_key
+        self.model = model
+        self.cache: ResponseCache = cache if cache is not None else NullCache()
+
+    @classmethod
+    @abstractmethod
+    def from_env(
+        cls, model: str, cache: ResponseCache | None = None
+    ) -> "BaseLLM":
+        """Build an instance using credentials/config from environment variables."""
+
+    def _cache_key(self, text: str, include_thinking: bool) -> str:
+        return f"{self.model}::{include_thinking}::{text}"
+
+    def call(self, text: str, include_thinking: bool = True) -> LLMCallReturn:
+        key = self._cache_key(text, include_thinking)
+        cached = self.cache.get(key)
+        if cached is not None:
+            return cached
+        call_return = self._call(text=text, include_thinking=include_thinking)
+        self.cache.put(key, call_return)
+        return call_return
+
+    @abstractmethod
+    def _call(self, text: str, include_thinking: bool = True) -> LLMCallReturn:
+        pass