diff --git a/AGENTS.md b/AGENTS.md index 20c8626..81ba178 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -38,6 +38,7 @@ This is `nui-python-shared-utils`, a Python package providing production-ready u - **Database** (`db_client.py`) - Connection pooling with retry logic - **Metrics** (`cloudwatch_metrics.py`) - Batched CloudWatch publishing with decorators - **JWT Authentication** (`jwt_auth.py`) - RS256 token validation for API Gateway Lambdas (uses `rsa` package) +- **LLM Helper** (`llm.py`) - Generic Anthropic (Claude) client plumbing: API-key/Bedrock auth, forced tool-use call, text call (uses `anthropic`) - **Error Handling** (`error_handler.py`) - Retry patterns with exponential backoff - **Timezone Utils** (`timezone.py`) - Timezone conversion and formatting utilities @@ -50,6 +51,7 @@ The package uses optional extras to minimize Lambda bundle size: - `slack` - Slack SDK - `jwt` - RS256 JWT validation (`rsa` package) - `snowflake` - Pure-Python Snowflake SQL API client (`snowflake-sql-api`) +- `llm` - Anthropic (Claude) client helper (`anthropic[bedrock]`, both API-key and Bedrock IAM auth) - `all` - All integrations - `dev` - Development and testing tools diff --git a/README.md b/README.md index c15b33d..10f256e 100644 --- a/README.md +++ b/README.md @@ -35,6 +35,7 @@ Production-ready shared Python utilities for AWS Lambda functions, CLI tools, an - **Database Connections** - Connection pooling, automatic retries, and transaction management - **CloudWatch Metrics** - Batched publishing with custom dimensions - **JWT Authentication** - RS256 token validation for API Gateway Lambdas (lightweight, no PyJWT needed) +- **Anthropic (Claude) Helper** - Generic client plumbing for LLM calls: API-key or Bedrock IAM auth, forced tool-use, and text calls (prompts and schemas stay in your code) - **Error Handling** - Intelligent retry patterns with exponential backoff - **Timezone Utilities** - Timezone handling and formatting - **Configurable Defaults** - Environment-aware configuration system @@ -65,13 +66,14 @@ Production-ready shared Python utilities for AWS Lambda functions, CLI tools, an pip install nui-python-shared-utils # With specific extras for optional dependencies -pip install nui-python-shared-utils[all] # Core optional integrations (excludes Snowflake) +pip install nui-python-shared-utils[all] # Core optional integrations (excludes Snowflake and LLM) pip install nui-python-shared-utils[powertools] # AWS Powertools only pip install nui-python-shared-utils[slack] # Slack only pip install nui-python-shared-utils[elasticsearch] # Elasticsearch only pip install nui-python-shared-utils[database] # Database only pip install nui-python-shared-utils[jwt] # JWT authentication only pip install nui-python-shared-utils[snowflake] # Snowflake SQL API client +pip install nui-python-shared-utils[llm] # Anthropic (Claude) client helper ``` ### Basic Configuration @@ -260,6 +262,41 @@ def lambda_handler(event, context): **[→ See JWT authentication guide](docs/guides/jwt-authentication.md)** +### Anthropic (Claude) Helper + +Generic plumbing for Claude calls: build a client (API-key or Bedrock IAM), make +a forced tool-use call or a text call, and get the parsed result back. Prompts, +tool schemas, model ids, and result post-processing stay in your code. + +```python +from nui_shared_utils import build_anthropic_client, call_tool + +# API-key auth: explicit key -> ANTHROPIC_API_KEY env -> Secrets Manager +client = build_anthropic_client(secret_name="my/anthropic-key") +# Or Bedrock IAM (no key): build_anthropic_client(mode="bedrock", region="us-east-1") + +tool = { + "name": "classify", + "description": "Classify the text.", + "input_schema": { + "type": "object", + "properties": {"label": {"type": "string"}, "score": {"type": "number"}}, + "required": ["label", "score"], + "additionalProperties": False, + }, +} + +# Forced tool-use; returns the tool's input dict, or None on any model/parse failure. +result = call_tool(client, tool=tool, prompt="Great product!", model="claude-haiku-4-5", max_tokens=256) +if result is not None: + print(result["label"], result["score"]) +``` + +Requires the `[llm]` extra. `call_tool` is best-effort (returns `None`, never +raises); `call_text` returns `{text, input_tokens, output_tokens}`. + +**[→ See full LLM integration guide](docs/guides/llm-integration.md)** + ### Error Handling ```python diff --git a/docs/README.md b/docs/README.md index 32acf16..197511e 100644 --- a/docs/README.md +++ b/docs/README.md @@ -22,6 +22,7 @@ Component-specific guides for major features: - **[Elasticsearch Integration](guides/elasticsearch-integration.md)** - Search, bulk indexing, health checks - **[Snowflake Integration](guides/snowflake-integration.md)** - SQL API client, Secrets Manager credentials, sync/async usage - **[JWT Authentication](guides/jwt-authentication.md)** - RS256 token validation for API Gateway Lambdas +- **[Anthropic (Claude) Integration](guides/llm-integration.md)** - Client helper for API-key/Bedrock auth, forced tool-use, and text calls - **[Log Processing](guides/log-processing.md)** - Kinesis log extraction and ES index naming - Database Connections (planned) - Error Handling Patterns (planned) @@ -127,6 +128,7 @@ When contributing to documentation: - Elasticsearch integration guide (guides/elasticsearch-integration.md) - Snowflake integration guide (guides/snowflake-integration.md) - JWT authentication guide (guides/jwt-authentication.md) +- Anthropic (Claude) integration guide (guides/llm-integration.md) - Shared types reference (guides/shared-types.md) - CLI tools guide (guides/cli-tools.md) - Testing guide (development/testing.md) diff --git a/docs/guides/llm-integration.md b/docs/guides/llm-integration.md new file mode 100644 index 0000000..bece0ca --- /dev/null +++ b/docs/guides/llm-integration.md @@ -0,0 +1,200 @@ +# Anthropic (Claude) Integration + +The LLM helper (`nui_shared_utils.llm`) is generic plumbing for calling +Anthropic's Claude models from Lambda functions and CLI tools. It builds a +client, makes a forced tool-use call or a plain text call, and hands back the +parsed result. + +It is deliberately thin. Prompts, tool schemas, model ids, and any +domain-specific processing of the result stay in your code. The helper owns auth, +the call shape, and result extraction, the parts that were being copy-pasted +across repos. + +Install the optional extra: + +```bash +pip install "nui-python-shared-utils[llm]" +``` + +`anthropic[bedrock]` covers both auth modes (API key and Bedrock IAM). The +`[bedrock]` sub-extra is what makes `anthropic.AnthropicBedrock` importable. + +## What this helper is (and is not) + +| In scope (lives here) | Out of scope (stays in your repo) | +| --------------------------------------------- | ---------------------------------------------- | +| Build a client (API key or Bedrock IAM) | Model id selection | +| Forced tool-use call + `tool_use` extraction | Tool schemas (`input_schema`) | +| Text call + token-usage extraction | Prompts and system prompts | +| Best-effort `None` on tool-call failure | Validation / coercion of the returned dict | + +If you find yourself wanting to add a default model, a prompt, or a tool schema +to this module, it belongs in the consumer instead. + +## Building a Client + +`build_anthropic_client(mode="api_key" | "bedrock", *, api_key=None, secret_name=None, region=None, max_retries=5)` + +### API-key auth (default) + +Returns an `anthropic.Anthropic`. The key resolves in order: + +1. explicit `api_key` argument +2. `ANTHROPIC_API_KEY` environment variable +3. AWS Secrets Manager via `secret_name` (read from the `api_key` field) + +```python +from nui_shared_utils import build_anthropic_client + +# Lambda: key in Secrets Manager +client = build_anthropic_client(secret_name="my-service/anthropic-key") + +# Local dev: key in the environment +client = build_anthropic_client() # reads ANTHROPIC_API_KEY + +# Explicit (e.g. a key resolved from a CLI keyring or a non-default secret field) +client = build_anthropic_client(api_key=my_resolved_key) +``` + +If your secret stores the key under a field other than `api_key`, resolve it +yourself and pass `api_key=`: + +```python +from nui_shared_utils import get_api_key, build_anthropic_client + +key = get_api_key("my-service/creds", key_field="anthropic_api_key") +client = build_anthropic_client(api_key=key) +``` + +### Bedrock IAM auth + +Returns an `anthropic.AnthropicBedrock`. No key, the Lambda's IAM role provides +access. `region` sets `aws_region`; it falls back to `AWS_REGION` / +`AWS_DEFAULT_REGION` and then to the SDK's own default region resolution. + +```python +client = build_anthropic_client(mode="bedrock", region="us-east-1") +``` + +## Forced Tool-Use: `call_tool` + +`call_tool(client, *, tool, prompt, model, max_tokens, system=None) -> dict | None` + +Forces the model to answer through the named tool (`tool_choice` of type `tool`) +and returns that tool's `tool_use.input` dict. The helper only forces the call +and returns the input; it does no validation or coercion itself. Setting +`"strict": True` on your tool's `input_schema` asks the API to constrain the tool +input to that schema, but you still own validating and coercing the returned dict +(value ranges, enum membership, types) in your code. + +```python +from nui_shared_utils import build_anthropic_client, call_tool + +client = build_anthropic_client(secret_name="my-service/anthropic-key") + +# The tool schema is yours; the helper just forces and extracts it. +classify_tool = { + "name": "classify_item", + "description": "Classify a support message.", + "strict": True, + "input_schema": { + "type": "object", + "properties": { + "category": {"type": "string", "enum": ["bug", "question", "feature"]}, + "urgency": {"type": "number", "description": "0.0 low to 1.0 urgent"}, + }, + "required": ["category", "urgency"], + "additionalProperties": False, + }, +} + +result = call_tool( + client, + tool=classify_tool, + prompt="The export button does nothing and I have a deadline.", + model="claude-haiku-4-5", # your choice; the helper imposes no default + max_tokens=256, + system="You triage inbound support messages.", # optional +) + +if result is None: + # Model error, no tool block, or a malformed result. Leave the item + # unprocessed and let the next run retry it. + ... +else: + category = result["category"] # validate/coerce in your code + urgency = result["urgency"] +``` + +### Best-effort contract + +`call_tool` is best-effort and **never raises** on a model or network failure. It +logs a warning and returns `None` when: + +- `client.messages.create` raises any exception (transport, timeout, rate-limit, etc.), +- the response has no `tool_use` block for the named tool, +- the tool's `input` is not an object, +- the tool definition is not a dict or has no `name`. + +This matches the dominant Lambda pattern: a single item failing to enrich must +not abort the batch. If a caller genuinely needs the exception (rather than a +`None`-check), that is a deliberate future addition, not the default. + +## Text Calls: `call_text` + +`call_text(client, *, prompt, model, max_tokens, system=None) -> dict` + +Returns `{"text", "input_tokens", "output_tokens"}`. `text` is the concatenation +of all text content blocks (empty string if the response carried none). Unlike +`call_tool`, this is **not** best-effort: a transport error propagates, because +the caller wanted the text or an exception. + +```python +from nui_shared_utils import build_anthropic_client, call_text + +client = build_anthropic_client(mode="bedrock", region="us-east-1") + +out = call_text( + client, + prompt="Summarize this PDF extract in two sentences:\n\n" + extract, + model="claude-haiku-4-5", + max_tokens=512, +) +print(out["text"]) +print(out["input_tokens"], out["output_tokens"]) # for cost/usage tracking +``` + +## Cold Start and the Optional Extra + +`anthropic` is imported at the top of `nui_shared_utils.llm`, so any use of the +helper requires the `[llm]` extra. Importing the package itself stays cheap: +`import nui_shared_utils` does not import this module (and therefore does not +import `anthropic`) until you first touch an `llm` attribute. This is the same +PEP 562 lazy-loading behaviour the rest of the package uses to keep Lambda +cold-start fast (enforced by `tests/test_lazy_imports.py`). + +Without the extra installed, the top-level lazy exports resolve to `None` +(matching the other optional integrations), so a missing dependency surfaces as a +clear `None` rather than a package import failure: + +```python +import nui_shared_utils as nui + +if nui.build_anthropic_client is None: + raise RuntimeError("install nui-python-shared-utils[llm] to use the LLM helper") +``` + +A direct `from nui_shared_utils.llm import build_anthropic_client` raises +`ImportError` when the extra is missing. + +## Credential Resolution Summary + +Consistent with the other shared clients (Slack, Elasticsearch, Snowflake): + +1. explicit argument (`api_key=`) +2. environment variable (`ANTHROPIC_API_KEY`) +3. AWS Secrets Manager (`secret_name`, `api_key` field) + +Bedrock mode uses IAM instead of any of the above. A consumer needing a CLI +keyring or GPG-backed key resolves it on its own and passes `api_key=`, the +helper does not pull in a CLI credential loader. diff --git a/nui_shared_utils/__init__.py b/nui_shared_utils/__init__.py index 122c556..4c1f6d8 100644 --- a/nui_shared_utils/__init__.py +++ b/nui_shared_utils/__init__.py @@ -116,6 +116,10 @@ "get_jwt_public_key": ("jwt_auth", "get_jwt_public_key"), "JWTValidationError": ("jwt_auth", "JWTValidationError"), "AuthenticationError": ("jwt_auth", "AuthenticationError"), + # Optional: Anthropic (Claude) LLM helper (anthropic[bedrock]) + "build_anthropic_client": ("llm", "build_anthropic_client"), + "call_tool": ("llm", "call_tool"), + "call_text": ("llm", "call_text"), } # Submodules that are optional integrations; ImportError during lazy load @@ -130,6 +134,7 @@ "powertools_helpers", "jwt_auth", "snowflake_client", + "llm", "slack_setup", } @@ -269,6 +274,7 @@ def __dir__() -> List[str]: require_auth, validate_jwt, ) + from .llm import build_anthropic_client, call_text, call_tool from . import slack_setup diff --git a/nui_shared_utils/llm.py b/nui_shared_utils/llm.py new file mode 100644 index 0000000..026b5f8 --- /dev/null +++ b/nui_shared_utils/llm.py @@ -0,0 +1,209 @@ +"""Generic Anthropic (Claude) client helper for AWS Lambda functions and CLI tools. + +Plumbing only. Three repos independently hand-rolled the same Anthropic client + +forced-tool-use call + ``tool_use``-block parse; this consolidates that into one +tested surface. Prompts, tool schemas, model ids, and any domain post-processing +of the result stay in the consuming repo. + +What lives here: + +- :func:`build_anthropic_client` - API-key auth (explicit -> ``ANTHROPIC_API_KEY`` + env -> Secrets Manager) or Bedrock IAM auth, one function for both. +- :func:`call_tool` - a forced tool-use call that returns the named tool's + ``input`` dict, or ``None`` on any model/parse failure (best-effort; never + raises, matching the dominant Lambda enrichment pattern). +- :func:`call_text` - a plain text call returning ``{text, input_tokens, + output_tokens}``. + +The ``anthropic`` SDK is a module-level import (matching the package's other +optional integrations), so the bare ``[llm]`` extra is required to use anything +here. The cold-start guarantee is preserved by the package ``__init__`` PEP 562 +lazy loader: ``import nui_shared_utils`` never imports this module (and so never +imports ``anthropic``) until a consumer first touches an ``llm`` attribute. See +``tests/test_lazy_imports.py``. + +Install the extra:: + + pip install 'nui-python-shared-utils[llm]' + +``anthropic[bedrock]`` covers both auth modes (``anthropic.Anthropic`` for the +API key, ``anthropic.AnthropicBedrock`` for IAM); the ``[bedrock]`` sub-extra is +what makes ``AnthropicBedrock`` available. +""" + +from __future__ import annotations + +import logging +import os +from typing import Any, Dict, Optional + +import anthropic + +from .secrets_helper import get_api_key + +log = logging.getLogger(__name__) + + +def build_anthropic_client( + mode: str = "api_key", + *, + api_key: Optional[str] = None, + secret_name: Optional[str] = None, + region: Optional[str] = None, + max_retries: int = 5, +) -> Any: + """Build an Anthropic SDK client for the requested auth mode. + + ``mode="api_key"`` (default) returns an :class:`anthropic.Anthropic`. The key + is resolved in order: the explicit ``api_key`` argument, then the + ``ANTHROPIC_API_KEY`` environment variable, then + :func:`nui_shared_utils.get_api_key` against ``secret_name`` (AWS Secrets + Manager, ``api_key`` field). A consumer whose key lives behind a CLI keyring + or a non-default secret field should resolve it and pass ``api_key=``. + + ``mode="bedrock"`` returns an :class:`anthropic.AnthropicBedrock` (IAM auth, + no key). ``region`` sets ``aws_region``; it falls back to ``AWS_REGION`` / + ``AWS_DEFAULT_REGION`` and then to the SDK's own default region resolution. + + Args: + mode: ``"api_key"`` or ``"bedrock"``. + api_key: Explicit Anthropic API key (api_key mode only). Skips env and + Secrets Manager when set. + secret_name: Secrets Manager secret name to read the key from when no + explicit key and no env var are present (api_key mode only). + region: AWS region for Bedrock (bedrock mode only). + max_retries: Passed through to the SDK client (default 5). + + Raises: + ValueError: Unknown ``mode``, or api_key mode with no key resolvable + (no ``api_key``, no env var, and no ``secret_name``). + """ + if mode == "bedrock": + aws_region = region or os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION") + return anthropic.AnthropicBedrock(aws_region=aws_region, max_retries=max_retries) + + if mode == "api_key": + key = api_key or os.environ.get("ANTHROPIC_API_KEY") + if not key: + if not secret_name: + raise ValueError( + "build_anthropic_client(mode='api_key'): no api_key argument, " + "no ANTHROPIC_API_KEY env var, and no secret_name to read from " + "Secrets Manager" + ) + key = get_api_key(secret_name, key_field="api_key") + return anthropic.Anthropic(api_key=key, max_retries=max_retries) + + raise ValueError(f"build_anthropic_client: unknown mode {mode!r} (expected 'api_key' or 'bedrock')") + + +def call_tool( + client: Any, + *, + tool: Dict[str, Any], + prompt: str, + model: str, + max_tokens: int, + system: Optional[str] = None, +) -> Optional[Dict[str, Any]]: + """Make a forced tool-use call and return the named tool's ``input`` dict. + + Best-effort by design: a transport error, a malformed tool definition, a + response with no matching ``tool_use`` block, or a non-object tool input + logs a warning and returns ``None`` so the caller can leave the item + unprocessed and retry on the next run. It never raises on a model or network + failure (the dominant Lambda enrichment pattern). Validation and any + coercion of the returned dict stay in the consumer. + + Args: + client: An Anthropic client (from :func:`build_anthropic_client` or any + object exposing ``messages.create``). + tool: A single Anthropic tool definition. Must carry a ``name``. + prompt: The user-turn prompt text. + model: Model id (kept in the consumer; this helper imposes no default). + max_tokens: Response token cap. + system: Optional system prompt. + + Returns: + The tool's ``input`` dict, or ``None`` on any failure. + """ + if not isinstance(tool, dict): + log.warning("call_tool: tool is not a dict; cannot force tool use") + return None + tool_name = str(tool.get("name") or "").strip() + if not tool_name: + log.warning("call_tool: tool definition has no usable 'name'; cannot force tool use") + return None + + kwargs: Dict[str, Any] = { + "model": model, + "max_tokens": max_tokens, + "messages": [{"role": "user", "content": prompt}], + "tools": [tool], + "tool_choice": {"type": "tool", "name": tool_name}, + } + if system is not None: + kwargs["system"] = system + + try: + response = client.messages.create(**kwargs) + except Exception as exc: # noqa: BLE001 - best-effort; never abort the caller's run + log.warning("call_tool: model call failed for tool '%s': %s", tool_name, exc) + return None + + for block in getattr(response, "content", None) or []: + if getattr(block, "type", None) == "tool_use" and getattr(block, "name", "") == tool_name: + tool_input = block.input + if isinstance(tool_input, dict): + return tool_input + log.warning("call_tool: tool_use input for '%s' was not an object", tool_name) + return None + + log.warning("call_tool: response contained no tool_use block for '%s'", tool_name) + return None + + +def call_text( + client: Any, + *, + prompt: str, + model: str, + max_tokens: int, + system: Optional[str] = None, +) -> Dict[str, Any]: + """Make a plain text call and return the text plus token usage. + + Returns ``{"text", "input_tokens", "output_tokens"}``. ``text`` is the + concatenation of all ``text`` content blocks (empty string if the response + carried none). Unlike :func:`call_tool` this is not best-effort: a transport + error propagates to the caller, who wanted the text or an exception. + + Args: + client: An Anthropic client (from :func:`build_anthropic_client` or any + object exposing ``messages.create``). + prompt: The user-turn prompt text. + model: Model id (kept in the consumer; this helper imposes no default). + max_tokens: Response token cap. + system: Optional system prompt. + + Returns: + ``{"text": str, "input_tokens": int | None, "output_tokens": int | None}``. + """ + kwargs: Dict[str, Any] = { + "model": model, + "max_tokens": max_tokens, + "messages": [{"role": "user", "content": prompt}], + } + if system is not None: + kwargs["system"] = system + + response = client.messages.create(**kwargs) + text = "".join( + block.text for block in (getattr(response, "content", None) or []) if getattr(block, "type", None) == "text" + ) + usage = getattr(response, "usage", None) + return { + "text": text, + "input_tokens": getattr(usage, "input_tokens", None), + "output_tokens": getattr(usage, "output_tokens", None), + } diff --git a/pyproject.toml b/pyproject.toml index 46ec43a..c897a58 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,6 +51,10 @@ jwt = ["rsa>=4.9"] snowflake = [ "snowflake-sql-api @ git+https://github.com/hampsterx/snowflake-sql-api.git@v0.1.1", ] +# Anthropic (Claude) helper. anthropic[bedrock] covers both auth modes (API key +# + Bedrock IAM). Kept out of `all` (same as snowflake): the SDK is specialized +# and opt-in. Keep in sync with setup.py. +llm = ["anthropic[bedrock]>=0.45.0,<1.0.0"] all = [ "elasticsearch>=7.17.0,<8.0.0", "pymysql>=1.0.0", @@ -75,6 +79,7 @@ dev = [ "rsa>=4.9", "cryptography>=41.0.0", "snowflake-sql-api @ git+https://github.com/hampsterx/snowflake-sql-api.git@v0.1.1", + "anthropic[bedrock]>=0.45.0,<1.0.0", ] [project.urls] diff --git a/requirements-test.txt b/requirements-test.txt index 2709268..18cffa1 100644 --- a/requirements-test.txt +++ b/requirements-test.txt @@ -28,3 +28,6 @@ cryptography>=41.0.0 # snowflake adapter (tag pin; see pyproject [snowflake] extra) snowflake-sql-api @ git+https://github.com/hampsterx/snowflake-sql-api.git@v0.1.1 + +# Anthropic (Claude) helper (see pyproject [llm] extra) +anthropic[bedrock]>=0.45.0,<1.0.0 diff --git a/setup.py b/setup.py index 4709b02..d6c0209 100644 --- a/setup.py +++ b/setup.py @@ -64,6 +64,10 @@ "snowflake": [ "snowflake-sql-api @ git+https://github.com/hampsterx/snowflake-sql-api.git@50205b0cb63df008c9c453ee05f0f130dcfe4805", ], + # Anthropic (Claude) helper. anthropic[bedrock] covers both auth modes + # (API key + Bedrock IAM). Kept out of "all" (same as snowflake): the + # SDK is specialized and opt-in. Keep in sync with pyproject.toml. + "llm": ["anthropic[bedrock]>=0.45.0,<1.0.0"], "all": [ "elasticsearch>=7.17.0,<8.0.0", "pymysql>=1.0.0", @@ -87,6 +91,7 @@ "build>=0.8.0", "rsa>=4.9", "cryptography>=41.0.0", + "anthropic[bedrock]>=0.45.0,<1.0.0", ], }, python_requires=">=3.9", diff --git a/tests/test_lazy_imports.py b/tests/test_lazy_imports.py index 2bce4b7..b27e7f7 100644 --- a/tests/test_lazy_imports.py +++ b/tests/test_lazy_imports.py @@ -34,6 +34,25 @@ def test_bare_package_import_does_not_load_boto3(): assert "botocore=False" in output +def test_bare_package_import_does_not_load_anthropic(): + """``import nui_shared_utils`` must not transitively import anthropic. + + The ``llm`` helper imports ``anthropic`` at module top, so the cold-start + guarantee rests entirely on the PEP 562 lazy loader: the ``llm`` submodule + (and therefore ``anthropic``) must not load until a consumer first touches + an ``llm`` attribute. anthropic is installed via the ``[dev]`` extra, so this + assertion is meaningful rather than trivially true. + """ + output = _run(""" + import sys + import nui_shared_utils # noqa: F401 + print(f"anthropic={'anthropic' in sys.modules}") + print(f"llm={'nui_shared_utils.llm' in sys.modules}") + """) + assert "anthropic=False" in output + assert "llm=False" in output + + def test_jwt_auth_import_does_not_load_boto3(): """``from nui_shared_utils.jwt_auth import check_auth`` must stay lazy. diff --git a/tests/test_llm.py b/tests/test_llm.py new file mode 100644 index 0000000..7110fe8 --- /dev/null +++ b/tests/test_llm.py @@ -0,0 +1,227 @@ +"""Tests for the Anthropic (Claude) helper (nui_shared_utils.llm). + +Covers: +- ``build_anthropic_client`` auth-mode selection (api_key explicit/env/secret, + bedrock IAM) and the no-key error path, +- ``call_tool`` forced tool-use happy path, the ``None`` paths (no tool block, + wrong tool name, non-object input, transport error, nameless tool), +- ``call_text`` text concatenation and usage extraction. + +The ``anthropic`` client is mocked for the call helpers (a ``MagicMock`` with a +``messages.create``), so these are fast and offline. ``build_anthropic_client`` +tests construct real SDK client objects, which is offline too (the SDK only +defers network calls to the first request); they need the ``[llm]`` extra, which +is in the ``[dev]`` extra used to run the suite. +""" + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import anthropic +import pytest + +from nui_shared_utils import llm + +pytestmark = pytest.mark.unit + +# A minimal, generic tool definition. No domain content (prompts/schemas stay in +# the consumer); this is just enough to exercise the forced tool-use path. +TOOL = { + "name": "record_thing", + "description": "Record a thing.", + "input_schema": { + "type": "object", + "properties": {"score": {"type": "number"}}, + "required": ["score"], + "additionalProperties": False, + }, +} + + +def _client_returning(response): + """A mock Anthropic client whose ``messages.create`` returns ``response``.""" + client = MagicMock() + client.messages.create.return_value = response + return client + + +@pytest.fixture +def no_anthropic_env(monkeypatch): + """Ensure ANTHROPIC_API_KEY does not leak in from the host environment.""" + monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False) + + +# --------------------------------------------------------------------------- +# build_anthropic_client +# --------------------------------------------------------------------------- + + +class TestBuildAnthropicClient: + def test_api_key_explicit_wins(self, no_anthropic_env): + client = llm.build_anthropic_client(api_key="sk-explicit") + assert isinstance(client, anthropic.Anthropic) + assert client.api_key == "sk-explicit" + + def test_api_key_from_env(self, monkeypatch): + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-from-env") + client = llm.build_anthropic_client() + assert isinstance(client, anthropic.Anthropic) + assert client.api_key == "sk-from-env" + + def test_api_key_from_secrets_manager(self, no_anthropic_env, mocker): + get_api_key = mocker.patch.object(llm, "get_api_key", return_value="sk-from-secret") + client = llm.build_anthropic_client(secret_name="my/anthropic-key") + assert isinstance(client, anthropic.Anthropic) + assert client.api_key == "sk-from-secret" + get_api_key.assert_called_once_with("my/anthropic-key", key_field="api_key") + + def test_explicit_key_skips_secrets_manager(self, no_anthropic_env, mocker): + get_api_key = mocker.patch.object(llm, "get_api_key", return_value="sk-from-secret") + client = llm.build_anthropic_client(api_key="sk-explicit", secret_name="my/anthropic-key") + assert client.api_key == "sk-explicit" + get_api_key.assert_not_called() + + def test_api_key_missing_raises(self, no_anthropic_env): + with pytest.raises(ValueError, match="no api_key"): + llm.build_anthropic_client() + + def test_bedrock_mode(self, monkeypatch): + # Dummy AWS env so boto3 (used by the Bedrock client for signing) has a + # region and credentials to resolve; construction stays offline. + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "testing") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "testing") + client = llm.build_anthropic_client(mode="bedrock", region="us-east-1") + assert isinstance(client, anthropic.AnthropicBedrock) + + def test_unknown_mode_raises(self, no_anthropic_env): + with pytest.raises(ValueError, match="unknown mode"): + llm.build_anthropic_client(mode="totally-not-a-mode") + + +# --------------------------------------------------------------------------- +# call_tool +# --------------------------------------------------------------------------- + + +class TestCallTool: + def test_happy_path_returns_tool_input(self): + response = SimpleNamespace( + content=[SimpleNamespace(type="tool_use", name="record_thing", input={"score": 0.9})], + ) + client = _client_returning(response) + + result = llm.call_tool(client, tool=TOOL, prompt="hi", model="claude-x", max_tokens=256) + + assert result == {"score": 0.9} + client.messages.create.assert_called_once() + kwargs = client.messages.create.call_args.kwargs + assert kwargs["model"] == "claude-x" + assert kwargs["max_tokens"] == 256 + assert kwargs["tools"] == [TOOL] + assert kwargs["tool_choice"] == {"type": "tool", "name": "record_thing"} + assert kwargs["messages"] == [{"role": "user", "content": "hi"}] + assert "system" not in kwargs + + def test_system_prompt_is_forwarded(self): + response = SimpleNamespace( + content=[SimpleNamespace(type="tool_use", name="record_thing", input={"score": 0.1})], + ) + client = _client_returning(response) + + llm.call_tool(client, tool=TOOL, prompt="hi", model="m", max_tokens=64, system="be terse") + + assert client.messages.create.call_args.kwargs["system"] == "be terse" + + def test_no_tool_block_returns_none(self): + response = SimpleNamespace(content=[SimpleNamespace(type="text", text="no tool here")]) + result = llm.call_tool(_client_returning(response), tool=TOOL, prompt="hi", model="m", max_tokens=64) + assert result is None + + def test_wrong_tool_name_returns_none(self): + response = SimpleNamespace( + content=[SimpleNamespace(type="tool_use", name="some_other_tool", input={"x": 1})], + ) + result = llm.call_tool(_client_returning(response), tool=TOOL, prompt="hi", model="m", max_tokens=64) + assert result is None + + def test_non_object_input_returns_none(self): + response = SimpleNamespace( + content=[SimpleNamespace(type="tool_use", name="record_thing", input="not-a-dict")], + ) + result = llm.call_tool(_client_returning(response), tool=TOOL, prompt="hi", model="m", max_tokens=64) + assert result is None + + def test_transport_error_returns_none(self): + client = MagicMock() + client.messages.create.side_effect = RuntimeError("boom") + result = llm.call_tool(client, tool=TOOL, prompt="hi", model="m", max_tokens=64) + assert result is None + + def test_nameless_tool_returns_none_without_calling(self): + client = MagicMock() + result = llm.call_tool(client, tool={"description": "no name"}, prompt="hi", model="m", max_tokens=64) + assert result is None + client.messages.create.assert_not_called() + + def test_null_name_returns_none_without_calling(self): + client = MagicMock() + result = llm.call_tool(client, tool={"name": None}, prompt="hi", model="m", max_tokens=64) + assert result is None + client.messages.create.assert_not_called() + + def test_non_dict_tool_returns_none_without_calling(self): + client = MagicMock() + result = llm.call_tool(client, tool="not-a-dict", prompt="hi", model="m", max_tokens=64) + assert result is None + client.messages.create.assert_not_called() + + +# --------------------------------------------------------------------------- +# call_text +# --------------------------------------------------------------------------- + + +class TestCallText: + def test_happy_path_concatenates_text_and_extracts_usage(self): + response = SimpleNamespace( + content=[ + SimpleNamespace(type="text", text="hello"), + SimpleNamespace(type="text", text=" world"), + ], + usage=SimpleNamespace(input_tokens=11, output_tokens=7), + ) + client = _client_returning(response) + + result = llm.call_text(client, prompt="hi", model="claude-x", max_tokens=512) + + assert result == {"text": "hello world", "input_tokens": 11, "output_tokens": 7} + kwargs = client.messages.create.call_args.kwargs + assert kwargs["model"] == "claude-x" + assert kwargs["messages"] == [{"role": "user", "content": "hi"}] + assert "tools" not in kwargs + assert "system" not in kwargs + + def test_system_prompt_is_forwarded(self): + response = SimpleNamespace( + content=[SimpleNamespace(type="text", text="ok")], + usage=SimpleNamespace(input_tokens=1, output_tokens=1), + ) + client = _client_returning(response) + + llm.call_text(client, prompt="hi", model="m", max_tokens=64, system="be terse") + + assert client.messages.create.call_args.kwargs["system"] == "be terse" + + def test_no_text_blocks_returns_empty_string(self): + response = SimpleNamespace( + content=[SimpleNamespace(type="tool_use", name="x", input={})], + usage=SimpleNamespace(input_tokens=3, output_tokens=0), + ) + result = llm.call_text(_client_returning(response), prompt="hi", model="m", max_tokens=64) + assert result == {"text": "", "input_tokens": 3, "output_tokens": 0} + + def test_transport_error_propagates(self): + client = MagicMock() + client.messages.create.side_effect = RuntimeError("boom") + with pytest.raises(RuntimeError, match="boom"): + llm.call_text(client, prompt="hi", model="m", max_tokens=64)