diff --git a/docs/src/embeddings/available_embedding_models/text_embedding_functions/oci_genai_embedding.md b/docs/src/embeddings/available_embedding_models/text_embedding_functions/oci_genai_embedding.md new file mode 100644 index 000000000..b2716a2fe --- /dev/null +++ b/docs/src/embeddings/available_embedding_models/text_embedding_functions/oci_genai_embedding.md @@ -0,0 +1,78 @@ +# OCI Generative AI Embeddings + +[Oracle Cloud Infrastructure (OCI) Generative AI](https://docs.oracle.com/en-us/iaas/Content/generative-ai/home.htm) hosts Cohere embedding models behind an OCI-native, IAM-authenticated API. + +Using the `oci-genai` embedding function requires the OCI Python SDK, which can be installed using `pip install oci`. No API key is stored in LanceDB: requests are signed by the SDK with one of the standard OCI authentication methods, so a compartment OCID and an identity that is allowed to `use generative-ai-family` in that compartment are all that is needed. + +Supported models are: + +- `cohere.embed-v4.0` (1536 dims by default; `output_dimensions` can request 256, 512, 1024 or 1536) +- `cohere.embed-english-v3.0`, `cohere.embed-multilingual-v3.0` (1024 dims) +- `cohere.embed-english-light-v3.0`, `cohere.embed-multilingual-light-v3.0` (384 dims) +- Any dedicated AI cluster endpoint OCID (`ocid1.generativeaiendpoint...`) hosting an embedding model + +Cohere embedding models are asymmetric: rows of the source column are embedded with `input_type="SEARCH_DOCUMENT"` and search queries with `input_type="SEARCH_QUERY"`. Inputs are sent in batches of at most 96 texts (the service limit), and the vector size is discovered with a single probe request unless `output_dimensions` is set. + +Supported parameters (to be passed in `create` method) are: + +| Parameter | Type | Default Value | Description | +|---|---|--------|---------| +| `name` | `str` | `"cohere.embed-v4.0"` | Model ID of an on-demand embedding model, or a dedicated AI cluster endpoint OCID. | +| `compartment_id` | `str` | `None` | OCID of the compartment the requests are authorized in and billed to. Required by the service; falls back to the `OCI_COMPARTMENT_ID` environment variable. | +| `region` | `str` | `"us-chicago-1"` | Region used to derive the inference endpoint `https://inference.generativeai..oci.oraclecloud.com`. | +| `service_endpoint` | `str` | `None` | Full inference endpoint. Overrides `region`; use it for realms with a different endpoint template. | +| `auth_type` | `str` | `"API_KEY"` | `API_KEY`, `SECURITY_TOKEN`, `INSTANCE_PRINCIPAL` or `RESOURCE_PRINCIPAL` (see below). | +| `auth_profile` | `str` | `"DEFAULT"` | Profile name in the OCI config file (`API_KEY` and `SECURITY_TOKEN` only). | +| `auth_file_location` | `str` | `"~/.oci/config"` | Path of the OCI config file (`API_KEY` and `SECURITY_TOKEN` only). | +| `truncate` | `str` | `"END"` | What to do with inputs longer than the model context: `NONE` (fail), `START` or `END` (drop tokens from that side). | +| `source_input_type` | `str` | `"SEARCH_DOCUMENT"` | `input_type` used when embedding the source column. | +| `query_input_type` | `str` | `"SEARCH_QUERY"` | `input_type` used when embedding queries. | +| `output_dimensions` | `int` | `None` | Embedding size for models that support it (`cohere.embed-v4.0`: 256, 512, 1024, 1536). | + +When a table is opened in a new process, an explicit `service_endpoint` or a non-default `region` restored from table metadata requires approval with `get_registry().allow_endpoint(...)`. The default `us-chicago-1` endpoint does not need approval. + +Authentication options: + +| `auth_type` | Identity used | +|---|---| +| `API_KEY` | User API key from profile `auth_profile` in `auth_file_location` (the setup produced by `oci setup config`). | +| `SECURITY_TOKEN` | Short-lived session token from the same profile, as written by `oci session authenticate`. | +| `INSTANCE_PRINCIPAL` | The Compute instance the code runs on; needs a dynamic group policy. | +| `RESOURCE_PRINCIPAL` | The OCI resource the code runs in (OKE workload identity, Functions, Data Science, ...). | + +Usage Example: + +```python + import lancedb + from lancedb.pydantic import LanceModel, Vector + from lancedb.embeddings import get_registry + + oci_genai = get_registry().get("oci-genai").create( + name="cohere.embed-v4.0", + compartment_id="ocid1.compartment.oc1..", + auth_type="API_KEY", + auth_profile="DEFAULT", + region="us-chicago-1", + ) + + class TextModel(LanceModel): + text: str = oci_genai.SourceField() + vector: Vector(oci_genai.ndims()) = oci_genai.VectorField() + + data = [{"text": "hello world"}, + {"text": "goodbye world"}] + + db = lancedb.connect("~/.lancedb") + tbl = db.create_table("test", schema=TextModel, mode="overwrite") + + tbl.add(data) + rs = tbl.search("hello").limit(1).to_pandas() +``` + +The compartment OCID can also be supplied at runtime instead of being stored in the table metadata, either through the `OCI_COMPARTMENT_ID` environment variable or with a registry variable: + +```python + registry = get_registry() + registry.set_var("oci_compartment", "ocid1.compartment.oc1..") + oci_genai = registry.get("oci-genai").create(compartment_id="$var:oci_compartment") +``` diff --git a/python/pyproject.toml b/python/pyproject.toml index 4a4e1d767..6294879f2 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -97,6 +97,7 @@ embeddings = [ "awscli>=1.44.38", "botocore>=1.31.57", 'ibm-watsonx-ai>=1.1.2; python_version >= "3.10"', + "oci>=2.165.1", "ollama>=0.3.0", "sentencepiece>=0.1.99" ] diff --git a/python/python/lancedb/embeddings/__init__.py b/python/python/lancedb/embeddings/__init__.py index 20e8c331b..679407610 100644 --- a/python/python/lancedb/embeddings/__init__.py +++ b/python/python/lancedb/embeddings/__init__.py @@ -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", diff --git a/python/python/lancedb/embeddings/oci_genai.py b/python/python/lancedb/embeddings/oci_genai.py new file mode 100644 index 000000000..382c3500d --- /dev/null +++ b/python/python/lancedb/embeddings/oci_genai.py @@ -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 `. + - Missing IAM policy: the caller needs + `allow group to use generative-ai-family in compartment `. + - 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 `. +""" + + +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..", + 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 diff --git a/python/python/tests/test_embeddings_oci_genai.py b/python/python/tests/test_embeddings_oci_genai.py new file mode 100644 index 000000000..c45598e82 --- /dev/null +++ b/python/python/tests/test_embeddings_oci_genai.py @@ -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", + } diff --git a/python/python/tests/test_embeddings_slow.py b/python/python/tests/test_embeddings_slow.py index 4de85070f..c2d26ec5e 100644 --- a/python/python/tests/test_embeddings_slow.py +++ b/python/python/tests/test_embeddings_slow.py @@ -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" diff --git a/python/python/tests/test_embeddings_trust.py b/python/python/tests/test_embeddings_trust.py index 92221c898..5eaed15a6 100644 --- a/python/python/tests/test_embeddings_trust.py +++ b/python/python/tests/test_embeddings_trust.py @@ -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(