@pentatonic-ai/ai-agent-sdk 0.9.5 → 0.10.0

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

@@ -1 +0,0 @@
-"""Shared utilities used across the memory-engine layer services."""

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

@@ -1,468 +0,0 @@
-"""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 = "",
-        max_batch: int = 5,
-    ) -> 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
-        # 0 = unlimited (no chunking). Positive = max texts per upstream call;
-        # larger inputs are split into multiple calls (concurrent in async path)
-        # and the results concatenated. Defaults to 5 because that's the per-call
-        # cap observed on Pentatonic AI Gateway — above which it 502s and the
-        # caller silently loses vector writes (see test_chunking_* tests).
-        self._max_batch = max(0, max_batch)
-
-    # ------------------------------------------------------------------
-    # 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'
-          {prefix}EMBED_MAX_BATCH    default '5' (gateway-safe; '0' disables chunking)
-        """
-        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"))
-        max_batch = int(os.environ.get(f"{prefix}EMBED_MAX_BATCH", "5"))
-
-        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,
-            max_batch=max_batch,
-        )
-
-    # ------------------------------------------------------------------
-    # 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.
-
-        Splits into chunks of `max_batch` (default 5) and posts each
-        sequentially when the input exceeds the limit. Results are
-        concatenated in input order. `max_batch=0` disables chunking.
-        """
-        if not texts:
-            return []
-        if self._max_batch == 0 or len(texts) <= self._max_batch:
-            return self._post_with_autodetect(texts, async_mode=False)
-        out: list[list[float]] = []
-        for start in range(0, len(texts), self._max_batch):
-            chunk = texts[start:start + self._max_batch]
-            out.extend(self._post_with_autodetect(chunk, async_mode=False))
-        return out
-
-    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]]:
-        """Async embed. Chunks are fired concurrently via asyncio.gather
-        when the input exceeds `max_batch`; raises the first error if any
-        chunk fails (matching the un-chunked semantics)."""
-        if not texts:
-            return []
-        if self._max_batch == 0 or len(texts) <= self._max_batch:
-            return await self._post_with_autodetect_async(texts)
-        import asyncio
-        chunks = [
-            texts[start:start + self._max_batch]
-            for start in range(0, len(texts), self._max_batch)
-        ]
-        results = await asyncio.gather(
-            *(self._post_with_autodetect_async(chunk) for chunk in chunks)
-        )
-        out: list[list[float]] = []
-        for r in results:
-            out.extend(r)
-        return out
-
-    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

@@ -1,50 +0,0 @@
-FROM python:3.12-slim
-
-WORKDIR /app
-
-# System deps for sentence-transformers + neo4j driver
-RUN apt-get update && apt-get install -y --no-install-recommends \
-        gcc g++ git curl \
-    && rm -rf /var/lib/apt/lists/*
-
-# Reranker = sentence-transformers MiniLM cross-encoder.
-# Torch CPU wheels are fine — reranker is small enough to be CPU-bound.
-#
-# sqlite-vec 0.1.9: native KNN over packed-f32 vectors stored in a vec0
-# virtual table. Replaces the legacy hand-rolled Python cosine loop over
-# JSON-serialised embeddings in search_qmd_informed (~15s timeout at 450k
-# rows → ~50ms native MATCH). Pin to 0.1.9 — that's the version probed
-# against L4 QMD's wire format (struct.pack f32 + cosine distance_metric).
-RUN pip install --no-cache-dir \
-        fastapi "uvicorn[standard]" httpx requests pydantic \
-        neo4j \
-        sqlite-vec==0.1.9 \
-        "sentence-transformers" \
-        "torch" --extra-index-url https://download.pytorch.org/whl/cpu
-
-# 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 && \
-    python init_databases.py /data
-
-# L2 proxy expects various data dirs — point them at /data
-ENV PME_WORKSPACE=/data/workspace
-ENV PME_QMD_DB=/data/qmd.sqlite
-ENV PME_MEMORY_DB=/data/workspace.db
-ENV HF_HOME=/data/.cache/huggingface
-
-# OSS-friendly defaults: HyDE off (per request), reranker on
-ENV PME_HYDE_ENABLED=false
-ENV PME_RERANK_ENABLED=true
-ENV PME_NV_EMBED_ENABLED=true
-
-EXPOSE 8031
-
-# Pre-download the MiniLM cross-encoder so cold-start is fast
-RUN python -c "from sentence_transformers import CrossEncoder; CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2', cache_folder='/data/.cache/huggingface')"
-
-CMD ["python", "server.py", "--port", "8031", "--host", "0.0.0.0"]

package/packages/memory-engine/engine/services/l2/init_databases.py

@@ -1,81 +0,0 @@
-"""Initialise empty L0 BM25 SQLite + L4 QMD SQLite at container startup.
-
-The L2 proxy reads from two SQLite files that don't otherwise exist
-in a fresh container — so opening them via sqlite3.connect() fails
-with "unable to open database file." Pre-creating them with the
-schemas the proxy expects lets the proxy come up cleanly and return
-empty results from those layers (instead of crashing) until the shim
-populates them through /index-batch.
-
-Schemas mirror what the production ingester writes.
-"""
-
-import os
-import sqlite3
-import sys
-from pathlib import Path
-
-
-def init_l0_bm25(db_path: Path) -> None:
-    """L0 — workspace BM25 over chunked markdown."""
-    db_path.parent.mkdir(parents=True, exist_ok=True)
-    conn = sqlite3.connect(str(db_path))
-    conn.execute("PRAGMA journal_mode=WAL")
-    conn.execute("""
-        CREATE TABLE IF NOT EXISTS chunks (
-            id INTEGER PRIMARY KEY,
-            path TEXT,
-            text TEXT,
-            file_type TEXT,
-            chunk_index INTEGER,
-            created_at TEXT
-        )
-    """)
-    conn.execute("""
-        CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5(
-            path, text, file_type,
-            content='chunks',
-            content_rowid='id'
-        )
-    """)
-    conn.execute("""
-        CREATE TRIGGER IF NOT EXISTS chunks_ai AFTER INSERT ON chunks BEGIN
-            INSERT INTO chunks_fts(rowid, path, text, file_type)
-            VALUES (new.id, new.path, new.text, new.file_type);
-        END
-    """)
-    conn.commit()
-    conn.close()
-
-
-def init_l4_qmd(db_path: Path) -> None:
-    """L4 — QMD vector store, JSONB embeddings stored alongside text."""
-    db_path.parent.mkdir(parents=True, exist_ok=True)
-    conn = sqlite3.connect(str(db_path))
-    conn.execute("PRAGMA journal_mode=WAL")
-    conn.execute("""
-        CREATE TABLE IF NOT EXISTS chunks (
-            id INTEGER PRIMARY KEY,
-            path TEXT,
-            text TEXT,
-            embedding TEXT,
-            embedding_model TEXT,
-            embedding_dim INTEGER,
-            chunk_index INTEGER,
-            created_at TEXT
-        )
-    """)
-    conn.commit()
-    conn.close()
-
-
-def main() -> int:
-    base = Path(sys.argv[1] if len(sys.argv) > 1 else "/data")
-    init_l0_bm25(base / "workspace.db")
-    init_l4_qmd(base / "qmd.sqlite")
-    print(f"initialised L0 + L4 QMD databases under {base}")
-    return 0
-
-
-if __name__ == "__main__":
-    sys.exit(main())