From e9b2f7e6e1f02bdae287359cff0ae99494b9b3bd Mon Sep 17 00:00:00 2001
From: alexcastgonz <82875882+pcbeingused333@users.noreply.github.com>
Date: Sat, 31 Oct 2026 20:49:49 -0400
Subject: [PATCH] fix(python): use dim as the embedding size for
 OpenAI-compatible models

Allow OpenAIEmbeddings.ndims() to use an explicitly configured dim for model
names outside the three built-in OpenAI names, so schemas can be built for
compatible models.

Preserve existing model defaults and request behavior. Add regression coverage
for the resulting Arrow schema dimensions and for all known OpenAI model
dimensions.
---
 docs/web/integrations/embedding/openai.mdx |  2 +-
 python/python/lancedb/embeddings/openai.py | 10 ++++-
 python/python/tests/test_embeddings.py     | 45 ++++++++++++++++++++++
 3 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/docs/web/integrations/embedding/openai.mdx b/docs/web/integrations/embedding/openai.mdx
index fbfa62d9d..4dd1ee8ba 100644
--- a/docs/web/integrations/embedding/openai.mdx
+++ b/docs/web/integrations/embedding/openai.mdx
@@ -10,7 +10,7 @@ LanceDB registers the OpenAI embeddings function in the registry by default, as
 | Parameter | Type | Default Value | Description |
 |---|---|---|---|
 | `name` | `str` | `"text-embedding-ada-002"` | The name of the model. |
-| `dim` | `int` |  Model default   | Requested output dimensions for models that support shortening embeddings. Use `func.ndims()` when defining the vector field. |
+| `dim` | `int` |  Model default   | Requested output dimensions for models that support shortening embeddings. Required for model names other than the OpenAI ones listed by `model_names()` (for example a model served behind an OpenAI-compatible `base_url`), where it is the model's embedding dimension. Use `func.ndims()` when defining the vector field. |
 | `use_azure` | bool | `False` | Set true to use Azure OpenAI SDK |
 
 Set your OpenAI API key before running the example:
diff --git a/python/python/lancedb/embeddings/openai.py b/python/python/lancedb/embeddings/openai.py
index d4ecbdf3f..8592598cc 100644
--- a/python/python/lancedb/embeddings/openai.py
+++ b/python/python/lancedb/embeddings/openai.py
@@ -36,6 +36,7 @@ class OpenAIEmbeddings(TextEmbeddingFunction):
     embedding_function = openai.create(
         name="<ollama-embedding-model-name>",
         base_url="http://localhost:11434",
+        dim=768,  # the model's embedding dimension; required for unlisted names
         )
     ```
 
@@ -77,8 +78,15 @@ class OpenAIEmbeddings(TextEmbeddingFunction):
             return self.dim or 3072
         elif self.name == "text-embedding-3-small":
             return self.dim or 1536
+        elif self.dim:
+            # A name not listed above, e.g. a model served behind an
+            # OpenAI-compatible `base_url`: `dim` is its embedding size
+            return self.dim
         else:
-            raise ValueError(f"Unknown model name {self.name}")
+            raise ValueError(
+                f"Unknown model name {self.name}: set `dim` to the model's "
+                "embedding dimension"
+            )
 
     def generate_embeddings(
         self, texts: Union[List[str], "np.ndarray"]
diff --git a/python/python/tests/test_embeddings.py b/python/python/tests/test_embeddings.py
index 32503c04f..c18ddd644 100644
--- a/python/python/tests/test_embeddings.py
+++ b/python/python/tests/test_embeddings.py
@@ -602,6 +602,51 @@ def test_openai_propagates_api_key(monkeypatch):
     assert len(actual.text) > 0
 
 
+@pytest.mark.parametrize("base_url", [None, "http://localhost:11434/v1"])
+def test_openai_compatible_model_uses_dim_for_ndims(base_url):
+    # Unlisted model names use the explicit dimension independently of base_url.
+    openai = get_registry().get("openai")
+    func = openai.create(
+        name="nomic-embed-text",
+        base_url=base_url,
+        dim=768,
+    )
+    assert func.ndims() == 768
+
+    class Schema(LanceModel):
+        text: str = func.SourceField()
+        vector: Vector(func.ndims()) = func.VectorField()
+
+    schema = Schema.to_arrow_schema()
+    assert schema.field("vector").type == pa.list_(pa.float32(), 768)
+
+
+def test_openai_unknown_model_without_dim_raises_clear_error():
+    openai = get_registry().get("openai")
+    func = openai.create(
+        name="nomic-embed-text",
+        base_url="http://localhost:11434/v1",
+    )
+    with pytest.raises(ValueError, match="dim"):
+        func.ndims()
+
+
+@pytest.mark.parametrize(
+    ("name", "dim", "expected_dim"),
+    [
+        ("text-embedding-ada-002", None, 1536),
+        ("text-embedding-ada-002", 512, 1536),
+        ("text-embedding-3-large", None, 3072),
+        ("text-embedding-3-large", 1024, 1024),
+        ("text-embedding-3-small", None, 1536),
+        ("text-embedding-3-small", 512, 512),
+    ],
+)
+def test_openai_known_model_ndims(name, dim, expected_dim):
+    func = get_registry().get("openai").create(name=name, dim=dim)
+    assert func.ndims() == expected_dim
+
+
 @patch("time.sleep")
 def test_openai_no_retry_on_401(mock_sleep):
     """