@pentatonic-ai/ai-agent-sdk 0.7.13 → 0.8.1

package/packages/memory-engine/engine/services/_shared/embed_provider.py

@@ -0,0 +1,431 @@
+"""Embedding provider abstraction for memory-engine layer services.
+
+Different embedding gateways disagree on three things:
+  1. Auth scheme    (Authorization: Bearer ... vs X-API-Key: ...)
+  2. URL path       (/v1/embeddings vs /v1/embed vs vendor-specific)
+  3. Request shape  (OpenAI {"input": [...]} vs Cohere {"texts": [...], "input_type": ...})
+  4. Response shape ({"data": [{"embedding": ...}]} vs {"embeddings": [...]})
+
+`EmbedProvider` captures all four for one gateway. Built-in profiles cover
+OpenAI-compatible, Pentatonic AI Gateway, and Cohere. A "custom" profile lets
+you override auth + path via env without code changes.
+
+Auto-detect: when the configured provider returns 401, the client probes the
+other built-in profiles once and caches the winner for the rest of the
+process. Disabled with `EMBED_AUTODETECT=false` per layer.
+
+Usage from a layer service:
+
+    from _shared.embed_provider import EmbedClient
+
+    embed = EmbedClient.from_env(prefix="L4_")
+    vectors = embed.embed_batch(["hello", "world"])
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from dataclasses import dataclass
+from typing import Any, Callable
+from urllib.parse import urlparse, urlunparse
+
+import httpx
+
+log = logging.getLogger("embed_provider")
+
+
+# ----------------------------------------------------------------------
+# Body builders + response parsers
+# ----------------------------------------------------------------------
+
+def _openai_body(texts: list[str], model: str) -> dict[str, Any]:
+    return {"input": texts, "model": model}
+
+
+def _openai_response(payload: dict[str, Any]) -> list[list[float]]:
+    return [d["embedding"] for d in payload["data"]]
+
+
+def _pentatonic_response(payload: dict[str, Any]) -> list[list[float]]:
+    """Pentatonic AI Gateway returns {"embeddings": [...]} on /v1/embed
+    and OpenAI-shaped {"data": [{"embedding": ...}]} on /v1/embeddings.
+    Accept either."""
+    if "data" in payload:
+        return [d["embedding"] for d in payload["data"]]
+    return payload["embeddings"]
+
+
+def _cohere_body(texts: list[str], model: str) -> dict[str, Any]:
+    return {"texts": texts, "model": model, "input_type": "search_document"}
+
+
+def _cohere_response(payload: dict[str, Any]) -> list[list[float]]:
+    return payload["embeddings"]
+
+
+_BODY_BUILDERS: dict[str, Callable[[list[str], str], dict[str, Any]]] = {
+    "openai": _openai_body,
+    "cohere": _cohere_body,
+}
+
+_RESPONSE_PARSERS: dict[str, Callable[[dict[str, Any]], list[list[float]]]] = {
+    "openai": _openai_response,
+    "pentatonic": _pentatonic_response,
+    "cohere": _cohere_response,
+}
+
+
+# ----------------------------------------------------------------------
+# Provider profiles
+# ----------------------------------------------------------------------
+
+@dataclass(frozen=True)
+class EmbedProvider:
+    name: str
+    auth_header: str
+    auth_format: str  # f-string-ish; "{key}" placeholder substituted at request time
+    path_default: str
+    body_builder: Callable[[list[str], str], dict[str, Any]]
+    response_parser: Callable[[dict[str, Any]], list[list[float]]]
+
+
+PROVIDERS: dict[str, EmbedProvider] = {
+    "openai": EmbedProvider(
+        name="openai",
+        auth_header="Authorization",
+        auth_format="Bearer {key}",
+        path_default="/v1/embeddings",
+        body_builder=_openai_body,
+        response_parser=_openai_response,
+    ),
+    "pentatonic-gateway": EmbedProvider(
+        name="pentatonic-gateway",
+        auth_header="X-API-Key",
+        auth_format="{key}",
+        path_default="/v1/embed",
+        body_builder=_openai_body,
+        response_parser=_pentatonic_response,
+    ),
+    "cohere": EmbedProvider(
+        name="cohere",
+        auth_header="Authorization",
+        auth_format="Bearer {key}",
+        path_default="/v1/embed",
+        body_builder=_cohere_body,
+        response_parser=_cohere_response,
+    ),
+}
+
+
+def _build_custom_provider(env_prefix: str) -> EmbedProvider:
+    """Build a custom provider from env vars.
+
+    Env vars (env_prefix is e.g. 'L4_'):
+        {prefix}EMBED_AUTH_HEADER     default Authorization
+        {prefix}EMBED_AUTH_FORMAT     default Bearer {key}
+        {prefix}EMBED_PATH_DEFAULT    default /v1/embeddings
+        {prefix}EMBED_BODY_SHAPE      "openai" | "cohere"  default openai
+        {prefix}EMBED_RESPONSE_SHAPE  "openai" | "pentatonic" | "cohere"  default openai
+    """
+    body_shape = os.environ.get(f"{env_prefix}EMBED_BODY_SHAPE", "openai")
+    response_shape = os.environ.get(f"{env_prefix}EMBED_RESPONSE_SHAPE", "openai")
+    return EmbedProvider(
+        name="custom",
+        auth_header=os.environ.get(f"{env_prefix}EMBED_AUTH_HEADER", "Authorization"),
+        auth_format=os.environ.get(f"{env_prefix}EMBED_AUTH_FORMAT", "Bearer {key}"),
+        path_default=os.environ.get(f"{env_prefix}EMBED_PATH_DEFAULT", "/v1/embeddings"),
+        body_builder=_BODY_BUILDERS.get(body_shape, _openai_body),
+        response_parser=_RESPONSE_PARSERS.get(response_shape, _openai_response),
+    )
+
+
+def resolve_provider(name: str, env_prefix: str = "") -> EmbedProvider:
+    """Look up a built-in provider, or build a custom one from env."""
+    if name in PROVIDERS:
+        return PROVIDERS[name]
+    if name == "custom":
+        return _build_custom_provider(env_prefix)
+    raise ValueError(
+        f"Unknown EMBED_PROVIDER {name!r}. "
+        f"Built-ins: {sorted(PROVIDERS)} or 'custom'."
+    )
+
+
+# ----------------------------------------------------------------------
+# URL helpers
+# ----------------------------------------------------------------------
+
+def _swap_path(url: str, new_path: str) -> str:
+    """Replace the path component of `url` with `new_path`. Empty path on the
+    input URL is treated as 'no path provided' and gets `new_path` appended."""
+    p = urlparse(url)
+    return urlunparse((p.scheme, p.netloc, new_path, "", p.query, p.fragment))
+
+
+def _resolved_url(url: str, provider: EmbedProvider) -> str:
+    """If `url` has no path, append the provider's default path; otherwise
+    leave as-is (operator chose a specific path)."""
+    p = urlparse(url)
+    if not p.path or p.path == "/":
+        return _swap_path(url, provider.path_default)
+    return url
+
+
+# ----------------------------------------------------------------------
+# Client
+# ----------------------------------------------------------------------
+
+class EmbedAuthError(RuntimeError):
+    """Raised when every candidate provider returns 401."""
+
+
+class EmbedHTTPError(RuntimeError):
+    """Raised on non-401 HTTP errors from the embedding endpoint."""
+
+    def __init__(self, status: int, body: str):
+        super().__init__(f"HTTP {status}: {body[:200]}")
+        self.status = status
+        self.body = body
+
+
+class EmbedClient:
+    """Sync + async embedding client with provider auto-detection.
+
+    Construct via `EmbedClient.from_env(prefix="L4_")` so each layer service
+    reads its own env-var namespace; or call the constructor directly for
+    tests.
+
+    The client is stateful: after a successful auto-detect, the winning
+    provider is cached for the rest of the process lifetime. Set
+    `autodetect=False` (or env `{prefix}EMBED_AUTODETECT=false`) to disable.
+    """
+
+    def __init__(
+        self,
+        *,
+        url: str,
+        api_key: str,
+        model: str,
+        provider: EmbedProvider,
+        autodetect: bool = True,
+        timeout: float = 120.0,
+        env_prefix: str = "",
+    ) -> None:
+        self._configured_provider = provider
+        self._provider = provider
+        self._configured_url = url
+        self._url = _resolved_url(url, provider)
+        self._api_key = api_key
+        self._model = model
+        self._timeout = timeout
+        self._autodetect = autodetect
+        self._env_prefix = env_prefix
+        self._detected = False
+
+    # ------------------------------------------------------------------
+    # Construction
+    # ------------------------------------------------------------------
+
+    @classmethod
+    def from_env(
+        cls,
+        prefix: str,
+        *,
+        url_var: str | None = None,
+        key_var: str | None = None,
+        model_var: str | None = None,
+        default_url: str = "",
+        default_model: str = "nv-embed-v2",
+    ) -> "EmbedClient":
+        """Build an EmbedClient from layer-prefixed env vars.
+
+        Layer services use `EmbedClient.from_env(prefix="L4_")` and the
+        client reads:
+          {prefix}NV_EMBED_URL  (override with `url_var=...`)
+          {prefix}EMBED_API_KEY
+          {prefix}EMBED_MODEL
+          {prefix}EMBED_PROVIDER     default 'openai'
+          {prefix}EMBED_AUTODETECT   default 'true'
+          {prefix}EMBED_TIMEOUT      default '120'
+        """
+        url_var = url_var or f"{prefix}NV_EMBED_URL"
+        key_var = key_var or f"{prefix}EMBED_API_KEY"
+        model_var = model_var or f"{prefix}EMBED_MODEL"
+
+        url = os.environ.get(url_var, default_url)
+        api_key = os.environ.get(key_var, "")
+        model = os.environ.get(model_var, default_model)
+        provider_name = os.environ.get(f"{prefix}EMBED_PROVIDER", "openai")
+        autodetect = os.environ.get(f"{prefix}EMBED_AUTODETECT", "true").lower() == "true"
+        timeout = float(os.environ.get(f"{prefix}EMBED_TIMEOUT", "120"))
+
+        provider = resolve_provider(provider_name, env_prefix=prefix)
+        return cls(
+            url=url,
+            api_key=api_key,
+            model=model,
+            provider=provider,
+            autodetect=autodetect,
+            timeout=timeout,
+            env_prefix=prefix,
+        )
+
+    # ------------------------------------------------------------------
+    # Internals
+    # ------------------------------------------------------------------
+
+    def _headers(self, provider: EmbedProvider) -> dict[str, str]:
+        if not self._api_key:
+            return {}
+        return {provider.auth_header: provider.auth_format.format(key=self._api_key)}
+
+    def _candidate_url(self, provider: EmbedProvider) -> str:
+        """URL to try for this provider. The configured provider keeps the
+        operator's chosen URL; auto-detect candidates swap in their own
+        path_default since the operator's path was wrong for them."""
+        if provider.name == self._configured_provider.name:
+            return self._url
+        return _swap_path(self._configured_url, provider.path_default)
+
+    def _switch_to(self, provider: EmbedProvider) -> None:
+        """Make `provider` the active provider for future requests."""
+        log.warning(
+            "embed_provider auto-detect switched: configured=%s -> active=%s. "
+            "Set %sEMBED_PROVIDER=%s to silence this.",
+            self._configured_provider.name,
+            provider.name,
+            self._env_prefix,
+            provider.name,
+        )
+        self._provider = provider
+        self._url = self._candidate_url(provider)
+        self._detected = True
+
+    # ------------------------------------------------------------------
+    # Sync API
+    # ------------------------------------------------------------------
+
+    def embed_batch(self, texts: list[str]) -> list[list[float]]:
+        """Embed a list of texts. Empty list returns empty list."""
+        if not texts:
+            return []
+        return self._post_with_autodetect(texts, async_mode=False)
+
+    def embed_one(self, text: str) -> list[float]:
+        return self.embed_batch([text])[0]
+
+    # ------------------------------------------------------------------
+    # Async API
+    # ------------------------------------------------------------------
+
+    async def embed_batch_async(self, texts: list[str]) -> list[list[float]]:
+        if not texts:
+            return []
+        return await self._post_with_autodetect_async(texts)
+
+    async def embed_one_async(self, text: str) -> list[float]:
+        out = await self.embed_batch_async([text])
+        return out[0]
+
+    # ------------------------------------------------------------------
+    # Request paths
+    # ------------------------------------------------------------------
+
+    def _post_with_autodetect(self, texts: list[str], *, async_mode: bool) -> list[list[float]]:
+        del async_mode  # kept for symmetry; sync path is its own method
+        body = self._provider.body_builder(texts, self._model)
+        headers = self._headers(self._provider)
+        try:
+            r = httpx.post(self._url, json=body, headers=headers, timeout=self._timeout)
+        except httpx.HTTPError as exc:
+            raise EmbedHTTPError(0, str(exc)) from exc
+
+        if r.status_code == 401 and self._autodetect and not self._detected:
+            return self._autodetect_and_retry(texts, last_body=r.text)
+
+        if r.status_code == 401:
+            raise EmbedAuthError(r.text)
+        if not r.is_success:
+            raise EmbedHTTPError(r.status_code, r.text)
+        return self._provider.response_parser(r.json())
+
+    async def _post_with_autodetect_async(self, texts: list[str]) -> list[list[float]]:
+        body = self._provider.body_builder(texts, self._model)
+        headers = self._headers(self._provider)
+        try:
+            async with httpx.AsyncClient(timeout=self._timeout) as client:
+                r = await client.post(self._url, json=body, headers=headers)
+        except httpx.HTTPError as exc:
+            raise EmbedHTTPError(0, str(exc)) from exc
+
+        if r.status_code == 401 and self._autodetect and not self._detected:
+            return await self._autodetect_and_retry_async(texts, last_body=r.text)
+
+        if r.status_code == 401:
+            raise EmbedAuthError(r.text)
+        if not r.is_success:
+            raise EmbedHTTPError(r.status_code, r.text)
+        return self._provider.response_parser(r.json())
+
+    # ------------------------------------------------------------------
+    # Auto-detect
+    # ------------------------------------------------------------------
+
+    def _candidates(self) -> list[EmbedProvider]:
+        return [p for p in PROVIDERS.values() if p.name != self._provider.name]
+
+    def _autodetect_and_retry(self, texts: list[str], *, last_body: str) -> list[list[float]]:
+        for candidate in self._candidates():
+            url = self._candidate_url(candidate)
+            body = candidate.body_builder(texts, self._model)
+            headers = (
+                {candidate.auth_header: candidate.auth_format.format(key=self._api_key)}
+                if self._api_key
+                else {}
+            )
+            try:
+                r = httpx.post(url, json=body, headers=headers, timeout=self._timeout)
+            except httpx.HTTPError:
+                continue
+            if r.is_success:
+                self._switch_to(candidate)
+                return candidate.response_parser(r.json())
+        raise EmbedAuthError(
+            f"all providers returned auth or transport errors. "
+            f"last 401 body: {last_body[:200]}"
+        )
+
+    async def _autodetect_and_retry_async(self, texts: list[str], *, last_body: str) -> list[list[float]]:
+        for candidate in self._candidates():
+            url = self._candidate_url(candidate)
+            body = candidate.body_builder(texts, self._model)
+            headers = (
+                {candidate.auth_header: candidate.auth_format.format(key=self._api_key)}
+                if self._api_key
+                else {}
+            )
+            try:
+                async with httpx.AsyncClient(timeout=self._timeout) as client:
+                    r = await client.post(url, json=body, headers=headers)
+            except httpx.HTTPError:
+                continue
+            if r.is_success:
+                self._switch_to(candidate)
+                return candidate.response_parser(r.json())
+        raise EmbedAuthError(
+            f"all providers returned auth or transport errors. "
+            f"last 401 body: {last_body[:200]}"
+        )
+
+    # ------------------------------------------------------------------
+    # Introspection (handy for /health and tests)
+    # ------------------------------------------------------------------
+
+    @property
+    def active_provider(self) -> str:
+        return self._provider.name
+
+    @property
+    def active_url(self) -> str:
+        return self._url

package/packages/memory-engine/engine/services/l2/Dockerfile

@@ -15,8 +15,10 @@ RUN pip install --no-cache-dir \
         "sentence-transformers" \
         "torch" --extra-index-url https://download.pytorch.org/whl/cpu
 
-COPY l2-hybridrag-proxy.py /app/server.py
-COPY init_databases.py /app/init_databases.py
+# Shared embed_provider module (build context is engine/services).
+COPY _shared /app/_shared
+COPY l2/l2-hybridrag-proxy.py /app/server.py
+COPY l2/init_databases.py /app/init_databases.py
 
 # Pre-create empty L0 BM25 SQLite + empty workspace
 RUN mkdir -p /data/workspace /data/.pentatonic/memory /data/.cache && \