@pentatonic-ai/ai-agent-sdk 0.7.12 → 0.8.0

package/packages/memory-engine/tests/test_embed_provider.py

@@ -0,0 +1,354 @@
+"""Unit tests for engine/services/_shared/embed_provider.py.
+
+Run with:
+    cd packages/memory-engine
+    python -m pytest tests/test_embed_provider.py -v
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+# Make the engine/services tree importable for tests without packaging it.
+ROOT = Path(__file__).parent.parent / "engine" / "services"
+sys.path.insert(0, str(ROOT))
+
+import json  # noqa: E402
+
+import httpx  # noqa: E402
+import pytest  # noqa: E402
+
+from _shared.embed_provider import (  # noqa: E402
+    PROVIDERS,
+    EmbedAuthError,
+    EmbedClient,
+    EmbedHTTPError,
+    EmbedProvider,
+    resolve_provider,
+)
+
+
+# ----------------------------------------------------------------------
+# Helpers — stub httpx so we can assert the request shape.
+# ----------------------------------------------------------------------
+
+class _FakeResponse:
+    def __init__(self, status_code: int, payload: dict | str = ""):
+        self.status_code = status_code
+        if isinstance(payload, dict):
+            self._json = payload
+            self.text = json.dumps(payload)
+        else:
+            self._json = None
+            self.text = payload
+
+    @property
+    def is_success(self) -> bool:
+        return 200 <= self.status_code < 300
+
+    def json(self) -> dict:
+        if self._json is None:
+            raise ValueError("not json")
+        return self._json
+
+
+class _Recorder:
+    """Records every httpx.post call and returns canned responses keyed by URL."""
+
+    def __init__(self):
+        self.calls: list[dict] = []
+        self.responses: dict[str, _FakeResponse] = {}
+
+    def respond(self, url: str, response: _FakeResponse) -> None:
+        self.responses[url] = response
+
+    def __call__(self, url, *, json, headers, timeout):
+        self.calls.append({"url": url, "json": json, "headers": headers, "timeout": timeout})
+        if url in self.responses:
+            return self.responses[url]
+        # default: 401 to flush out unmatched URLs
+        return _FakeResponse(401, "no stub for this url")
+
+
+@pytest.fixture
+def recorder(monkeypatch):
+    rec = _Recorder()
+    monkeypatch.setattr(httpx, "post", rec)
+    return rec
+
+
+# ----------------------------------------------------------------------
+# Provider resolution
+# ----------------------------------------------------------------------
+
+def test_resolve_built_in_providers():
+    for name in ("openai", "pentatonic-gateway", "cohere"):
+        p = resolve_provider(name)
+        assert p.name == name
+
+
+def test_resolve_unknown_provider_raises():
+    with pytest.raises(ValueError):
+        resolve_provider("not-a-provider")
+
+
+def test_resolve_custom_provider_from_env(monkeypatch):
+    monkeypatch.setenv("L4_EMBED_AUTH_HEADER", "X-Custom-Auth")
+    monkeypatch.setenv("L4_EMBED_AUTH_FORMAT", "Token {key}")
+    monkeypatch.setenv("L4_EMBED_PATH_DEFAULT", "/embed")
+    monkeypatch.setenv("L4_EMBED_BODY_SHAPE", "cohere")
+    monkeypatch.setenv("L4_EMBED_RESPONSE_SHAPE", "cohere")
+    p = resolve_provider("custom", env_prefix="L4_")
+    assert p.auth_header == "X-Custom-Auth"
+    assert p.auth_format == "Token {key}"
+    assert p.path_default == "/embed"
+    # body shape produces Cohere-style "texts" field
+    body = p.body_builder(["hi"], "model-x")
+    assert body == {"texts": ["hi"], "model": "model-x", "input_type": "search_document"}
+
+
+# ----------------------------------------------------------------------
+# Request shape
+# ----------------------------------------------------------------------
+
+def test_openai_provider_request_shape(recorder):
+    recorder.respond(
+        "https://gw/v1/embeddings",
+        _FakeResponse(200, {"data": [{"embedding": [0.1, 0.2]}]}),
+    )
+    client = EmbedClient(
+        url="https://gw/v1/embeddings",
+        api_key="k",
+        model="m",
+        provider=PROVIDERS["openai"],
+    )
+    out = client.embed_batch(["hello"])
+    assert out == [[0.1, 0.2]]
+    call = recorder.calls[0]
+    assert call["url"] == "https://gw/v1/embeddings"
+    assert call["json"] == {"input": ["hello"], "model": "m"}
+    assert call["headers"] == {"Authorization": "Bearer k"}
+
+
+def test_pentatonic_provider_request_shape(recorder):
+    recorder.respond(
+        "https://lambda-gateway.pentatonic.com/v1/embed",
+        _FakeResponse(200, {"data": [{"embedding": [1.0, 2.0]}]}),
+    )
+    client = EmbedClient(
+        url="https://lambda-gateway.pentatonic.com/v1/embed",
+        api_key="secret",
+        model="nv-embed-v2",
+        provider=PROVIDERS["pentatonic-gateway"],
+    )
+    out = client.embed_batch(["t1"])
+    assert out == [[1.0, 2.0]]
+    call = recorder.calls[0]
+    assert call["url"] == "https://lambda-gateway.pentatonic.com/v1/embed"
+    assert call["json"] == {"input": ["t1"], "model": "nv-embed-v2"}
+    assert call["headers"] == {"X-API-Key": "secret"}
+
+
+def test_pentatonic_response_parser_handles_both_shapes(recorder):
+    """Pentatonic Gateway has historically returned both {"data":[...]} and
+    {"embeddings":[...]} on different endpoints. Parser accepts either."""
+    p = PROVIDERS["pentatonic-gateway"]
+    assert p.response_parser({"data": [{"embedding": [1.0]}]}) == [[1.0]]
+    assert p.response_parser({"embeddings": [[1.0]]}) == [[1.0]]
+
+
+def test_cohere_provider_request_shape(recorder):
+    recorder.respond(
+        "https://api.cohere.ai/v1/embed",
+        _FakeResponse(200, {"embeddings": [[3.0, 4.0]]}),
+    )
+    client = EmbedClient(
+        url="https://api.cohere.ai/v1/embed",
+        api_key="cohere-key",
+        model="embed-english-v3.0",
+        provider=PROVIDERS["cohere"],
+    )
+    out = client.embed_batch(["hi"])
+    assert out == [[3.0, 4.0]]
+    call = recorder.calls[0]
+    assert call["json"] == {
+        "texts": ["hi"],
+        "model": "embed-english-v3.0",
+        "input_type": "search_document",
+    }
+    assert call["headers"] == {"Authorization": "Bearer cohere-key"}
+
+
+# ----------------------------------------------------------------------
+# Auto-detect
+# ----------------------------------------------------------------------
+
+def test_autodetect_on_401_falls_back_to_pentatonic(recorder):
+    """Operator configured openai but the URL+key actually belong to
+    Pentatonic Gateway. First call 401s, auto-detect probes pentatonic
+    and succeeds."""
+    recorder.respond(
+        "https://lambda-gateway.pentatonic.com/v1/embeddings",
+        _FakeResponse(401, '{"error":"Invalid or missing API key"}'),
+    )
+    recorder.respond(
+        "https://lambda-gateway.pentatonic.com/v1/embed",
+        _FakeResponse(200, {"data": [{"embedding": [9.0]}]}),
+    )
+    client = EmbedClient(
+        url="https://lambda-gateway.pentatonic.com/v1/embeddings",
+        api_key="k",
+        model="nv-embed-v2",
+        provider=PROVIDERS["openai"],
+    )
+    out = client.embed_batch(["x"])
+    assert out == [[9.0]]
+    assert client.active_provider == "pentatonic-gateway"
+    # First call uses configured (openai) shape, second uses pentatonic
+    assert recorder.calls[0]["headers"] == {"Authorization": "Bearer k"}
+    assert recorder.calls[1]["headers"] == {"X-API-Key": "k"}
+
+
+def test_autodetect_caches_after_first_success(recorder):
+    """Once auto-detect picks a winner, subsequent calls go straight to it
+    without retrying the original 401."""
+    recorder.respond(
+        "https://gw/v1/embeddings",
+        _FakeResponse(401, "wrong scheme"),
+    )
+    recorder.respond(
+        "https://gw/v1/embed",
+        _FakeResponse(200, {"data": [{"embedding": [1.0]}]}),
+    )
+    client = EmbedClient(
+        url="https://gw/v1/embeddings",
+        api_key="k",
+        model="m",
+        provider=PROVIDERS["openai"],
+    )
+    client.embed_batch(["a"])  # triggers detect
+    n_after_first = len(recorder.calls)
+    client.embed_batch(["b"])  # should go straight to /v1/embed
+    assert len(recorder.calls) == n_after_first + 1
+    assert recorder.calls[-1]["url"] == "https://gw/v1/embed"
+
+
+def test_autodetect_disabled_raises(recorder):
+    recorder.respond("https://gw/v1/embeddings", _FakeResponse(401, "no auth"))
+    client = EmbedClient(
+        url="https://gw/v1/embeddings",
+        api_key="k",
+        model="m",
+        provider=PROVIDERS["openai"],
+        autodetect=False,
+    )
+    with pytest.raises(EmbedAuthError):
+        client.embed_batch(["x"])
+    # Only one call: no probing happened.
+    assert len(recorder.calls) == 1
+
+
+def test_autodetect_all_fail_raises(recorder):
+    """Every candidate also 401s — raise EmbedAuthError."""
+    recorder.respond("https://gw/v1/embeddings", _FakeResponse(401, "x"))
+    recorder.respond("https://gw/v1/embed", _FakeResponse(401, "x"))
+    client = EmbedClient(
+        url="https://gw/v1/embeddings",
+        api_key="k",
+        model="m",
+        provider=PROVIDERS["openai"],
+    )
+    with pytest.raises(EmbedAuthError):
+        client.embed_batch(["x"])
+
+
+# ----------------------------------------------------------------------
+# Error handling
+# ----------------------------------------------------------------------
+
+def test_non_401_http_error_does_not_trigger_autodetect(recorder):
+    recorder.respond(
+        "https://gw/v1/embeddings",
+        _FakeResponse(503, "upstream down"),
+    )
+    client = EmbedClient(
+        url="https://gw/v1/embeddings",
+        api_key="k",
+        model="m",
+        provider=PROVIDERS["openai"],
+    )
+    with pytest.raises(EmbedHTTPError) as exc:
+        client.embed_batch(["x"])
+    assert exc.value.status == 503
+    assert len(recorder.calls) == 1
+
+
+def test_empty_input_returns_empty(recorder):
+    client = EmbedClient(
+        url="https://gw/v1/embeddings",
+        api_key="k",
+        model="m",
+        provider=PROVIDERS["openai"],
+    )
+    assert client.embed_batch([]) == []
+    assert recorder.calls == []
+
+
+# ----------------------------------------------------------------------
+# from_env construction
+# ----------------------------------------------------------------------
+
+def test_from_env_reads_layer_prefix(monkeypatch, recorder):
+    monkeypatch.setenv("L4_NV_EMBED_URL", "https://lambda-gateway.pentatonic.com/v1/embed")
+    monkeypatch.setenv("L4_EMBED_API_KEY", "real-key")
+    monkeypatch.setenv("L4_EMBED_MODEL", "nv-embed-v2")
+    monkeypatch.setenv("L4_EMBED_PROVIDER", "pentatonic-gateway")
+    recorder.respond(
+        "https://lambda-gateway.pentatonic.com/v1/embed",
+        _FakeResponse(200, {"data": [{"embedding": [42.0]}]}),
+    )
+    client = EmbedClient.from_env(prefix="L4_")
+    out = client.embed_batch(["t"])
+    assert out == [[42.0]]
+    assert client.active_provider == "pentatonic-gateway"
+    assert recorder.calls[0]["headers"] == {"X-API-Key": "real-key"}
+
+
+def test_from_env_default_provider_is_openai(monkeypatch):
+    monkeypatch.setenv("L5_NV_EMBED_URL", "https://gw/v1/embeddings")
+    monkeypatch.setenv("L5_EMBED_API_KEY", "k")
+    client = EmbedClient.from_env(prefix="L5_")
+    assert client.active_provider == "openai"
+
+
+def test_from_env_autodetect_opt_out(monkeypatch, recorder):
+    monkeypatch.setenv("L4_NV_EMBED_URL", "https://gw/v1/embeddings")
+    monkeypatch.setenv("L4_EMBED_API_KEY", "k")
+    monkeypatch.setenv("L4_EMBED_AUTODETECT", "false")
+    recorder.respond("https://gw/v1/embeddings", _FakeResponse(401, "x"))
+    client = EmbedClient.from_env(prefix="L4_")
+    with pytest.raises(EmbedAuthError):
+        client.embed_batch(["x"])
+    assert len(recorder.calls) == 1
+
+
+# ----------------------------------------------------------------------
+# URL handling
+# ----------------------------------------------------------------------
+
+def test_url_without_path_gets_provider_default(recorder):
+    """If operator provides only a base URL, the provider's path_default
+    is appended."""
+    recorder.respond(
+        "https://lambda-gateway.pentatonic.com/v1/embed",
+        _FakeResponse(200, {"data": [{"embedding": [0.0]}]}),
+    )
+    client = EmbedClient(
+        url="https://lambda-gateway.pentatonic.com",
+        api_key="k",
+        model="m",
+        provider=PROVIDERS["pentatonic-gateway"],
+    )
+    client.embed_batch(["x"])
+    assert recorder.calls[0]["url"] == "https://lambda-gateway.pentatonic.com/v1/embed"