mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-13 14:16:54 +00:00
feat(python): add OCI Generative AI text embedding function
Adds `OCIGenAIEmbeddings`, registered as `"oci-genai"`, embedding text with the Cohere models hosted on Oracle Cloud Infrastructure Generative AI (default `cohere.embed-v4.0`) via the `oci` SDK, alongside the existing Bedrock, Gemini and watsonx functions. - Auth delegated to the OCI SDK: `API_KEY` / `SECURITY_TOKEN` profiles from the OCI config file, or `INSTANCE_PRINCIPAL` / `RESOURCE_PRINCIPAL`. No credentials are stored on the function or in table metadata; `compartment_id` can come from `OCI_COMPARTMENT_ID` so the OCID need not be baked into metadata either. - Endpoint derived from `region` (default `us-chicago-1`) or overridden with `service_endpoint`; dedicated AI cluster endpoint OCIDs use `DedicatedServingMode`. - Asymmetric Cohere input types (`SEARCH_DOCUMENT` for the source column, `SEARCH_QUERY` for queries); `truncate` and `output_dimensions` (embed-v4.0: 256/512/1024/1536) forwarded. - Batched at the service limit of 96 inputs; empty inputs are not sent and come back as `None`. - `ndims()` uses `output_dimensions` when set, otherwise a single cached probe request, so there is no model→dims table to keep in sync with the catalog. - SDK retries disabled in favour of LanceDB's `max_retries` backoff; 401/403 service errors carry `status_code` so they surface without retrying, with an actionable message (expired session token, IAM policy, region, model). - `oci>=2.165.1` added to the `embeddings` extra (the floor where `EmbedTextDetails.output_dimensions` exists). Tests: 22 mocked unit tests (auth paths, batching, input types, config validation, retry behaviour, pickling, table-metadata round trip) and a live test skipped unless `OCI_COMPARTMENT_ID` is set, which passes against us-chicago-1 for embed-v4.0 at 1536 and 512 dims and embed-multilingual-v3.0 at 1024. Verified end to end with a real LanceDB table (`LanceModel` with `SourceField`/`VectorField`, 6 rows, two semantic queries ranked correctly). `ruff format --check`, `ruff check`, pyright and typos are clean. Docs: a page under `docs/src/embeddings/available_embedding_models/text_embedding_functions/` in the shape of the remaining `voyageai_embedding.md`; the function is exported from `lancedb.embeddings` so the API reference picks it up. The hosted docs' available-models table lives outside this repo and would need a separate update.
This commit is contained in:
Fede Kamelhar
@@ -13,6 +13,7 @@ from .bedrock import BedRockText
|
||||
from .cohere import CohereEmbeddingFunction
|
||||
from .gemini_text import GeminiText
|
||||
from .instructor import InstructorEmbeddingFunction
|
||||
from .oci_genai import OCIGenAIEmbeddings
|
||||
from .ollama import OllamaEmbeddings
|
||||
from .open_clip import OpenClipEmbeddings
|
||||
from .openai import OpenAIEmbeddings
|
||||
@@ -49,6 +50,7 @@ __all__ = [
|
||||
"GteEmbeddings",
|
||||
"InstructorEmbeddingFunction",
|
||||
"JinaEmbeddings",
|
||||
"OCIGenAIEmbeddings",
|
||||
"OllamaEmbeddings",
|
||||
"TransformersEmbeddingFunction",
|
||||
"ColbertEmbeddings",
|
||||
|
||||
@@ -0,0 +1,320 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
|
||||
import os
|
||||
from functools import cached_property
|
||||
from typing import Any, List, Literal, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
from pydantic import ConfigDict
|
||||
|
||||
from ..util import attempt_import_or_raise
|
||||
from .base import TextEmbeddingFunction
|
||||
from .registry import register
|
||||
from .utils import TEXT
|
||||
|
||||
# The Generative AI embed-text API accepts at most 96 inputs per request:
|
||||
# https://docs.oracle.com/en-us/iaas/api/#/en/generative-ai-inference/latest/EmbedTextResult/EmbedText
|
||||
EMBEDDING_BATCH_SIZE = 96
|
||||
|
||||
# A ``name`` that is a dedicated AI cluster endpoint OCID is routed with
|
||||
# DedicatedServingMode instead of the on-demand model catalog.
|
||||
_DEDICATED_ENDPOINT_PREFIX = "ocid1.generativeaiendpoint"
|
||||
_DEFAULT_REGION = "us-chicago-1"
|
||||
|
||||
_SERVICE_ERROR_HELP = """
|
||||
Common causes:
|
||||
- Expired session token (auth_type="SECURITY_TOKEN"): run
|
||||
`oci session authenticate --profile-name <profile>`.
|
||||
- Missing IAM policy: the caller needs
|
||||
`allow group <group> to use generative-ai-family in compartment <name>`.
|
||||
- Model not offered in the region behind `region` / `service_endpoint`, or
|
||||
a wrong `name` / `compartment_id`. List models with
|
||||
`oci generative-ai model-collection list-models --compartment-id <ocid>`.
|
||||
"""
|
||||
|
||||
|
||||
class OCIGenAIServiceError(ValueError):
|
||||
"""Raised when the Generative AI service rejects an ``embed_text`` request.
|
||||
|
||||
``status_code`` mirrors the HTTP status so LanceDB's retry helper treats
|
||||
401/403 responses as permanent instead of retrying them.
|
||||
"""
|
||||
|
||||
def __init__(self, message: str, status_code: Optional[int] = None):
|
||||
super().__init__(message)
|
||||
self.status_code = status_code
|
||||
|
||||
|
||||
@register("oci-genai")
|
||||
class OCIGenAIEmbeddings(TextEmbeddingFunction):
|
||||
"""
|
||||
An embedding function that uses the Oracle Cloud Infrastructure (OCI)
|
||||
Generative AI service.
|
||||
|
||||
https://docs.oracle.com/en-us/iaas/Content/generative-ai/embed-models.htm
|
||||
|
||||
Requires the OCI Python SDK (``pip install oci``). No credentials are stored
|
||||
on this object: requests are signed by the SDK using an API-key or
|
||||
session-token profile from the OCI config file, or the instance / resource
|
||||
principal of the host, so there are no sensitive keys to protect with
|
||||
``$var:``.
|
||||
|
||||
Cohere embedding models are asymmetric: the source column is embedded with
|
||||
``input_type="SEARCH_DOCUMENT"`` and queries with ``"SEARCH_QUERY"``.
|
||||
Inputs are sent in batches of at most 96 texts, the service limit.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
name : str, default "cohere.embed-v4.0"
|
||||
Model ID of an on-demand embedding model: ``cohere.embed-v4.0``
|
||||
(1536 dims), ``cohere.embed-english-v3.0`` /
|
||||
``cohere.embed-multilingual-v3.0`` (1024 dims) or their ``-light-v3.0``
|
||||
variants (384 dims). A dedicated AI cluster endpoint OCID
|
||||
(``ocid1.generativeaiendpoint...``) is also accepted.
|
||||
compartment_id : str, optional
|
||||
OCID of the compartment the requests are authorized in and billed to.
|
||||
The service requires it; when not passed, the ``OCI_COMPARTMENT_ID``
|
||||
environment variable is used.
|
||||
region : str, default "us-chicago-1"
|
||||
OCI region used to derive the inference endpoint when
|
||||
``service_endpoint`` is not set.
|
||||
service_endpoint : str, optional
|
||||
Full inference endpoint, e.g.
|
||||
``https://inference.generativeai.us-chicago-1.oci.oraclecloud.com``.
|
||||
Takes precedence over ``region``; needed for realms with a different
|
||||
endpoint template.
|
||||
auth_type : str, default "API_KEY"
|
||||
One of ``"API_KEY"``, ``"SECURITY_TOKEN"``, ``"INSTANCE_PRINCIPAL"`` or
|
||||
``"RESOURCE_PRINCIPAL"``. The first two read profile ``auth_profile``
|
||||
from the OCI config file ``auth_file_location``; ``SECURITY_TOKEN``
|
||||
additionally uses the profile's ``security_token_file`` written by
|
||||
``oci session authenticate``. The principal types use the identity of
|
||||
the host (Compute instance, OKE workload, Functions, ...) and ignore
|
||||
the config file.
|
||||
auth_profile : str, default "DEFAULT"
|
||||
Profile name in the OCI config file.
|
||||
auth_file_location : str, default "~/.oci/config"
|
||||
Path of the OCI config file.
|
||||
truncate : str, default "END"
|
||||
How inputs longer than the model context are handled: ``"NONE"``
|
||||
(fail), ``"START"`` or ``"END"`` (drop tokens from that side).
|
||||
source_input_type : str, default "SEARCH_DOCUMENT"
|
||||
``input_type`` sent when embedding the source column.
|
||||
query_input_type : str, default "SEARCH_QUERY"
|
||||
``input_type`` sent when embedding queries.
|
||||
output_dimensions : int, optional
|
||||
Embedding size for models that support it (``cohere.embed-v4.0``:
|
||||
256, 512, 1024 or 1536). When set, ``ndims()`` returns it directly;
|
||||
otherwise the size is discovered with one probe request and cached.
|
||||
|
||||
Examples
|
||||
--------
|
||||
import lancedb
|
||||
from lancedb.embeddings import get_registry
|
||||
from lancedb.pydantic import LanceModel, Vector
|
||||
|
||||
model = get_registry().get("oci-genai").create(
|
||||
name="cohere.embed-v4.0",
|
||||
compartment_id="ocid1.compartment.oc1..<your-compartment>",
|
||||
auth_profile="DEFAULT",
|
||||
)
|
||||
|
||||
class TextModel(LanceModel):
|
||||
text: str = model.SourceField()
|
||||
vector: Vector(model.ndims()) = model.VectorField()
|
||||
|
||||
db = lancedb.connect("~/.lancedb")
|
||||
tbl = db.create_table("docs", schema=TextModel, mode="overwrite")
|
||||
tbl.add([{"text": "hello world"}, {"text": "goodbye world"}])
|
||||
rs = tbl.search("hello").limit(1).to_pandas()
|
||||
|
||||
When restored from table metadata in a new process, an explicit
|
||||
``service_endpoint`` or non-default ``region`` must be approved with
|
||||
[allow_endpoint][lancedb.embeddings.EmbeddingFunctionRegistry.allow_endpoint].
|
||||
The default region does not need approval.
|
||||
"""
|
||||
|
||||
name: str = "cohere.embed-v4.0"
|
||||
compartment_id: Optional[str] = None
|
||||
region: str = _DEFAULT_REGION
|
||||
service_endpoint: Optional[str] = None
|
||||
auth_type: Literal[
|
||||
"API_KEY", "SECURITY_TOKEN", "INSTANCE_PRINCIPAL", "RESOURCE_PRINCIPAL"
|
||||
] = "API_KEY"
|
||||
auth_profile: str = "DEFAULT"
|
||||
auth_file_location: str = "~/.oci/config"
|
||||
truncate: Literal["NONE", "START", "END"] = "END"
|
||||
source_input_type: str = "SEARCH_DOCUMENT"
|
||||
query_input_type: str = "SEARCH_QUERY"
|
||||
output_dimensions: Optional[int] = None
|
||||
|
||||
model_config = ConfigDict(ignored_types=(cached_property,))
|
||||
|
||||
@staticmethod
|
||||
def model_names() -> List[str]:
|
||||
"""On-demand text embedding models offered by the service."""
|
||||
return [
|
||||
"cohere.embed-v4.0",
|
||||
"cohere.embed-english-v3.0",
|
||||
"cohere.embed-multilingual-v3.0",
|
||||
"cohere.embed-english-light-v3.0",
|
||||
"cohere.embed-multilingual-light-v3.0",
|
||||
]
|
||||
|
||||
def endpoints(self) -> List[str]:
|
||||
# Metadata only selects the default endpoint by omitting both fields.
|
||||
if self.service_endpoint or self.region != _DEFAULT_REGION:
|
||||
return [self._resolve_service_endpoint()]
|
||||
return []
|
||||
|
||||
def ndims(self) -> int:
|
||||
return self._resolved_ndims
|
||||
|
||||
@cached_property
|
||||
def _resolved_ndims(self) -> int:
|
||||
if self.output_dimensions is not None:
|
||||
return self.output_dimensions
|
||||
# embed-v4.0 is variable-width and the catalog changes, so probe the
|
||||
# service once instead of hardcoding a model -> dimensions table.
|
||||
return len(self._embed_batch(["lancedb"], self.query_input_type)[0])
|
||||
|
||||
def compute_query_embeddings(
|
||||
self, query: str, *args, **kwargs
|
||||
) -> List[Union[List[float], None]]:
|
||||
return self.compute_source_embeddings(query, input_type=self.query_input_type)
|
||||
|
||||
def compute_source_embeddings(
|
||||
self, texts: TEXT, *args, **kwargs
|
||||
) -> List[Union[List[float], None]]:
|
||||
texts = self.sanitize_input(texts)
|
||||
# assume the source input type unless `compute_query_embeddings` set one
|
||||
kwargs["input_type"] = kwargs.get("input_type") or self.source_input_type
|
||||
return self.generate_embeddings(texts, **kwargs)
|
||||
|
||||
def generate_embeddings(
|
||||
self, texts: Union[List[str], np.ndarray], *args, **kwargs
|
||||
) -> List[Union[List[float], None]]:
|
||||
"""
|
||||
Get the embeddings for the given texts
|
||||
|
||||
Parameters
|
||||
----------
|
||||
texts: list[str] or np.ndarray (of str)
|
||||
The texts to embed
|
||||
|
||||
Returns
|
||||
-------
|
||||
list[list[float] | None]
|
||||
One embedding per input, in order. Empty inputs are not sent to
|
||||
the service (it rejects them) and come back as ``None``.
|
||||
"""
|
||||
input_type = kwargs.get("input_type") or self.source_input_type
|
||||
texts = list(texts)
|
||||
valid = [(idx, text) for idx, text in enumerate(texts) if text]
|
||||
results: List[Union[List[float], None]] = [None] * len(texts)
|
||||
for start in range(0, len(valid), EMBEDDING_BATCH_SIZE):
|
||||
batch = valid[start : start + EMBEDDING_BATCH_SIZE]
|
||||
embeddings = self._embed_batch([text for _, text in batch], input_type)
|
||||
for (idx, _), embedding in zip(batch, embeddings):
|
||||
results[idx] = embedding
|
||||
return results
|
||||
|
||||
def _embed_batch(self, texts: List[str], input_type: str) -> List[List[float]]:
|
||||
"""Send one embed_text request (at most EMBEDDING_BATCH_SIZE inputs)."""
|
||||
oci = attempt_import_or_raise("oci")
|
||||
models = oci.generative_ai_inference.models
|
||||
if self.name.startswith(_DEDICATED_ENDPOINT_PREFIX):
|
||||
serving_mode = models.DedicatedServingMode(endpoint_id=self.name)
|
||||
else:
|
||||
serving_mode = models.OnDemandServingMode(model_id=self.name)
|
||||
details = {
|
||||
"inputs": texts,
|
||||
"serving_mode": serving_mode,
|
||||
"compartment_id": self._resolve_compartment_id(),
|
||||
"truncate": self.truncate,
|
||||
"input_type": input_type,
|
||||
}
|
||||
if self.output_dimensions is not None:
|
||||
details["output_dimensions"] = self.output_dimensions
|
||||
try:
|
||||
response = self._client.embed_text(models.EmbedTextDetails(**details))
|
||||
except oci.exceptions.ServiceError as e:
|
||||
raise OCIGenAIServiceError(
|
||||
f"OCI Generative AI embed_text failed with HTTP {e.status} "
|
||||
f"({e.code}): {e.message}\n{_SERVICE_ERROR_HELP}",
|
||||
status_code=e.status,
|
||||
) from e
|
||||
return response.data.embeddings
|
||||
|
||||
def _resolve_compartment_id(self) -> str:
|
||||
compartment_id = self.compartment_id or os.environ.get("OCI_COMPARTMENT_ID")
|
||||
if not compartment_id:
|
||||
raise ValueError(
|
||||
"compartment_id is required by the OCI Generative AI service. "
|
||||
"Pass `compartment_id=` to OCIGenAIEmbeddings or set the "
|
||||
"OCI_COMPARTMENT_ID environment variable."
|
||||
)
|
||||
return compartment_id
|
||||
|
||||
def _resolve_service_endpoint(self) -> str:
|
||||
if self.service_endpoint:
|
||||
return self.service_endpoint
|
||||
return f"https://inference.generativeai.{self.region}.oci.oraclecloud.com"
|
||||
|
||||
def _build_auth(self, oci) -> tuple[dict, Any]:
|
||||
"""Return the ``(config, signer)`` pair for the configured ``auth_type``."""
|
||||
if self.auth_type in ("API_KEY", "SECURITY_TOKEN"):
|
||||
config = oci.config.from_file(
|
||||
file_location=self.auth_file_location, profile_name=self.auth_profile
|
||||
)
|
||||
token_file = config.get("security_token_file")
|
||||
if self.auth_type == "API_KEY":
|
||||
if token_file:
|
||||
raise ValueError(
|
||||
f"Profile '{self.auth_profile}' in {self.auth_file_location} "
|
||||
"is a session-token profile; use auth_type='SECURITY_TOKEN'."
|
||||
)
|
||||
return config, None
|
||||
if not token_file:
|
||||
raise ValueError(
|
||||
f"Profile '{self.auth_profile}' in {self.auth_file_location} "
|
||||
"has no security_token_file; run `oci session authenticate` "
|
||||
"or use auth_type='API_KEY'."
|
||||
)
|
||||
with open(os.path.expanduser(token_file)) as f:
|
||||
token = f.read().strip()
|
||||
private_key = oci.signer.load_private_key_from_file(
|
||||
os.path.expanduser(config["key_file"]), config.get("pass_phrase")
|
||||
)
|
||||
return config, oci.auth.signers.SecurityTokenSigner(token, private_key)
|
||||
if self.auth_type == "INSTANCE_PRINCIPAL":
|
||||
return {}, oci.auth.signers.InstancePrincipalsSecurityTokenSigner()
|
||||
if self.auth_type == "RESOURCE_PRINCIPAL":
|
||||
return {}, oci.auth.signers.get_resource_principals_signer()
|
||||
raise ValueError(f"Unsupported auth_type '{self.auth_type}'")
|
||||
|
||||
@cached_property
|
||||
def _client(self):
|
||||
oci = attempt_import_or_raise("oci")
|
||||
config, signer = self._build_auth(oci)
|
||||
kwargs = {
|
||||
"config": config,
|
||||
"service_endpoint": self._resolve_service_endpoint(),
|
||||
# LanceDB already retries with exponential backoff (`max_retries`),
|
||||
# so the SDK's own retry loop is disabled to avoid compounding it.
|
||||
"retry_strategy": oci.retry.NoneRetryStrategy(),
|
||||
"timeout": (10, 240),
|
||||
}
|
||||
if signer is not None:
|
||||
kwargs["signer"] = signer
|
||||
return oci.generative_ai_inference.GenerativeAiInferenceClient(**kwargs)
|
||||
|
||||
def __getstate__(self) -> dict[str, Any]:
|
||||
# The SDK client holds an HTTP session and cannot be pickled.
|
||||
state = super().__getstate__()
|
||||
state["__dict__"] = {
|
||||
k: v for k, v in state["__dict__"].items() if k != "_client"
|
||||
}
|
||||
return state
|
||||
@@ -0,0 +1,351 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
"""Unit tests for the OCI Generative AI embedding function.
|
||||
|
||||
The ``oci`` SDK is replaced with a small fake so these run without the package
|
||||
or cloud credentials. Live coverage lives in ``test_embeddings_slow.py``.
|
||||
"""
|
||||
|
||||
import os
|
||||
import pickle
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import lancedb
|
||||
import pytest
|
||||
from lancedb.embeddings import OCIGenAIEmbeddings, get_registry
|
||||
from lancedb.embeddings.oci_genai import EMBEDDING_BATCH_SIZE, OCIGenAIServiceError
|
||||
from lancedb.pydantic import LanceModel, Vector
|
||||
from pydantic import ValidationError
|
||||
|
||||
DIMS = 8
|
||||
COMPARTMENT = "ocid1.compartment.oc1..test"
|
||||
|
||||
|
||||
class FakeServiceError(Exception):
|
||||
"""Mimics ``oci.exceptions.ServiceError``."""
|
||||
|
||||
def __init__(self, status, code, message):
|
||||
super().__init__(message)
|
||||
self.status = status
|
||||
self.code = code
|
||||
self.message = message
|
||||
|
||||
|
||||
def _fake_embed_text(details):
|
||||
# Deterministic embeddings: first component encodes len(text) so ordering
|
||||
# and nearest-neighbour behaviour can be asserted.
|
||||
embeddings = [[float(len(text))] + [0.0] * (DIMS - 1) for text in details.inputs]
|
||||
return SimpleNamespace(data=SimpleNamespace(embeddings=embeddings))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_oci():
|
||||
"""Stand-in for the ``oci`` package, injected via attempt_import_or_raise."""
|
||||
oci = MagicMock(name="oci")
|
||||
oci.exceptions.ServiceError = FakeServiceError
|
||||
models = oci.generative_ai_inference.models
|
||||
models.EmbedTextDetails.side_effect = lambda **kw: SimpleNamespace(**kw)
|
||||
models.OnDemandServingMode.side_effect = lambda **kw: SimpleNamespace(
|
||||
serving_type="ON_DEMAND", **kw
|
||||
)
|
||||
models.DedicatedServingMode.side_effect = lambda **kw: SimpleNamespace(
|
||||
serving_type="DEDICATED", **kw
|
||||
)
|
||||
oci.config.from_file.return_value = {"key_file": "~/.oci/key.pem"}
|
||||
client = oci.generative_ai_inference.GenerativeAiInferenceClient.return_value
|
||||
client.embed_text.side_effect = _fake_embed_text
|
||||
with patch(
|
||||
"lancedb.embeddings.oci_genai.attempt_import_or_raise", return_value=oci
|
||||
):
|
||||
yield oci
|
||||
|
||||
|
||||
def _create(**kwargs):
|
||||
kwargs.setdefault("compartment_id", COMPARTMENT)
|
||||
kwargs.setdefault("max_retries", 0)
|
||||
return get_registry().get("oci-genai").create(**kwargs)
|
||||
|
||||
|
||||
def _client(fake_oci):
|
||||
return fake_oci.generative_ai_inference.GenerativeAiInferenceClient.return_value
|
||||
|
||||
|
||||
def _client_kwargs(fake_oci):
|
||||
return fake_oci.generative_ai_inference.GenerativeAiInferenceClient.call_args.kwargs
|
||||
|
||||
|
||||
def _embed_requests(fake_oci):
|
||||
return [call.args[0] for call in _client(fake_oci).embed_text.call_args_list]
|
||||
|
||||
|
||||
def test_oci_genai_registered():
|
||||
assert get_registry().get("oci-genai") is OCIGenAIEmbeddings
|
||||
assert "cohere.embed-v4.0" in OCIGenAIEmbeddings.model_names()
|
||||
func = _create()
|
||||
assert func.name == "cohere.embed-v4.0"
|
||||
assert func.auth_type == "API_KEY"
|
||||
assert func.truncate == "END"
|
||||
assert func.sensitive_keys() == []
|
||||
|
||||
|
||||
def test_oci_genai_service_endpoint(fake_oci):
|
||||
default = _create()
|
||||
assert default.endpoints() == []
|
||||
default._client
|
||||
assert (
|
||||
_client_kwargs(fake_oci)["service_endpoint"]
|
||||
== "https://inference.generativeai.us-chicago-1.oci.oraclecloud.com"
|
||||
)
|
||||
|
||||
regional = _create(region="eu-frankfurt-1")
|
||||
regional_endpoint = (
|
||||
"https://inference.generativeai.eu-frankfurt-1.oci.oraclecloud.com"
|
||||
)
|
||||
assert regional.endpoints() == [regional_endpoint]
|
||||
regional._client
|
||||
assert _client_kwargs(fake_oci)["service_endpoint"] == regional_endpoint
|
||||
|
||||
custom = "https://inference.generativeai.us-langley-1.oci.oraclegovcloud.com"
|
||||
explicit = _create(region="us-chicago-1", service_endpoint=custom)
|
||||
assert explicit.endpoints() == [custom]
|
||||
explicit._client
|
||||
assert _client_kwargs(fake_oci)["service_endpoint"] == custom
|
||||
|
||||
|
||||
def test_oci_genai_client_disables_sdk_retries(fake_oci):
|
||||
_create()._client
|
||||
kwargs = _client_kwargs(fake_oci)
|
||||
assert kwargs["retry_strategy"] is fake_oci.retry.NoneRetryStrategy.return_value
|
||||
|
||||
|
||||
def test_oci_genai_api_key_auth(fake_oci, tmp_path):
|
||||
config_file = str(tmp_path / "config")
|
||||
func = _create(auth_profile="MYPROFILE", auth_file_location=config_file)
|
||||
func._client
|
||||
fake_oci.config.from_file.assert_called_once_with(
|
||||
file_location=config_file, profile_name="MYPROFILE"
|
||||
)
|
||||
kwargs = _client_kwargs(fake_oci)
|
||||
assert kwargs["config"] is fake_oci.config.from_file.return_value
|
||||
assert "signer" not in kwargs
|
||||
|
||||
|
||||
def test_oci_genai_api_key_rejects_session_profile(fake_oci):
|
||||
fake_oci.config.from_file.return_value = {
|
||||
"key_file": "~/.oci/sessions/x/key.pem",
|
||||
"security_token_file": "~/.oci/sessions/x/token",
|
||||
}
|
||||
with pytest.raises(ValueError, match="SECURITY_TOKEN"):
|
||||
_create(auth_type="API_KEY")._client
|
||||
|
||||
|
||||
def test_oci_genai_security_token_auth(fake_oci, tmp_path):
|
||||
token_file = tmp_path / "token"
|
||||
token_file.write_text("session-token\n")
|
||||
fake_oci.config.from_file.return_value = {
|
||||
"key_file": "~/.oci/sessions/x/key.pem",
|
||||
"security_token_file": str(token_file),
|
||||
}
|
||||
_create(auth_type="SECURITY_TOKEN", auth_profile="SESSION")._client
|
||||
|
||||
fake_oci.signer.load_private_key_from_file.assert_called_once_with(
|
||||
os.path.expanduser("~/.oci/sessions/x/key.pem"), None
|
||||
)
|
||||
fake_oci.auth.signers.SecurityTokenSigner.assert_called_once_with(
|
||||
"session-token", fake_oci.signer.load_private_key_from_file.return_value
|
||||
)
|
||||
kwargs = _client_kwargs(fake_oci)
|
||||
assert kwargs["signer"] is fake_oci.auth.signers.SecurityTokenSigner.return_value
|
||||
assert kwargs["config"] is fake_oci.config.from_file.return_value
|
||||
|
||||
|
||||
def test_oci_genai_security_token_requires_token_file(fake_oci):
|
||||
with pytest.raises(ValueError, match="security_token_file"):
|
||||
_create(auth_type="SECURITY_TOKEN")._client
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"auth_type,factory",
|
||||
[
|
||||
("INSTANCE_PRINCIPAL", "InstancePrincipalsSecurityTokenSigner"),
|
||||
("RESOURCE_PRINCIPAL", "get_resource_principals_signer"),
|
||||
],
|
||||
)
|
||||
def test_oci_genai_principal_auth(fake_oci, auth_type, factory):
|
||||
_create(auth_type=auth_type)._client
|
||||
signer_factory = getattr(fake_oci.auth.signers, factory)
|
||||
signer_factory.assert_called_once_with()
|
||||
kwargs = _client_kwargs(fake_oci)
|
||||
assert kwargs["config"] == {}
|
||||
assert kwargs["signer"] is signer_factory.return_value
|
||||
fake_oci.config.from_file.assert_not_called()
|
||||
|
||||
|
||||
def test_oci_genai_rejects_invalid_config():
|
||||
with pytest.raises(ValidationError, match="auth_type"):
|
||||
_create(auth_type="PASSWORD")
|
||||
with pytest.raises(ValidationError, match="truncate"):
|
||||
_create(truncate="MIDDLE")
|
||||
|
||||
|
||||
def test_oci_genai_compartment_id_required(fake_oci, monkeypatch):
|
||||
monkeypatch.delenv("OCI_COMPARTMENT_ID", raising=False)
|
||||
func = get_registry().get("oci-genai").create(max_retries=0)
|
||||
with pytest.raises(ValueError, match="OCI_COMPARTMENT_ID"):
|
||||
func.compute_source_embeddings(["hello"])
|
||||
assert _embed_requests(fake_oci) == []
|
||||
|
||||
monkeypatch.setenv("OCI_COMPARTMENT_ID", "ocid1.compartment.oc1..from-env")
|
||||
func.compute_source_embeddings(["hello"])
|
||||
assert _embed_requests(fake_oci)[0].compartment_id == (
|
||||
"ocid1.compartment.oc1..from-env"
|
||||
)
|
||||
|
||||
|
||||
def test_oci_genai_request_shape(fake_oci):
|
||||
func = _create(name="cohere.embed-multilingual-v3.0", truncate="START")
|
||||
func.compute_source_embeddings(["hello"])
|
||||
(details,) = _embed_requests(fake_oci)
|
||||
assert details.inputs == ["hello"]
|
||||
assert details.compartment_id == COMPARTMENT
|
||||
assert details.truncate == "START"
|
||||
assert details.input_type == "SEARCH_DOCUMENT"
|
||||
assert details.serving_mode.serving_type == "ON_DEMAND"
|
||||
assert details.serving_mode.model_id == "cohere.embed-multilingual-v3.0"
|
||||
# only sent when configured so older SDKs / models keep working
|
||||
assert not hasattr(details, "output_dimensions")
|
||||
|
||||
|
||||
def test_oci_genai_dedicated_endpoint(fake_oci):
|
||||
endpoint = "ocid1.generativeaiendpoint.oc1.us-chicago-1.test"
|
||||
_create(name=endpoint).compute_source_embeddings(["hello"])
|
||||
(details,) = _embed_requests(fake_oci)
|
||||
assert details.serving_mode.serving_type == "DEDICATED"
|
||||
assert details.serving_mode.endpoint_id == endpoint
|
||||
|
||||
|
||||
def test_oci_genai_query_and_source_input_types(fake_oci):
|
||||
func = _create()
|
||||
func.compute_source_embeddings(["a document"])
|
||||
func.compute_query_embeddings("a query")
|
||||
source, query = _embed_requests(fake_oci)
|
||||
assert source.input_type == "SEARCH_DOCUMENT"
|
||||
assert query.input_type == "SEARCH_QUERY"
|
||||
|
||||
custom = _create(source_input_type="CLUSTERING", query_input_type="CLASSIFICATION")
|
||||
custom.compute_source_embeddings(["a document"])
|
||||
custom.compute_query_embeddings("a query")
|
||||
source, query = _embed_requests(fake_oci)[2:]
|
||||
assert source.input_type == "CLUSTERING"
|
||||
assert query.input_type == "CLASSIFICATION"
|
||||
|
||||
|
||||
def test_oci_genai_batches_at_service_limit(fake_oci):
|
||||
assert EMBEDDING_BATCH_SIZE == 96
|
||||
texts = [f"text {i:03d}" for i in range(2 * EMBEDDING_BATCH_SIZE + 8)]
|
||||
embeddings = _create().compute_source_embeddings(texts)
|
||||
|
||||
requests = _embed_requests(fake_oci)
|
||||
assert [len(r.inputs) for r in requests] == [96, 96, 8]
|
||||
assert [t for r in requests for t in r.inputs] == texts
|
||||
assert all(r.input_type == "SEARCH_DOCUMENT" for r in requests)
|
||||
assert len(embeddings) == len(texts)
|
||||
assert [e[0] for e in embeddings] == [float(len(t)) for t in texts]
|
||||
|
||||
|
||||
def test_oci_genai_empty_inputs_are_skipped(fake_oci):
|
||||
func = _create()
|
||||
assert func.compute_source_embeddings([]) == []
|
||||
assert _embed_requests(fake_oci) == []
|
||||
|
||||
embeddings = func.compute_source_embeddings(["hello", "", None, "world"])
|
||||
(details,) = _embed_requests(fake_oci)
|
||||
assert details.inputs == ["hello", "world"]
|
||||
assert embeddings[1] is None
|
||||
assert embeddings[2] is None
|
||||
assert embeddings[0][0] == 5.0
|
||||
assert embeddings[3][0] == 5.0
|
||||
|
||||
|
||||
def test_oci_genai_ndims_probe_is_cached(fake_oci):
|
||||
func = _create()
|
||||
assert func.ndims() == DIMS
|
||||
assert func.ndims() == DIMS
|
||||
assert len(_embed_requests(fake_oci)) == 1
|
||||
|
||||
|
||||
def test_oci_genai_output_dimensions(fake_oci):
|
||||
func = _create(output_dimensions=512)
|
||||
assert func.ndims() == 512
|
||||
assert _embed_requests(fake_oci) == [] # no probe request needed
|
||||
|
||||
func.compute_source_embeddings(["hello"])
|
||||
(details,) = _embed_requests(fake_oci)
|
||||
assert details.output_dimensions == 512
|
||||
|
||||
|
||||
@patch("time.sleep")
|
||||
def test_oci_genai_auth_errors_are_not_retried(mock_sleep, fake_oci):
|
||||
_client(fake_oci).embed_text.side_effect = FakeServiceError(
|
||||
401, "NotAuthenticated", "The required information was not provided."
|
||||
)
|
||||
func = _create(max_retries=3)
|
||||
with pytest.raises(OCIGenAIServiceError, match="HTTP 401") as exc_info:
|
||||
func.compute_source_embeddings_with_retry(["hello"])
|
||||
|
||||
assert exc_info.value.status_code == 401
|
||||
assert "oci session authenticate" in str(exc_info.value)
|
||||
assert isinstance(exc_info.value.__cause__, FakeServiceError)
|
||||
assert _client(fake_oci).embed_text.call_count == 1
|
||||
mock_sleep.assert_not_called()
|
||||
|
||||
|
||||
@patch("time.sleep")
|
||||
def test_oci_genai_server_errors_are_retried(mock_sleep, fake_oci):
|
||||
_client(fake_oci).embed_text.side_effect = FakeServiceError(
|
||||
500, "InternalError", "try again"
|
||||
)
|
||||
func = _create(max_retries=2)
|
||||
with pytest.raises(Exception, match="Maximum number of retries"):
|
||||
func.compute_source_embeddings_with_retry(["hello"])
|
||||
assert _client(fake_oci).embed_text.call_count == 3
|
||||
|
||||
|
||||
def test_oci_genai_pickle_after_client_created(fake_oci):
|
||||
func = _create(auth_profile="MYPROFILE")
|
||||
func.compute_source_embeddings(["hello"])
|
||||
assert "_client" in func.__dict__
|
||||
|
||||
restored = pickle.loads(pickle.dumps(func))
|
||||
assert "_client" not in restored.__dict__
|
||||
assert restored.compartment_id == COMPARTMENT
|
||||
assert restored.auth_profile == "MYPROFILE"
|
||||
|
||||
|
||||
def test_oci_genai_table_round_trip(fake_oci):
|
||||
func = _create(region="us-ashburn-1", auth_profile="MYPROFILE")
|
||||
|
||||
class Docs(LanceModel):
|
||||
text: str = func.SourceField()
|
||||
vector: Vector(func.ndims()) = func.VectorField()
|
||||
|
||||
db = lancedb.connect("memory://")
|
||||
tbl = db.create_table("docs", schema=Docs, mode="overwrite")
|
||||
tbl.add([{"text": "hello"}, {"text": "goodbye world"}])
|
||||
assert len(tbl.to_pandas()["vector"][0]) == DIMS
|
||||
|
||||
# the fake embeds len(text), so an equally long query lands on that row
|
||||
hits = tbl.search("goodbye world").limit(1).to_pandas()
|
||||
assert hits["text"][0] == "goodbye world"
|
||||
assert _embed_requests(fake_oci)[-1].input_type == "SEARCH_QUERY"
|
||||
|
||||
parsed = tbl.embedding_functions["vector"].function
|
||||
assert isinstance(parsed, OCIGenAIEmbeddings)
|
||||
assert parsed.safe_model_dump() == {
|
||||
"compartment_id": COMPARTMENT,
|
||||
"max_retries": 0,
|
||||
"region": "us-ashburn-1",
|
||||
"auth_profile": "MYPROFILE",
|
||||
}
|
||||
@@ -379,6 +379,58 @@ def test_bedrock_embedding(tmp_path):
|
||||
assert len(tbl.to_pandas()["vector"][0]) == model.ndims()
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.skipif(
|
||||
os.environ.get("OCI_COMPARTMENT_ID") is None, reason="OCI_COMPARTMENT_ID not set"
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"model_name,output_dimensions,expected_dims",
|
||||
[
|
||||
("cohere.embed-v4.0", None, 1536),
|
||||
("cohere.embed-v4.0", 512, 512),
|
||||
("cohere.embed-multilingual-v3.0", None, 1024),
|
||||
],
|
||||
)
|
||||
def test_oci_genai_embedding(model_name, output_dimensions, expected_dims, tmp_path):
|
||||
pytest.importorskip("oci")
|
||||
# compartment_id comes from OCI_COMPARTMENT_ID; the profile / auth follow the
|
||||
# OCI CLI's own OCI_CLI_PROFILE / OCI_CLI_AUTH / OCI_CLI_REGION variables so
|
||||
# both API keys and `oci session authenticate` sessions work.
|
||||
security_token = os.environ.get("OCI_CLI_AUTH") == "security_token"
|
||||
model = (
|
||||
get_registry()
|
||||
.get("oci-genai")
|
||||
.create(
|
||||
name=model_name,
|
||||
output_dimensions=output_dimensions,
|
||||
auth_type="SECURITY_TOKEN" if security_token else "API_KEY",
|
||||
auth_profile=os.environ.get("OCI_CLI_PROFILE", "DEFAULT"),
|
||||
region=os.environ.get("OCI_CLI_REGION", "us-chicago-1"),
|
||||
max_retries=0,
|
||||
)
|
||||
)
|
||||
assert model.ndims() == expected_dims
|
||||
|
||||
class TextModel(LanceModel):
|
||||
text: str = model.SourceField()
|
||||
vector: Vector(model.ndims()) = model.VectorField()
|
||||
|
||||
db = lancedb.connect(tmp_path)
|
||||
tbl = db.create_table("oci_genai_test", schema=TextModel, mode="overwrite")
|
||||
df = pd.DataFrame(
|
||||
{"text": ["hello world", "goodbye world", "a recipe for lasagna", ""]}
|
||||
)
|
||||
tbl.add(df, on_bad_vectors="null")
|
||||
|
||||
tb = tbl.to_arrow()
|
||||
assert tb.schema.field_by_name("vector").type == pa.list_(
|
||||
pa.float32(), expected_dims
|
||||
)
|
||||
# the empty input is not sent to the service and is stored as null
|
||||
assert tb["vector"].is_null().to_pylist() == [False, False, False, True]
|
||||
assert tbl.search("hello").limit(1).to_pandas()["text"][0] == "hello world"
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.skipif(
|
||||
os.environ.get("OPENAI_API_KEY") is None, reason="OPENAI_API_KEY not set"
|
||||
|
||||
@@ -15,6 +15,7 @@ import pyarrow as pa
|
||||
import pytest
|
||||
from lancedb.embeddings import (
|
||||
EmbeddingFunctionConfig,
|
||||
OCIGenAIEmbeddings,
|
||||
OllamaEmbeddings,
|
||||
OpenAIEmbeddings,
|
||||
SentenceTransformerEmbeddings,
|
||||
@@ -220,6 +221,9 @@ def test_trusting_one_model_does_not_trust_a_swapped_name():
|
||||
|
||||
|
||||
ATTACKER_URL = "https://attacker.example/v1"
|
||||
OCI_REGION_ENDPOINT = (
|
||||
"https://inference.generativeai.eu-frankfurt-1.oci.oraclecloud.com"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -247,6 +251,31 @@ def test_default_endpoints_need_no_allowlist():
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model,endpoint",
|
||||
[
|
||||
({"service_endpoint": ATTACKER_URL}, ATTACKER_URL),
|
||||
({"region": "eu-frankfurt-1"}, OCI_REGION_ENDPOINT),
|
||||
],
|
||||
)
|
||||
def test_restored_oci_custom_endpoint_requires_allowlist(model, endpoint):
|
||||
func = restore("oci-genai", model)
|
||||
assert isinstance(func, UntrustedEmbeddingFunction)
|
||||
with pytest.raises(UntrustedEmbeddingConfigError, match="allow_endpoint"):
|
||||
func.compute_query_embeddings("hello")
|
||||
|
||||
get_registry().allow_endpoint(endpoint)
|
||||
restored = restore("oci-genai", model)
|
||||
assert isinstance(restored, OCIGenAIEmbeddings)
|
||||
assert restored.endpoints() == [endpoint]
|
||||
|
||||
|
||||
def test_restored_oci_default_endpoint_needs_no_allowlist():
|
||||
func = restore("oci-genai", {})
|
||||
assert isinstance(func, OCIGenAIEmbeddings)
|
||||
assert func.endpoints() == []
|
||||
|
||||
|
||||
def test_endpoint_allowlist_ignores_trailing_slash():
|
||||
get_registry().allow_endpoint("http://my-host:11434/")
|
||||
assert isinstance(
|
||||
|
||||
Reference in New Issue
Block a user