feat: support Anthropic-compatible endpoints for benchmark LLMs

Author: zhangchi47

pyproject.toml

@@ -4,6 +4,7 @@ version = "0.1.0"
 description = "Open Memory Benchmark"
 requires-python = ">=3.11"
 dependencies = [
+    "anthropic>=0.84.0",
     "datasets>=2.0",
     "typer>=0.12",
     "rich>=13",

src/memory_bench/cli.py

@@ -22,12 +22,46 @@
 console = Console()
 
 
-def _resolve_gemini_key() -> None:
-    key = os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY")
-    if not key:
-        typer.echo("Error: GEMINI_API_KEY environment variable is not set.", err=True)
-        raise typer.Exit(1)
-    os.environ["GOOGLE_API_KEY"] = key
+def _ensure_provider_env(provider: str, role: str) -> None:
+    if provider == "anthropic":
+        if not os.environ.get("ANTHROPIC_API_KEY"):
+            typer.echo(f"Error: {role} LLM provider '{provider}' requires ANTHROPIC_API_KEY.", err=True)
+            raise typer.Exit(1)
+        return
+
+    if provider == "gemini":
+        key = os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY")
+        if not key:
+            typer.echo(f"Error: {role} LLM provider '{provider}' requires GEMINI_API_KEY.", err=True)
+            raise typer.Exit(1)
+        os.environ["GOOGLE_API_KEY"] = key
+        return
+
+    if provider == "groq":
+        if not os.environ.get("GROQ_API_KEY"):
+            typer.echo(f"Error: {role} LLM provider '{provider}' requires GROQ_API_KEY.", err=True)
+            raise typer.Exit(1)
+        return
+
+    if provider == "openai":
+        if not os.environ.get("OPENAI_API_KEY"):
+            typer.echo(f"Error: {role} LLM provider '{provider}' requires OPENAI_API_KEY.", err=True)
+            raise typer.Exit(1)
+        return
+
+
+def _validate_run_env(memory: str) -> None:
+    answer_provider = os.environ.get("OMB_ANSWER_LLM", "groq")
+    judge_provider = os.environ.get("OMB_JUDGE_LLM", "gemini")
+    _ensure_provider_env(answer_provider, "Answer")
+    _ensure_provider_env(judge_provider, "Judge")
+
+    if memory == "hindsight":
+        key = os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY")
+        if not key:
+            typer.echo("Error: memory provider 'hindsight' requires GEMINI_API_KEY for embedded extraction.", err=True)
+            raise typer.Exit(1)
+        os.environ["GOOGLE_API_KEY"] = key
 
 
 @app.command()
@@ -53,7 +87,7 @@ def run(
     description: str = typer.Option(None, "--description", "-d", help="Optional description for this run (stored in the result JSON)"),
 ) -> None:
     """Run an evaluation on a single split (optionally filtered to a category)."""
-    _resolve_gemini_key()
+    _validate_run_env(memory)
 
     ds = get_dataset(dataset)
 

src/memory_bench/llm/__init__.py

@@ -1,11 +1,13 @@
 import os
 
+from .anthropic import AnthropicLLM
 from .base import LLM, Schema
 from .gemini import GeminiLLM
 from .groq import GroqLLM
 from .openai import OpenAILLM
 
 REGISTRY: dict[str, type[LLM]] = {
+    "anthropic": AnthropicLLM,
     "gemini": GeminiLLM,
     "groq": GroqLLM,
     "openai": OpenAILLM,

src/memory_bench/llm/anthropic.py

@@ -0,0 +1,142 @@
+import json
+import os
+import re
+import time
+
+from .base import LLM, Schema
+
+_MAX_RETRIES = 6
+_RETRY_BASE_DELAY = 5
+
+
+def _parse_json_payload(text: str) -> dict:
+    text = text.strip()
+
+    try:
+        return json.loads(text)
+    except json.JSONDecodeError:
+        pass
+
+    fenced = re.search(r"```(?:json)?\s*(\{.*\})\s*```", text, flags=re.DOTALL | re.IGNORECASE)
+    if fenced:
+        return json.loads(fenced.group(1))
+
+    start = text.find("{")
+    end = text.rfind("}")
+    if start != -1 and end != -1 and end > start:
+        return json.loads(text[start : end + 1])
+
+    raise json.JSONDecodeError("Could not find JSON object in model response", text, 0)
+
+
+def _coerce_text_payload(text: str, schema: Schema) -> dict | None:
+    text = text.strip()
+    if not text:
+        return None
+
+    result: dict = {}
+    for field in schema.required:
+        spec = schema.properties.get(field, {})
+        field_type = spec.get("type", "string")
+        lowered = text.lower()
+
+        if field_type == "string":
+            result[field] = text
+            continue
+
+        if field_type == "boolean":
+            if re.search(r"\btrue\b", lowered):
+                result[field] = True
+                continue
+            if re.search(r"\bfalse\b", lowered):
+                result[field] = False
+                continue
+            if re.search(r"\b(correct|yes)\b", lowered) and not re.search(r"\b(incorrect|wrong|no)\b", lowered):
+                result[field] = True
+                continue
+            if re.search(r"\b(incorrect|wrong|no)\b", lowered):
+                result[field] = False
+                continue
+            return None
+
+        return None
+
+    return result
+
+
+class AnthropicLLM(LLM):
+    def __init__(self, model: str | None = None):
+        from anthropic import Anthropic
+
+        api_key = os.environ.get("ANTHROPIC_API_KEY")
+        if not api_key:
+            raise RuntimeError("Anthropic provider requires ANTHROPIC_API_KEY")
+
+        base_url = os.environ.get("ANTHROPIC_BASE_URL")
+        self._client = Anthropic(
+            api_key=api_key,
+            base_url=base_url or None,
+            max_retries=0,
+        )
+        self._model = (
+            model
+            or os.environ.get("ANTHROPIC_MODEL")
+            or "claude-sonnet-4-5"
+        )
+
+    @property
+    def model_id(self) -> str:
+        return f"anthropic:{self._model}"
+
+    def generate(self, prompt: str, schema: Schema) -> dict:
+        from anthropic import APIConnectionError, APIStatusError, RateLimitError
+
+        schema_json = {
+            "type": "object",
+            "properties": schema.properties,
+            "required": schema.required,
+            "additionalProperties": False,
+        }
+        system_prompt = (
+            "Return only a valid JSON object matching this schema. "
+            "Do not wrap JSON in markdown fences.\n\n"
+            f"{json.dumps(schema_json, ensure_ascii=False)}"
+        )
+
+        delay = _RETRY_BASE_DELAY
+        last_exc = None
+
+        for attempt in range(_MAX_RETRIES):
+            try:
+                response = self._client.messages.create(
+                    model=self._model,
+                    max_tokens=4096,
+                    temperature=0.0,
+                    system=system_prompt,
+                    messages=[{"role": "user", "content": prompt}],
+                )
+                text = "".join(block.text for block in response.content if getattr(block, "type", None) == "text")
+                try:
+                    return _parse_json_payload(text)
+                except json.JSONDecodeError:
+                    coerced = _coerce_text_payload(text, schema)
+                    if coerced is not None:
+                        return coerced
+                    raise
+            except (RateLimitError, APIConnectionError) as e:
+                last_exc = e
+            except APIStatusError as e:
+                last_exc = e
+                if e.status_code not in (429, 500, 502, 503, 504):
+                    raise
+            except Exception as e:
+                last_exc = e
+                msg = str(e)
+                if "429" not in msg and "rate" not in msg.lower():
+                    raise
+
+            if attempt < _MAX_RETRIES - 1:
+                time.sleep(delay)
+                delay *= 2
+
+        raise RuntimeError(f"Anthropic request failed after {_MAX_RETRIES} retries: {last_exc}")