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.

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"]

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):
     """