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
2026-11-03 16:08:15 +08:00
committed by GitHub
parent c16216e1c4
commit 8be1f8e00f
7 changed files with 833 additions and 0 deletions
+1
View File
@@ -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"
]