Add vector embedding UDF using Qwen3 model

Author: BohuTANG

.github/workflows/ci.yml

@@ -0,0 +1,26 @@
+name: CI
+
+on:
+  push:
+    branches: ["main"]
+  pull_request:
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+
+      - name: Install uv
+        run: pip install uv
+
+      - name: Run tests
+        env:
+          HF_HOME: ${{ github.workspace }}/.hf-cache
+          TRANSFORMERS_CACHE: ${{ github.workspace }}/.hf-cache
+        run: uv run pytest -m "not slow"

.gitignore

@@ -2,3 +2,4 @@
 __pycache__/
 .pytest_cache/
 *.egg-info/
+.hf-cache/

README.md

@@ -1,41 +1,32 @@
-# ai-server
+# databend-aiserver
 
-Stage aware Databend UDF server providing storage centric helpers used by the AI
-extensions. The initial release ships with:
+Databend AI Server fuses object storage and AI so everything—files, vectors,
+and structured rows—can be orchestrated inside Databend with standard SQL.
 
-- `list_stage_files`: enumerate objects inside an external S3 stage via
-  [Apache OpenDAL](https://opendal.apache.org/docs/python).
-- `read_pdf`: pull down a PDF object from a stage and return its extracted text
-  as a `STRING`.
-- `read_docx`: fetch Microsoft Word files (`.docx`) and expose their textual
-  content as a `STRING`.
+## UDFs (prefix `aiserver_`)
+- `list_stage_files(stage, limit)` – list objects from external stages.
+- `read_pdf(stage, path)` – return PDF text.
+- `read_docx(stage, path)` – return DOCX text.
+- `vector_embed_text_1024(model, text)` – generate 1024-dim embeddings (supports batches).
 
-## Getting started
+Embeddings currently accept alias `qwen` (`Qwen/Qwen3-Embedding-0.6B`). Models
+download to `.hf-cache/` and auto-select GPU/MPS/CPU.
 
+## Quickstart
 ```bash
-python3 -m venv .venv
-source .venv/bin/activate
-pip install -e .[dev]
-ai-udf-server --port 8815 --metrics-port 9091
+uv sync
+uv run databend-aiserver --port 8815
 ```
 
-Databend can now connect to the running Flight server and call the registered
-functions. `list_stage_files` expects a `STAGE_LOCATION` argument supplied by
-Databend plus a numeric limit. Example:
-
+## Sample SQL
 ```sql
-SELECT
-    list_stage_files(
-        @stage_location,
-        50
-    );
+SELECT aiserver_list_stage_files(@docs_stage, 50);
+SELECT aiserver_read_pdf(@docs_stage, 'reports/q1.pdf');
+SELECT aiserver_vector_embed_text_1024('qwen', doc_body) FROM docs_tbl;
 ```
 
-To retrieve file contents call the document readers with an explicit path. The
-PDF reader optionally accepts `NULL` for the `max_pages` argument to stream the
-entire document:
-
-```sql
-SELECT read_pdf(@stage_location, 'inbox/manual.pdf', NULL);
-SELECT read_docx(@stage_location, 'reports/summary.docx');
+## Tests
+```bash
+uv run pytest
+uv run pytest -m "not slow"
 ```

ai_server/__init__.py

@@ -0,0 +1,19 @@
+# Copyright 2025 Databend Labs
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""AI-enhanced Databend UDF server utilities."""
+
+from .server import create_server
+
+__all__ = ["create_server"]

ai_server/server.py

@@ -1,3 +1,17 @@
+# Copyright 2025 Databend Labs
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 """Command line entrypoint for the Databend AI UDF server."""
 
 from __future__ import annotations
@@ -11,7 +25,7 @@
 
 from prometheus_client import start_http_server as start_prometheus_server
 
-from ai_server.server import create_server
+from databend_aiserver.server import create_server
 
 logger = logging.getLogger(__name__)
 

ai_server/udfs/__init__.py

@@ -0,0 +1,56 @@
+# Copyright 2025 Databend Labs
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Entrypoint for the Databend AI UDF server."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+from databend_udf import UDFServer
+
+from databend_aiserver.udfs import (
+    aiserver_list_stage_files,
+    aiserver_read_docx,
+    aiserver_read_pdf,
+    aiserver_vector_embed_text_1024,
+)
+
+
+def create_server(
+    host: str = "0.0.0.0", port: int = 8815, metric_port: Optional[int] = None
+) -> UDFServer:
+    """
+    Create a configured UDF server instance.
+
+    Parameters
+    ----------
+    host:
+        Bind address for the Flight server.
+    port:
+        Bind port for the Flight server.
+    metric_port:
+        Optional metrics port for Prometheus exporter. When provided the
+        databend-udf server will expose metrics via Prometheus.
+    """
+    location = f"{host}:{port}"
+    metric_location = (
+        f"{host}:{metric_port}" if metric_port is not None else None
+    )
+    server = UDFServer(location, metric_location=metric_location)
+    server.add_function(aiserver_list_stage_files)
+    server.add_function(aiserver_read_pdf)
+    server.add_function(aiserver_read_docx)
+    server.add_function(aiserver_vector_embed_text_1024)
+    return server

databend_aiserver/__init__.py

@@ -1,3 +1,17 @@
+# Copyright 2025 Databend Labs
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 """Helpers to translate Databend stage metadata into OpenDAL operators."""
 
 from __future__ import annotations