vera-embed-openai 0.3.1__tar.gz

vera_embed_openai-0.3.1/.gitignore

@@ -0,0 +1,16 @@
+.venv/
+__pycache__/
+.pytest_cache/
+*.pyc
+*.vera
+dist/
+dist-electron/
+release/
+build/
+node_modules/
+*.egg-info/
+examples/*.pdf
+site/
+.env
+.env.local
+.env.*.local

vera_embed_openai-0.3.1/PKG-INFO

@@ -0,0 +1,77 @@
+Metadata-Version: 2.5
+Name: vera-embed-openai
+Version: 0.3.1
+Summary: Official OpenAI embeddings provider for VERA
+Project-URL: Homepage, https://github.com/dkylewillis/vera
+Project-URL: Repository, https://github.com/dkylewillis/vera
+Project-URL: Documentation, https://dkylewillis.github.io/vera/packages/vera-embed-openai/
+Author: Kyle Willis
+License: Apache-2.0
+Keywords: embeddings,openai,semantic-search,vera
+Classifier: Development Status :: 3 - Alpha
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Text Processing :: Indexing
+Requires-Python: >=3.10
+Requires-Dist: numpy>=1.24
+Requires-Dist: vera-doc>=0.3.1
+Description-Content-Type: text/markdown
+
+# vera-embed-openai
+
+Official OpenAI embeddings plugin for VERA. Registers the `openai` provider
+under the `vera.embedders` entry-point group.
+
+`vera-cli` and `vera-app` depend on this package so hosted OpenAI conversion
+works out of the box. The client uses stdlib `urllib` — there is no `openai`
+SDK dependency.
+
+## Install
+
+```bash
+python -m pip install "vera-embed-openai>=0.3.0"
+```
+
+From a repository checkout with uv, the workspace installs it by default
+(via `vera-cli` / `vera-app`):
+
+```bash
+uv sync
+```
+
+## Usage
+
+```bash
+set OPENAI_API_KEY=...
+vera convert "manual.pdf" --model openai:text-embedding-3-small
+```
+
+```python
+import os
+
+from vera_ingest import convert
+
+os.environ["OPENAI_API_KEY"] = "..."
+convert("manual.pdf", "manual.vera", model="openai:text-embedding-3-small")
+```
+
+Keep the API key in `OPENAI_API_KEY`. Optional `OPENAI_BASE_URL` (default
+`https://api.openai.com/v1`) points at Azure, OpenRouter, or a local
+OpenAI-compatible server. Archives still record `openai:<model-id>`, so two
+endpoints that embed different models under the same id are indistinguishable
+at search time.
+
+## Notes
+
+- Known dimensions (no network in the constructor): `text-embedding-3-small`
+  1536, `text-embedding-3-large` 3072, `text-embedding-ada-002` 1536.
+  Unrecognized model ids probe once on first use.
+- Convert-time options: `batch_size` (1–2048) and `timeout` seconds.
+  Search resolves `get_embedder(stored_model_name)` with defaults.
+- Semantic search of a hosted archive needs the same provider and credentials
+  on the searching machine. Keyword search still works without a key.
+- Desktop Convert Cancel does not interrupt an in-flight embeddings HTTP
+  batch; conversion checks cancellation after `embed()` returns.
+
+See the [vera-embed-openai documentation](https://dkylewillis.github.io/vera/packages/vera-embed-openai/)
+and [conversion guide](https://github.com/dkylewillis/vera/blob/main/docs/conversion.md).

vera_embed_openai-0.3.1/README.md

@@ -0,0 +1,58 @@
+# vera-embed-openai
+
+Official OpenAI embeddings plugin for VERA. Registers the `openai` provider
+under the `vera.embedders` entry-point group.
+
+`vera-cli` and `vera-app` depend on this package so hosted OpenAI conversion
+works out of the box. The client uses stdlib `urllib` — there is no `openai`
+SDK dependency.
+
+## Install
+
+```bash
+python -m pip install "vera-embed-openai>=0.3.0"
+```
+
+From a repository checkout with uv, the workspace installs it by default
+(via `vera-cli` / `vera-app`):
+
+```bash
+uv sync
+```
+
+## Usage
+
+```bash
+set OPENAI_API_KEY=...
+vera convert "manual.pdf" --model openai:text-embedding-3-small
+```
+
+```python
+import os
+
+from vera_ingest import convert
+
+os.environ["OPENAI_API_KEY"] = "..."
+convert("manual.pdf", "manual.vera", model="openai:text-embedding-3-small")
+```
+
+Keep the API key in `OPENAI_API_KEY`. Optional `OPENAI_BASE_URL` (default
+`https://api.openai.com/v1`) points at Azure, OpenRouter, or a local
+OpenAI-compatible server. Archives still record `openai:<model-id>`, so two
+endpoints that embed different models under the same id are indistinguishable
+at search time.
+
+## Notes
+
+- Known dimensions (no network in the constructor): `text-embedding-3-small`
+  1536, `text-embedding-3-large` 3072, `text-embedding-ada-002` 1536.
+  Unrecognized model ids probe once on first use.
+- Convert-time options: `batch_size` (1–2048) and `timeout` seconds.
+  Search resolves `get_embedder(stored_model_name)` with defaults.
+- Semantic search of a hosted archive needs the same provider and credentials
+  on the searching machine. Keyword search still works without a key.
+- Desktop Convert Cancel does not interrupt an in-flight embeddings HTTP
+  batch; conversion checks cancellation after `embed()` returns.
+
+See the [vera-embed-openai documentation](https://dkylewillis.github.io/vera/packages/vera-embed-openai/)
+and [conversion guide](https://github.com/dkylewillis/vera/blob/main/docs/conversion.md).

vera_embed_openai-0.3.1/pyproject.toml

@@ -0,0 +1,40 @@
+[project]
+name = "vera-embed-openai"
+version = "0.3.1"
+description = "Official OpenAI embeddings provider for VERA"
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+  "vera-doc>=0.3.1",
+  "numpy>=1.24",
+]
+authors = [{name = "Kyle Willis"}]
+license = {text = "Apache-2.0"}
+keywords = ["openai", "embeddings", "semantic-search", "vera"]
+classifiers = [
+  "Development Status :: 3 - Alpha",
+  "License :: OSI Approved :: Apache Software License",
+  "Programming Language :: Python :: 3",
+  "Topic :: Text Processing :: Indexing",
+]
+
+[project.urls]
+Homepage = "https://github.com/dkylewillis/vera"
+Repository = "https://github.com/dkylewillis/vera"
+Documentation = "https://dkylewillis.github.io/vera/packages/vera-embed-openai/"
+
+[project.entry-points."vera.embedders"]
+openai = "vera_embed_openai:create_embedder"
+
+[project.entry-points."vera.embedder_descriptors"]
+openai = "vera_embed_openai:create_descriptor"
+
+[project.entry-points."vera.embedder_models"]
+openai = "vera_embed_openai:list_models"
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/vera_embed_openai"]

vera_embed_openai-0.3.1/src/vera_embed_openai/__init__.py

@@ -0,0 +1,81 @@
+"""Official OpenAI embeddings provider for VERA."""
+
+from __future__ import annotations
+
+from vera_doc import (
+    EmbedderDescriptor,
+    EmbeddingModelInfo,
+    register_embedder,
+    register_embedder_descriptor,
+    register_embedder_models,
+)
+
+from .options import (
+    BASE_URL_ENV,
+    CREDENTIAL_ENV,
+    DEFAULT_BASE_URL,
+    DEFAULT_MODEL_ID,
+    MODEL_DIMENSIONS,
+    PROVIDER,
+    OpenAIOptions,
+    describe_provider,
+)
+from .options import list_models as _list_models
+from .provider import (
+    MAX_INPUT_TOKENS,
+    MAX_REQUEST_TOKENS,
+    OpenAIEmbedder,
+    OpenAIEmbedderError,
+    create_embedder,
+    embeddings_url,
+    estimate_tokens,
+    iter_embed_batches,
+)
+
+__all__ = [
+    "BASE_URL_ENV",
+    "CREDENTIAL_ENV",
+    "DEFAULT_BASE_URL",
+    "DEFAULT_MODEL_ID",
+    "MAX_INPUT_TOKENS",
+    "MAX_REQUEST_TOKENS",
+    "MODEL_DIMENSIONS",
+    "OpenAIEmbedder",
+    "OpenAIEmbedderError",
+    "OpenAIOptions",
+    "PROVIDER",
+    "create_descriptor",
+    "create_embedder",
+    "describe_provider",
+    "embeddings_url",
+    "ensure_registered",
+    "estimate_tokens",
+    "iter_embed_batches",
+    "list_models",
+]
+
+
+def create_descriptor() -> EmbedderDescriptor:
+    """Entry-point factory for ``vera.embedder_descriptors``."""
+    return describe_provider()
+
+
+def list_models() -> tuple[EmbeddingModelInfo, ...]:
+    """Entry-point factory for ``vera.embedder_models``."""
+    return _list_models()
+
+
+def ensure_registered(*, replace: bool = True) -> None:
+    """Register the ``openai`` embedder without relying on package metadata.
+
+    Entry-point discovery fails in PyInstaller freezes and PYTHONPATH-only
+    source runs that never install ``vera-embed-openai`` dist-info. Callers
+    that already import this package (CLI, the desktop sidecar) should invoke
+    this so Convert and search still resolve the provider.
+    """
+    register_embedder(PROVIDER, create_embedder, replace=replace)
+    register_embedder_descriptor(PROVIDER, create_descriptor, replace=replace)
+    register_embedder_models(PROVIDER, list_models, replace=replace)
+
+
+ensure_registered()

vera_embed_openai-0.3.1/src/vera_embed_openai/options.py

@@ -0,0 +1,124 @@
+"""Typed options, descriptor, and model list for the OpenAI embedder."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from vera_doc import (
+    EmbedderCapabilities,
+    EmbedderDescriptor,
+    EmbedderOptions,
+    EmbeddingModelInfo,
+)
+from vera_doc.embedder_descriptors import fields_from_dataclass
+
+PROVIDER = "openai"
+DEFAULT_MODEL_ID = "text-embedding-3-small"
+CREDENTIAL_ENV = "OPENAI_API_KEY"
+BASE_URL_ENV = "OPENAI_BASE_URL"
+DEFAULT_BASE_URL = "https://api.openai.com/v1"
+
+MODEL_DIMENSIONS: dict[str, int] = {
+    "text-embedding-3-small": 1536,
+    "text-embedding-3-large": 3072,
+    "text-embedding-ada-002": 1536,
+}
+
+
+@dataclass(frozen=True)
+class OpenAIOptions(EmbedderOptions):
+    """OpenAI-owned convert-time embedding settings.
+
+    Credentials and the API root stay in the environment
+    (``OPENAI_API_KEY``, optional ``OPENAI_BASE_URL``), not in this dataclass.
+    """
+
+    batch_size: int = field(
+        default=128,
+        metadata={
+            "label": "Batch size",
+            "description": (
+                "Maximum texts per OpenAI request. Requests also split when "
+                "the estimated token budget would be exceeded. Convert-time "
+                "only — search uses the default."
+            ),
+            "minimum": 1,
+            "maximum": 2048,
+            "step": 1,
+            "scope": "convert",
+        },
+    )
+    timeout: int = field(
+        default=60,
+        metadata={
+            "label": "Timeout",
+            "description": (
+                "Seconds to wait for each embeddings HTTP response. "
+                "Convert-time only — search uses the default."
+            ),
+            "unit": "seconds",
+            "minimum": 1,
+            "maximum": 600,
+            "step": 1,
+            "scope": "convert",
+        },
+    )
+
+
+def describe_provider() -> EmbedderDescriptor:
+    return EmbedderDescriptor(
+        provider=PROVIDER,
+        label="openai — hosted embeddings",
+        description=(
+            "OpenAI embeddings API (text-embedding-3-* and similar). "
+            "Requires OPENAI_API_KEY. Archives converted with this provider "
+            "need the same key for later semantic or hybrid search; keyword "
+            "search still works offline. Conversion bills per request."
+        ),
+        default_model_id=DEFAULT_MODEL_ID,
+        example_specs=(
+            "openai:text-embedding-3-small",
+            "openai:text-embedding-3-large",
+        ),
+        capabilities=EmbedderCapabilities(
+            requires_network=True,
+            requires_api_key=True,
+            credential_env=CREDENTIAL_ENV,
+            local_model=False,
+            configurable_dimension=False,
+            supports_model_listing=True,
+        ),
+        fields=fields_from_dataclass(OpenAIOptions),
+        notes=(
+            "Set OPENAI_API_KEY (desktop: File > Settings → Embeddings). "
+            "Optional OPENAI_BASE_URL overrides the API root "
+            f"(default {DEFAULT_BASE_URL}); archives still record "
+            "openai:<model-id>, so a custom endpoint that serves a different "
+            "model is not detectable at search time. "
+            "batch_size and timeout are convert-time options; search uses defaults. "
+            "Cancel does not interrupt an in-flight embeddings request.",
+        ),
+    )
+
+
+def list_models() -> tuple[EmbeddingModelInfo, ...]:
+    return (
+        EmbeddingModelInfo(
+            model_id="text-embedding-3-small",
+            label="text-embedding-3-small",
+            spec="openai:text-embedding-3-small",
+            description="OpenAI text-embedding-3-small (1536-d).",
+        ),
+        EmbeddingModelInfo(
+            model_id="text-embedding-3-large",
+            label="text-embedding-3-large",
+            spec="openai:text-embedding-3-large",
+            description="OpenAI text-embedding-3-large (3072-d).",
+        ),
+        EmbeddingModelInfo(
+            model_id="text-embedding-ada-002",
+            label="text-embedding-ada-002",
+            spec="openai:text-embedding-ada-002",
+            description="Legacy OpenAI ada embedding model (1536-d).",
+        ),
+    )

vera_embed_openai-0.3.1/src/vera_embed_openai/provider.py

@@ -0,0 +1,290 @@
+"""Stdlib HTTP OpenAI embeddings client.
+
+Batch splitting, retries, and URL construction stay in this module. The
+``EmbeddingFunction`` contract is ``embed(texts)``, so convert and search
+never see request boundaries.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+import threading
+import time
+import urllib.error
+import urllib.request
+from collections.abc import Iterator, Sequence
+from typing import Any
+
+import numpy as np
+
+from .options import (
+    BASE_URL_ENV,
+    CREDENTIAL_ENV,
+    DEFAULT_BASE_URL,
+    DEFAULT_MODEL_ID,
+    MODEL_DIMENSIONS,
+    PROVIDER,
+    OpenAIOptions,
+)
+
+# OpenAI embeddings caps (re-check against current docs if these start failing).
+# Per-input limit for text-embedding-3-* and ada-002.
+MAX_INPUT_TOKENS = 8192
+# Published per-request cap is ~300k tokens; stay under it to absorb the
+# character heuristic's estimation error.
+MAX_REQUEST_TOKENS = 250_000
+MAX_BATCH_ITEMS = 2048
+MAX_RETRIES = 4
+_RETRY_STATUS = frozenset({429, 500, 502, 503, 504})
+_SURROGATE_RE = re.compile(r"[\ud800-\udfff]")
+
+
+class OpenAIEmbedderError(RuntimeError):
+    """Raised when the OpenAI embeddings API rejects or cannot complete a request."""
+
+
+def estimate_tokens(text: str) -> int:
+    """Conservative character heuristic (~3 chars/token). Over-estimates on purpose."""
+    if not text:
+        return 1
+    return max(1, (len(text) + 2) // 3)
+
+
+def iter_embed_batches(
+    texts: Sequence[str],
+    *,
+    batch_size: int,
+    max_request_tokens: int = MAX_REQUEST_TOKENS,
+    max_input_tokens: int = MAX_INPUT_TOKENS,
+) -> Iterator[list[str]]:
+    """Yield text batches that fit both the item cap and the token budget.
+
+    A single text over ``max_input_tokens`` raises rather than truncating.
+    """
+    max_items = max(1, min(int(batch_size), MAX_BATCH_ITEMS))
+    budget = max(1, int(max_request_tokens))
+    batch: list[str] = []
+    batch_tokens = 0
+    for index, text in enumerate(texts):
+        value = text if isinstance(text, str) else ("" if text is None else str(text))
+        tokens = estimate_tokens(value)
+        if tokens > max_input_tokens:
+            raise OpenAIEmbedderError(
+                f"chunk {index} is ~{tokens} tokens, over the {max_input_tokens} "
+                "per-input OpenAI embeddings limit; lower the pipeline chunk_size"
+            )
+        if batch and (len(batch) >= max_items or batch_tokens + tokens > budget):
+            yield batch
+            batch = []
+            batch_tokens = 0
+        batch.append(value)
+        batch_tokens += tokens
+    if batch:
+        yield batch
+
+
+def embeddings_url(base_url: str) -> str:
+    """Join ``base_url`` with ``/embeddings``, inserting ``/v1`` when missing."""
+    root = (base_url or DEFAULT_BASE_URL).strip() or DEFAULT_BASE_URL
+    root = root.rstrip("/")
+    if root.endswith("/embeddings"):
+        return root
+    if not root.endswith("/v1") and "/v1/" not in root + "/":
+        root = f"{root}/v1"
+    return f"{root}/embeddings"
+
+
+def _sanitize_unicode(value: Any) -> Any:
+    if isinstance(value, str):
+        return _SURROGATE_RE.sub("\ufffd", value)
+    if isinstance(value, list):
+        return [_sanitize_unicode(item) for item in value]
+    if isinstance(value, dict):
+        return {
+            _sanitize_unicode(key) if isinstance(key, str) else key: _sanitize_unicode(item)
+            for key, item in value.items()
+        }
+    return value
+
+
+def _l2_normalize(vector: np.ndarray) -> np.ndarray:
+    array = np.asarray(vector, dtype=np.float32)
+    norm = float(np.linalg.norm(array))
+    if norm:
+        array = array / norm
+    return array.astype(np.float32)
+
+
+def _retry_delay_seconds(headers: Any, attempt: int) -> float:
+    raw = ""
+    if headers is not None:
+        raw = str(headers.get("Retry-After") or headers.get("retry-after") or "")
+    if raw:
+        try:
+            return min(60.0, max(0.5, float(raw.strip())))
+        except ValueError:
+            pass
+    return min(32.0, 0.5 * (2**attempt))
+
+
+def _error_detail(exc: urllib.error.HTTPError) -> str:
+    try:
+        body = exc.read().decode("utf-8", errors="replace")
+    except Exception:  # noqa: BLE001 - error path must not hide the status
+        body = ""
+    snippet = " ".join(body.split())
+    if len(snippet) > 500:
+        snippet = snippet[:497] + "..."
+    return snippet
+
+
+def _post_embeddings(
+    *,
+    url: str,
+    api_key: str,
+    model: str,
+    inputs: list[str],
+    timeout: float,
+) -> list[np.ndarray]:
+    payload = json.dumps(
+        _sanitize_unicode({"model": model, "input": inputs}),
+        ensure_ascii=False,
+    ).encode("utf-8")
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {api_key}",
+    }
+    last_error: Exception | None = None
+    for attempt in range(MAX_RETRIES):
+        request = urllib.request.Request(url, data=payload, headers=headers, method="POST")
+        try:
+            with urllib.request.urlopen(request, timeout=timeout) as response:
+                raw = json.loads(response.read().decode("utf-8"))
+            return _vectors_from_payload(raw, expected=len(inputs))
+        except urllib.error.HTTPError as exc:
+            detail = _error_detail(exc)
+            if exc.code in _RETRY_STATUS and attempt + 1 < MAX_RETRIES:
+                time.sleep(_retry_delay_seconds(exc.headers, attempt))
+                last_error = OpenAIEmbedderError(
+                    f"OpenAI embeddings HTTP {exc.code}: {detail or exc.reason}"
+                )
+                continue
+            raise OpenAIEmbedderError(
+                f"OpenAI embeddings HTTP {exc.code}: {detail or exc.reason}"
+            ) from exc
+        except urllib.error.URLError as exc:
+            reason = getattr(exc, "reason", exc)
+            last_error = OpenAIEmbedderError(f"Unable to reach OpenAI embeddings API: {reason}")
+            if attempt + 1 < MAX_RETRIES:
+                time.sleep(_retry_delay_seconds(None, attempt))
+                continue
+            raise last_error from exc
+        except (json.JSONDecodeError, TypeError, ValueError, KeyError) as exc:
+            raise OpenAIEmbedderError(
+                f"OpenAI embeddings returned an invalid payload: {exc}"
+            ) from exc
+    raise last_error or OpenAIEmbedderError("OpenAI embeddings request failed")
+
+
+def _vectors_from_payload(payload: Any, *, expected: int) -> list[np.ndarray]:
+    if not isinstance(payload, dict):
+        raise OpenAIEmbedderError("OpenAI embeddings response is not a JSON object")
+    rows = payload.get("data")
+    if not isinstance(rows, list) or len(rows) != expected:
+        raise OpenAIEmbedderError(
+            f"OpenAI embeddings returned {0 if not isinstance(rows, list) else len(rows)} "
+            f"vectors; expected {expected}"
+        )
+    indexed: list[tuple[int, np.ndarray]] = []
+    for position, row in enumerate(rows):
+        if not isinstance(row, dict):
+            raise OpenAIEmbedderError("OpenAI embeddings data entry is not an object")
+        embedding = row.get("embedding")
+        if not isinstance(embedding, list) or not embedding:
+            raise OpenAIEmbedderError("OpenAI embeddings data entry is missing embedding")
+        index = row.get("index", position)
+        try:
+            order = int(index)
+        except (TypeError, ValueError) as exc:
+            raise OpenAIEmbedderError("OpenAI embeddings data entry has an invalid index") from exc
+        indexed.append((order, _l2_normalize(np.asarray(embedding, dtype=np.float32))))
+    indexed.sort(key=lambda item: item[0])
+    return [vector for _index, vector in indexed]
+
+
+class OpenAIEmbedder:
+    """Hosted OpenAI embeddings. Constructor does not touch the network."""
+
+    normalization = "l2"
+
+    def __init__(
+        self,
+        model_id: str,
+        *,
+        api_key: str,
+        base_url: str,
+        batch_size: int,
+        timeout: int,
+    ) -> None:
+        model = (model_id or "").strip() or DEFAULT_MODEL_ID
+        self.model_name = f"{PROVIDER}:{model}"
+        self._model = model
+        self._api_key = api_key
+        self._url = embeddings_url(base_url)
+        self._batch_size = max(1, min(int(batch_size), MAX_BATCH_ITEMS))
+        self._timeout = float(timeout)
+        self._lock = threading.Lock()
+        self._dimension = MODEL_DIMENSIONS.get(model)
+
+    @property
+    def dimension(self) -> int:
+        if self._dimension is None:
+            self._dimension = int(len(self._embed_one_batch(["ping"])[0]))
+        return self._dimension
+
+    def embed(self, texts: list[str]) -> list[np.ndarray]:
+        if not texts:
+            return []
+        vectors: list[np.ndarray] = []
+        for batch in iter_embed_batches(texts, batch_size=self._batch_size):
+            vectors.extend(self._embed_one_batch(batch))
+        if self._dimension is None and vectors:
+            self._dimension = int(vectors[0].shape[0])
+        return vectors
+
+    def _embed_one_batch(self, batch: list[str]) -> list[np.ndarray]:
+        with self._lock:
+            return _post_embeddings(
+                url=self._url,
+                api_key=self._api_key,
+                model=self._model,
+                inputs=batch,
+                timeout=self._timeout,
+            )
+
+
+def _require_api_key() -> str:
+    api_key = os.environ.get(CREDENTIAL_ENV, "").strip()
+    if not api_key:
+        raise OpenAIEmbedderError(
+            f"Set the {CREDENTIAL_ENV} environment variable before converting or searching "
+            "with an OpenAI embedding model. In the desktop app, save it under "
+            "File > Settings → Embeddings."
+        )
+    return api_key
+
+
+def create_embedder(model_id: str, **config: Any) -> OpenAIEmbedder:
+    """Entry-point factory for ``vera.embedders`` provider ``openai``."""
+    options = OpenAIOptions.from_mapping(config)
+    model = (model_id or "").strip() or DEFAULT_MODEL_ID
+    base_url = os.environ.get(BASE_URL_ENV, "").strip() or DEFAULT_BASE_URL
+    return OpenAIEmbedder(
+        model,
+        api_key=_require_api_key(),
+        base_url=base_url,
+        batch_size=options.batch_size,
+        timeout=options.timeout,
+    )

vera_embed_openai-0.3.1/tests/test_openai_embedder.py

@@ -0,0 +1,275 @@
+"""Unit tests for the official OpenAI embeddings plugin."""
+
+from __future__ import annotations
+
+import io
+import json
+from unittest.mock import patch
+
+import numpy as np
+import pytest
+
+from vera_doc import (
+    clear_embedder_cache,
+    describe_embedder,
+    get_embedder,
+    list_embedding_models,
+    preflight_embedder,
+)
+from vera_embed_openai import (
+    CREDENTIAL_ENV,
+    DEFAULT_MODEL_ID,
+    MAX_INPUT_TOKENS,
+    OpenAIEmbedder,
+    OpenAIEmbedderError,
+    OpenAIOptions,
+    create_embedder,
+    embeddings_url,
+    ensure_registered,
+    estimate_tokens,
+    iter_embed_batches,
+)
+
+
+def _response(vectors: list[list[float]], *, indexes: list[int] | None = None):
+    payload = {
+        "data": [
+            {"index": (indexes[i] if indexes else i), "embedding": vector}
+            for i, vector in enumerate(vectors)
+        ]
+    }
+    raw = json.dumps(payload).encode("utf-8")
+
+    class _Handle:
+        def __enter__(self):
+            return self
+
+        def __exit__(self, *exc):
+            return False
+
+        def read(self) -> bytes:
+            return raw
+
+    return _Handle()
+
+
+@pytest.fixture
+def openai_key(monkeypatch):
+    monkeypatch.setenv(CREDENTIAL_ENV, "sk-test")
+    return "sk-test"
+
+
+def test_estimate_tokens_is_conservative():
+    assert estimate_tokens("") == 1
+    assert estimate_tokens("abcd") == 2
+    assert estimate_tokens("a" * 30) == 10
+
+
+def test_iter_embed_batches_splits_on_item_cap_and_token_budget():
+    texts = ["aa", "bb", "cc", "dd"]
+    batches = list(iter_embed_batches(texts, batch_size=2, max_request_tokens=1000))
+    assert batches == [["aa", "bb"], ["cc", "dd"]]
+
+    # Each "aaaaaa" is 2 tokens with the heuristic; budget 3 forces size 1.
+    longish = ["aaaaaa", "aaaaaa", "aaaaaa"]
+    token_batches = list(iter_embed_batches(longish, batch_size=10, max_request_tokens=3))
+    assert token_batches == [["aaaaaa"], ["aaaaaa"], ["aaaaaa"]]
+
+
+def test_iter_embed_batches_rejects_oversized_chunk():
+    huge = "a" * (MAX_INPUT_TOKENS * 3 + 3)
+    with pytest.raises(OpenAIEmbedderError, match="chunk 0"):
+        list(iter_embed_batches([huge], batch_size=8))
+
+
+def test_embeddings_url_joins_v1_and_embeddings():
+    assert embeddings_url("https://api.openai.com/v1") == "https://api.openai.com/v1/embeddings"
+    assert embeddings_url("https://api.openai.com") == "https://api.openai.com/v1/embeddings"
+    assert (
+        embeddings_url("https://example.test/v1/embeddings") == "https://example.test/v1/embeddings"
+    )
+
+
+def test_constructor_does_not_touch_the_network(openai_key):
+    calls: list[object] = []
+
+    def boom(*args, **kwargs):
+        calls.append((args, kwargs))
+        raise AssertionError("constructor must not call urlopen")
+
+    with patch("vera_embed_openai.provider.urllib.request.urlopen", boom):
+        embedder = create_embedder("text-embedding-3-small")
+    assert embedder.model_name == "openai:text-embedding-3-small"
+    assert embedder.dimension == 1536
+    assert embedder.normalization == "l2"
+    assert calls == []
+
+
+def test_known_model_dimensions_need_no_probe(openai_key):
+    small = OpenAIEmbedder(
+        "text-embedding-3-small",
+        api_key=openai_key,
+        base_url="https://api.openai.com/v1",
+        batch_size=8,
+        timeout=5,
+    )
+    large = OpenAIEmbedder(
+        "text-embedding-3-large",
+        api_key=openai_key,
+        base_url="https://api.openai.com/v1",
+        batch_size=8,
+        timeout=5,
+    )
+    ada = OpenAIEmbedder(
+        "text-embedding-ada-002",
+        api_key=openai_key,
+        base_url="https://api.openai.com/v1",
+        batch_size=8,
+        timeout=5,
+    )
+    assert small.dimension == 1536
+    assert large.dimension == 3072
+    assert ada.dimension == 1536
+
+
+def test_embed_normalizes_and_preserves_order(openai_key):
+    embedder = create_embedder("text-embedding-3-small", batch_size=8)
+    captured: dict[str, object] = {}
+
+    def fake_urlopen(request, timeout=None):
+        captured["url"] = request.full_url
+        captured["timeout"] = timeout
+        body = json.loads(request.data.decode("utf-8"))
+        captured["payload"] = body
+        auth = request.headers.get("Authorization") or request.get_header("Authorization")
+        captured["authorization"] = auth
+        # Unnormalized, reversed indexes — client must L2-normalize and reorder.
+        return _response(
+            [[0.0, 3.0], [4.0, 0.0]],
+            indexes=[1, 0],
+        )
+
+    with patch("vera_embed_openai.provider.urllib.request.urlopen", fake_urlopen):
+        vectors = embedder.embed(["one", "two"])
+
+    assert captured["url"] == "https://api.openai.com/v1/embeddings"
+    assert captured["payload"] == {"model": "text-embedding-3-small", "input": ["one", "two"]}
+    assert str(captured["authorization"]).endswith("sk-test")
+    assert pytest.approx(float(np.linalg.norm(vectors[0])), abs=1e-5) == 1.0
+    assert pytest.approx(float(np.linalg.norm(vectors[1])), abs=1e-5) == 1.0
+    # After reorder, index 0 is [4,0] -> [1,0]; index 1 is [0,3] -> [0,1].
+    assert pytest.approx(float(vectors[0][0]), abs=1e-5) == 1.0
+    assert pytest.approx(float(vectors[1][1]), abs=1e-5) == 1.0
+
+
+def test_embed_splits_http_requests_on_batch_size(openai_key):
+    embedder = create_embedder("text-embedding-3-small", batch_size=2)
+    payloads: list[list[str]] = []
+
+    def fake_urlopen(request, timeout=None):
+        body = json.loads(request.data.decode("utf-8"))
+        payloads.append(list(body["input"]))
+        unit = [[1.0, 0.0] for _ in body["input"]]
+        return _response(unit)
+
+    with patch("vera_embed_openai.provider.urllib.request.urlopen", fake_urlopen):
+        vectors = embedder.embed(["a", "b", "c"])
+    assert payloads == [["a", "b"], ["c"]]
+    assert len(vectors) == 3
+
+
+def test_retries_on_429_then_succeeds(openai_key, monkeypatch):
+    embedder = create_embedder("text-embedding-3-small")
+    attempts = {"n": 0}
+
+    def fake_urlopen(request, timeout=None):
+        attempts["n"] += 1
+        if attempts["n"] == 1:
+            import vera_embed_openai.provider as provider_mod
+
+            raise provider_mod.urllib.error.HTTPError(
+                request.full_url,
+                429,
+                "rate limited",
+                hdrs={"Retry-After": "0"},
+                fp=io.BytesIO(b'{"error":"slow down"}'),
+            )
+        return _response([[1.0, 0.0]])
+
+    monkeypatch.setattr("vera_embed_openai.provider.time.sleep", lambda _seconds: None)
+    with patch("vera_embed_openai.provider.urllib.request.urlopen", fake_urlopen):
+        vectors = embedder.embed(["hello"])
+    assert attempts["n"] == 2
+    assert vectors[0].shape == (2,)
+
+
+def test_http_error_includes_status_and_body(openai_key):
+    embedder = create_embedder("text-embedding-3-small")
+
+    def fake_urlopen(request, timeout=None):
+        import vera_embed_openai.provider as provider_mod
+
+        raise provider_mod.urllib.error.HTTPError(
+            request.full_url,
+            400,
+            "bad request",
+            hdrs={},
+            fp=io.BytesIO(b'{"error":{"message":"max tokens"}}'),
+        )
+
+    with patch("vera_embed_openai.provider.urllib.request.urlopen", fake_urlopen):
+        with pytest.raises(OpenAIEmbedderError, match="HTTP 400") as exc:
+            embedder.embed(["hello"])
+    assert "max tokens" in str(exc.value)
+
+
+def test_missing_api_key_fails_in_factory(monkeypatch):
+    monkeypatch.delenv(CREDENTIAL_ENV, raising=False)
+    with pytest.raises(OpenAIEmbedderError, match="OPENAI_API_KEY"):
+        create_embedder("text-embedding-3-small")
+
+
+def test_options_validate_bounds():
+    options = OpenAIOptions.from_mapping({"batch_size": 16, "timeout": 30})
+    assert options.batch_size == 16
+    assert options.timeout == 30
+    with pytest.raises(ValueError, match="batch_size"):
+        OpenAIOptions.from_mapping({"batch_size": 0})
+    with pytest.raises(ValueError, match="timeout"):
+        OpenAIOptions.from_mapping({"timeout": 0})
+
+
+def test_ensure_registered_descriptor_and_preflight(monkeypatch, openai_key):
+    ensure_registered()
+    try:
+        descriptor = describe_embedder("openai")
+        assert descriptor.provider == "openai"
+        assert descriptor.capabilities.credential_env == CREDENTIAL_ENV
+        assert descriptor.capabilities.requires_api_key is True
+        models = list_embedding_models("openai")
+        assert any(item.model_id == DEFAULT_MODEL_ID for item in models)
+        assert preflight_embedder("openai:text-embedding-3-small").ok is True
+        monkeypatch.delenv(CREDENTIAL_ENV, raising=False)
+        failed = preflight_embedder("openai:text-embedding-3-large")
+        assert failed.ok is False
+        assert failed.missing_credential_env == CREDENTIAL_ENV
+        monkeypatch.setenv(CREDENTIAL_ENV, openai_key)
+        embedder = get_embedder("openai:text-embedding-3-small")
+        assert embedder.model_name == "openai:text-embedding-3-small"
+        assert embedder.dimension == 1536
+    finally:
+        clear_embedder_cache()
+
+
+def test_unknown_model_probes_dimension_lazily(openai_key):
+    calls = {"n": 0}
+
+    def fake_urlopen(request, timeout=None):
+        calls["n"] += 1
+        return _response([[0.6, 0.8]])
+
+    embedder = create_embedder("text-embedding-custom")
+    assert calls["n"] == 0
+    with patch("vera_embed_openai.provider.urllib.request.urlopen", fake_urlopen):
+        assert embedder.dimension == 2
+    assert calls["n"] == 1