diff --git a/README.md b/README.md index 9e282b3e10..9ea5ee52e9 100644 --- a/README.md +++ b/README.md @@ -89,6 +89,10 @@ That's it! - Login with username: `unstract` password: `unstract` - Start extracting data! +The Compose quickstart also provisions local Qdrant, pgvector/PostgreSQL, +Weaviate, and Milvus services. See [local vector database connection +settings](docker/README.md#local-vector-databases) for the adapter values. + ## 📦 Other Deployment Options ### Docker Compose diff --git a/docker/README.md b/docker/README.md index be71603c89..0e06aecbe9 100644 --- a/docker/README.md +++ b/docker/README.md @@ -30,6 +30,141 @@ VERSION=dev docker compose -f docker-compose.yaml --profile optional up -d Now access frontend at http://frontend.unstract.localhost +## Local vector databases + +The default development Compose stack starts the open-source vector database +backends supported by Unstract: + +- Qdrant +- PostgreSQL with the `pgvector` extension +- Weaviate +- Milvus Standalone (with private etcd and MinIO dependencies) + +Pinecone is supported as a hosted provider and is intentionally not included +in the local stack because it is not a self-hosted open-source service. + +The services use persistent named volumes and are bound to loopback on the host. +Qdrant and Weaviate use anonymous access because this is a local development +stack; do not expose these ports beyond the local machine without adding +authentication and TLS. The Unstract workers connect over the Compose network, +so use the internal addresses below when creating an adapter in the UI: + +| Adapter | UI fields | Address from Unstract containers | Host address | Local credentials | +|---------|-----------|----------------------------------|--------------|-------------------| +| Qdrant | URL, API Key | `http://qdrant:6333` | `http://localhost:6333` | Leave API Key empty | +| Postgres | Database, Host, Port, User, Password, Enable SSL | Host `postgres-vector`, port `5432` | Host `localhost`, port `5433` | Values from `docker/essentials.env`; set Enable SSL to `false` | +| Weaviate | URL, API Key | `http://weaviate:8080` | `http://localhost:8084` | Leave API Key empty | +| Milvus | URI, Token | `http://milvus:19530` | `http://localhost:19530` | Leave Token empty | + +For the Postgres adapter, use the `POSTGRES_USER`, `POSTGRES_PASSWORD`, and +`POSTGRES_DB` values in `docker/essentials.env`; the vector database has its own +container and volume even though it reuses the platform's local development +credentials. The initialization script enables `vector` and creates the +`POSTGRES_SCHEMA` schema on first boot. + +The default `run-platform.sh` flow creates `docker/.env` and +`docker/essentials.env` before starting Compose. To start only the vector +services during development, run: + +```bash +cd docker +VERSION=dev docker compose -f docker-compose.yaml up -d \ + qdrant postgres-vector weaviate milvus +``` + +Do not use `docker compose down -v` unless deleting local vector data is +intentional. + +## Local Qwen3 embeddings + +The Compose stack includes six optional Hugging Face Text Embeddings +Inference (TEI) services for the three selected Qwen3 embedding models. The +CPU and GPU services are separate so they can be benchmarked against the same +model and workload. Model files are cached in persistent, model-specific +volumes. + +| Model | Vector dimensions | CPU service / host port | GPU service / host port | +|-------|-------------------:|-------------------------|-------------------------| +| Qwen3-Embedding-0.6B | 1024 | qwen3-embedding-06b-cpu / 8101 | qwen3-embedding-06b-gpu / 8201 | +| Qwen3-Embedding-4B | 2560 | qwen3-embedding-4b-cpu / 8102 | qwen3-embedding-4b-gpu / 8202 | +| Qwen3-Embedding-8B | 4096 | qwen3-embedding-8b-cpu / 8103 | qwen3-embedding-8b-gpu / 8203 | + +Start the CPU, GPU, or complete comparison matrix from the docker directory: + +```bash +cd docker + +# CPU services +VERSION=dev docker compose -f docker-compose.yaml --profile embeddings-cpu up -d \ + qwen3-embedding-06b-cpu qwen3-embedding-4b-cpu qwen3-embedding-8b-cpu + +# GPU services +VERSION=dev docker compose -f docker-compose.yaml --profile embeddings-gpu up -d \ + qwen3-embedding-06b-gpu qwen3-embedding-4b-gpu qwen3-embedding-8b-gpu + +# All six services +VERSION=dev docker compose -f docker-compose.yaml --profile embeddings-both up -d \ + qwen3-embedding-06b-cpu qwen3-embedding-4b-cpu qwen3-embedding-8b-cpu \ + qwen3-embedding-06b-gpu qwen3-embedding-4b-gpu qwen3-embedding-8b-gpu +``` + +GPU services require the NVIDIA Container Toolkit and a compatible NVIDIA +driver. The default CUDA image targets the TEI CUDA 1.9 runtime; set +QWEN3_TEI_GPU_IMAGE when an architecture-specific image is needed. +The default CPU image targets x86_64; set QWEN3_TEI_CPU_IMAGE to the TEI +cpu-arm64-1.9 image on ARM64 hosts. The embeddings-both profile starts every +service in the matrix and may exceed available GPU memory if all six are +launched together. For a fair comparison, start one profile at a time or +benchmark endpoints sequentially. + +Each service exposes an OpenAI-compatible endpoint. From an Unstract +container, use the internal URL; from the host, use the localhost URL: + +| Service family | Internal API base | Host API base | +|----------------|-------------------|---------------| +| 0.6B CPU | http://qwen3-embedding-06b-cpu/v1 | http://localhost:8101/v1 | +| 0.6B GPU | http://qwen3-embedding-06b-gpu/v1 | http://localhost:8201/v1 | +| 4B CPU | http://qwen3-embedding-4b-cpu/v1 | http://localhost:8102/v1 | +| 4B GPU | http://qwen3-embedding-4b-gpu/v1 | http://localhost:8202/v1 | +| 8B CPU | http://qwen3-embedding-8b-cpu/v1 | http://localhost:8103/v1 | +| 8B GPU | http://qwen3-embedding-8b-gpu/v1 | http://localhost:8203/v1 | + +For the OpenAI Compatible Embedding adapter, set Model to the matching +served-model alias, API Base to the appropriate URL above, and API Key to the +local placeholder accepted by the adapter. Qwen3 retrieval expects this query +prefix: + +```text +Instruct: Given a web search query, retrieve relevant passages that answer the query +Query: +``` + +Leave Passage Prefix empty. Append one ASCII space after Query: before the +query text. These fields are available in the adapter schema and are also +recorded with model IDs, dimensions, and endpoints in +docker/local-embeddings.config.json. Use one model consistently for indexing +and querying. A model with different vector dimensions requires a separate +collection and a full reindex. + +Run the benchmark harness after starting the desired services: + +```bash +python3 docker/benchmarks/embedding_benchmark.py --mode cpu \ + --output /tmp/qwen3-cpu.json --strict +python3 docker/benchmarks/embedding_benchmark.py --mode gpu \ + --output /tmp/qwen3-gpu.json --strict +python3 docker/benchmarks/embedding_benchmark.py --mode both \ + --output /tmp/qwen3-both.json --strict +``` + +The harness waits for health, warms each endpoint, measures single-request +and batch latency/throughput, checks the expected dimensions, and reports +fixture retrieval metrics. Replace +docker/benchmarks/qwen3_smoke_dataset.json with a representative corpus and +labeled queries before using quality scores to select a production model. +This change intentionally retains all six variants; unselected services can +be removed in a follow-up after the benchmark review. + ## Overriding a service's config By making use of the [merge compose files](https://docs.docker.com/compose/how-tos/multiple-compose-files/merge/) feature its possible to override some configuration that's used by the services. diff --git a/docker/benchmarks/embedding_benchmark.py b/docker/benchmarks/embedding_benchmark.py new file mode 100644 index 0000000000..ac0a9d228d --- /dev/null +++ b/docker/benchmarks/embedding_benchmark.py @@ -0,0 +1,434 @@ +#!/usr/bin/env python3 +"""Benchmark local OpenAI-compatible Qwen3 embedding endpoints.""" + +from __future__ import annotations + +import argparse +import json +import math +import statistics +import sys +import time +from datetime import UTC, datetime +from pathlib import Path +from typing import Any +from urllib.error import HTTPError, URLError +from urllib.request import Request, urlopen + +DEFAULT_CONFIG = Path(__file__).resolve().parents[1] / "local-embeddings.config.json" +DEFAULT_DATASET = Path(__file__).resolve().parent / "qwen3_smoke_dataset.json" +DEFAULT_TIMEOUT = 300.0 +DEFAULT_HEALTH_TIMEOUT = 600.0 +MODES = ("cpu", "gpu") +QUALITY_K_VALUES = (1, 3, 5) + + +def load_json(path: Path) -> dict[str, Any]: + """Load a JSON object from disk.""" + with path.open(encoding="utf-8") as file: + value = json.load(file) + if not isinstance(value, dict): + raise ValueError(f"{path} must contain a JSON object") + return value + + +def percentile(values: list[float], quantile: float) -> float: + """Return a nearest-rank percentile in milliseconds.""" + if not values: + return 0.0 + ordered = sorted(values) + index = max(0, math.ceil(quantile * len(ordered)) - 1) + return ordered[index] + + +def request_json( + url: str, + payload: dict[str, Any] | None, + timeout: float, +) -> dict[str, Any]: + """Send a JSON request and return its JSON object response.""" + body = None if payload is None else json.dumps(payload).encode("utf-8") + request = Request( + url, + data=body, + headers={"Content-Type": "application/json"} if body is not None else {}, + method="POST" if body is not None else "GET", + ) + try: + with urlopen(request, timeout=timeout) as response: # noqa: S310 + value = json.loads(response.read().decode("utf-8")) + except HTTPError as error: + detail = error.read().decode("utf-8", errors="replace")[:500] + raise RuntimeError(f"{error.code} from {url}: {detail}") from error + except URLError as error: + raise RuntimeError(f"request to {url} failed: {error.reason}") from error + if not isinstance(value, dict): + raise ValueError(f"{url} returned a non-object JSON response") + return value + + +def service_root(base_url: str) -> str: + """Convert an OpenAI base URL to the service root.""" + root = base_url.rstrip("/") + if root.endswith("/v1"): + return root[:-3].rstrip("/") + return root + + +def wait_for_health(base_url: str, timeout: float) -> None: + """Wait until the embedding service reports healthy.""" + health_url = f"{service_root(base_url)}/health" + deadline = time.monotonic() + timeout + last_error = "no response" + while time.monotonic() < deadline: + try: + with urlopen( # noqa: S310 + Request(health_url, method="GET"), timeout=10.0 + ) as response: + if response.status == 200: + return + last_error = f"HTTP {response.status}" + except HTTPError as error: + last_error = f"HTTP {error.code}" + except (OSError, URLError) as error: + last_error = str(error) + time.sleep(5.0) + raise TimeoutError(f"{health_url} did not become healthy: {last_error}") + + +def embed( + base_url: str, + model_id: str, + texts: list[str], + timeout: float, +) -> list[list[float]]: + """Generate embeddings through the OpenAI-compatible endpoint.""" + response = request_json( + f"{base_url.rstrip('/')}/embeddings", + { + "model": model_id, + "input": texts, + "encoding_format": "float", + }, + timeout, + ) + data = response.get("data") + if not isinstance(data, list) or len(data) != len(texts): + raise ValueError( + f"expected {len(texts)} embeddings from {base_url}, got {data!r}" + ) + vectors: list[list[float]] = [] + for item in data: + if not isinstance(item, dict) or not isinstance(item.get("embedding"), list): + raise ValueError(f"invalid embedding response from {base_url}") + vectors.append(item["embedding"]) + return vectors + + +def embed_in_batches( + base_url: str, + model_id: str, + texts: list[str], + batch_size: int, + timeout: float, +) -> list[list[float]]: + """Generate embeddings in bounded requests for backends with batch caps.""" + vectors: list[list[float]] = [] + for start in range(0, len(texts), batch_size): + vectors.extend( + embed(base_url, model_id, texts[start : start + batch_size], timeout) + ) + return vectors + + +def prefixed(text: str, prefix: str) -> str: + """Apply a configured query/document prefix.""" + return f"{prefix}{text}" if prefix else text + + +def cosine(left: list[float], right: list[float]) -> float: + """Calculate cosine similarity without third-party dependencies.""" + left_norm = math.sqrt(sum(value * value for value in left)) + right_norm = math.sqrt(sum(value * value for value in right)) + if not left_norm or not right_norm: + return 0.0 + return sum(a * b for a, b in zip(left, right, strict=True)) / (left_norm * right_norm) + + +def quality_metrics( + query_vectors: list[list[float]], + document_vectors: list[list[float]], + queries: list[dict[str, Any]], + documents: list[dict[str, Any]], +) -> dict[str, float]: + """Calculate retrieval metrics for the supplied labeled fixture.""" + document_ids = [str(document["id"]) for document in documents] + reciprocal_ranks: list[float] = [] + recalls = {k: [] for k in QUALITY_K_VALUES} + ndcgs = {k: [] for k in QUALITY_K_VALUES} + + for query_vector, query in zip(query_vectors, queries, strict=True): + relevant = {str(item) for item in query["relevant"]} + ranked = sorted( + ( + (cosine(query_vector, document_vector), document_id) + for document_id, document_vector in zip( + document_ids, document_vectors, strict=True + ) + ), + reverse=True, + ) + ranked_ids = [document_id for _, document_id in ranked] + positions = [ + index + 1 + for index, document_id in enumerate(ranked_ids) + if document_id in relevant + ] + reciprocal_ranks.append(1.0 / positions[0] if positions else 0.0) + + for k in QUALITY_K_VALUES: + top_ids = ranked_ids[:k] + recalls[k].append(len(set(top_ids) & relevant) / max(1, len(relevant))) + dcg = sum( + 1.0 / math.log2(index + 2) + for index, document_id in enumerate(top_ids) + if document_id in relevant + ) + ideal_hits = min(k, len(relevant)) + ideal_dcg = sum(1.0 / math.log2(index + 2) for index in range(ideal_hits)) + ndcgs[k].append(dcg / ideal_dcg if ideal_dcg else 0.0) + + result = {"mrr": statistics.mean(reciprocal_ranks)} + result.update( + {f"recall_at_{k}": statistics.mean(values) for k, values in recalls.items()} + ) + result.update( + {f"ndcg_at_{k}": statistics.mean(values) for k, values in ndcgs.items()} + ) + return {key: round(value, 6) for key, value in result.items()} + + +def validate_vectors(vectors: list[list[float]], expected_dimension: int) -> None: + """Ensure an endpoint returns the configured vector dimension.""" + actual_dimensions = {len(vector) for vector in vectors} + if actual_dimensions != {expected_dimension}: + raise ValueError( + f"expected {expected_dimension} dimensions, got {sorted(actual_dimensions)}" + ) + + +def benchmark_endpoint( + model: dict[str, Any], + mode: str, + config: dict[str, Any], + dataset: dict[str, Any], + repetitions: int, + batch_size: int, + timeout: float, + health_timeout: float, + skip_quality: bool, +) -> dict[str, Any]: + """Benchmark one model/mode endpoint.""" + endpoint = model[mode] + base_url = str(endpoint["base_url"]) + model_id = str(model["id"]) + expected_dimension = int(model["dimension"]) + documents = dataset["documents"] + queries = dataset["queries"] + query_prefix = str(config.get("query_prefix", "")) + passage_prefix = str(config.get("passage_prefix", "")) + query_text = prefixed(str(queries[0]["text"]), query_prefix) + passage_texts = [ + prefixed(str(document["text"]), passage_prefix) + for document in documents[:batch_size] + ] + if len(passage_texts) < batch_size: + passage_texts *= math.ceil(batch_size / len(passage_texts)) + passage_texts = passage_texts[:batch_size] + + wait_for_health(base_url, health_timeout) + for _ in range(2): + validate_vectors( + embed(base_url, model_id, [query_text], timeout), expected_dimension + ) + + single_latencies: list[float] = [] + for _ in range(repetitions): + started = time.perf_counter() + validate_vectors( + embed(base_url, model_id, [query_text], timeout), expected_dimension + ) + single_latencies.append((time.perf_counter() - started) * 1000) + + batch_latencies: list[float] = [] + for _ in range(repetitions): + started = time.perf_counter() + vectors = embed(base_url, model_id, passage_texts, timeout) + validate_vectors(vectors, expected_dimension) + batch_latencies.append((time.perf_counter() - started) * 1000) + + result: dict[str, Any] = { + "model": model_id, + "repository": model["repository"], + "revision": model.get("revision"), + "mode": mode, + "service": endpoint["service"], + "base_url": base_url, + "dimension": expected_dimension, + "single_latency_ms": { + "p50": round(percentile(single_latencies, 0.50), 3), + "p95": round(percentile(single_latencies, 0.95), 3), + }, + "batch_latency_ms": { + "p50": round(percentile(batch_latencies, 0.50), 3), + "p95": round(percentile(batch_latencies, 0.95), 3), + }, + "batch_size": batch_size, + "batch_throughput_items_per_second": round( + batch_size / (statistics.mean(batch_latencies) / 1000), 3 + ), + } + + if not skip_quality: + document_vectors = embed_in_batches( + base_url, + model_id, + [prefixed(str(document["text"]), passage_prefix) for document in documents], + batch_size, + timeout, + ) + query_vectors = embed_in_batches( + base_url, + model_id, + [prefixed(str(query["text"]), query_prefix) for query in queries], + batch_size, + timeout, + ) + validate_vectors(document_vectors, expected_dimension) + validate_vectors(query_vectors, expected_dimension) + result["quality"] = quality_metrics( + query_vectors, document_vectors, queries, documents + ) + + return result + + +def parse_args() -> argparse.Namespace: + """Parse command-line arguments.""" + parser = argparse.ArgumentParser( + description="Benchmark the local Qwen3 CPU/GPU embedding matrix." + ) + parser.add_argument( + "--mode", + choices=("cpu", "gpu", "both"), + default="both", + help="Endpoint variation to test; both tests CPU and GPU endpoints.", + ) + parser.add_argument( + "--config", + type=Path, + default=DEFAULT_CONFIG, + help=f"Endpoint config JSON (default: {DEFAULT_CONFIG})", + ) + parser.add_argument( + "--dataset", + type=Path, + default=DEFAULT_DATASET, + help=f"Labeled benchmark fixture (default: {DEFAULT_DATASET})", + ) + parser.add_argument("--repetitions", type=int, default=5) + parser.add_argument("--batch-size", type=int, default=8) + parser.add_argument("--timeout", type=float, default=DEFAULT_TIMEOUT) + parser.add_argument("--health-timeout", type=float, default=DEFAULT_HEALTH_TIMEOUT) + parser.add_argument("--skip-quality", action="store_true") + parser.add_argument("--strict", action="store_true") + parser.add_argument("--output", type=Path) + return parser.parse_args() + + +def main() -> int: + """Run the selected benchmark matrix.""" + args = parse_args() + if args.repetitions < 1: + raise SystemExit("--repetitions must be at least 1") + if args.batch_size < 1: + raise SystemExit("--batch-size must be at least 1") + + config = load_json(args.config) + dataset = load_json(args.dataset) + modes = MODES if args.mode == "both" else (args.mode,) + results: list[dict[str, Any]] = [] + models = config.get("models") + if not isinstance(models, list) or not models: + raise SystemExit("config must define at least one model") + + for model in models: + for mode in modes: + print(f"Benchmarking {model['id']} ({mode})...", flush=True) + try: + results.append( + benchmark_endpoint( + model, + mode, + config, + dataset, + args.repetitions, + args.batch_size, + args.timeout, + args.health_timeout, + args.skip_quality, + ) + ) + except Exception as error: # noqa: BLE001 - continue the matrix + results.append( + { + "model": model["id"], + "repository": model["repository"], + "revision": model.get("revision"), + "mode": mode, + "service": model[mode]["service"], + "base_url": model[mode]["base_url"], + "status": "error", + "error": str(error), + } + ) + print(f" ERROR: {error}", file=sys.stderr) + + report = { + "generated_at_utc": datetime.now(UTC).isoformat(), + "mode": args.mode, + "config": str(args.config), + "dataset": str(args.dataset), + "settings": { + "repetitions": args.repetitions, + "batch_size": args.batch_size, + "timeout_seconds": args.timeout, + "health_timeout_seconds": args.health_timeout, + "quality_fixture": not args.skip_quality, + }, + "results": results, + } + if args.output: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8") + print(f"Report written to {args.output}") + + print("\nSummary") + print("mode model status dimension single_p50_ms batch_items_per_second mrr") + for result in results: + quality = result.get("quality", {}) + print( + f"{result['mode']} {result['model']} " + f"{result.get('status', 'ok')} " + f"{result.get('dimension', '-')} " + f"{result.get('single_latency_ms', {}).get('p50', '-')} " + f"{result.get('batch_throughput_items_per_second', '-')} " + f"{quality.get('mrr', '-')}" + ) + + has_errors = any(result.get("status") == "error" for result in results) + return 1 if args.strict and has_errors else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/docker/benchmarks/qwen3_smoke_dataset.json b/docker/benchmarks/qwen3_smoke_dataset.json new file mode 100644 index 0000000000..700bc7b1c9 --- /dev/null +++ b/docker/benchmarks/qwen3_smoke_dataset.json @@ -0,0 +1,76 @@ +{ + "documents": [ + { + "id": "qdrant", + "text": "Qdrant is an open-source vector database used to store and search dense embeddings. In the local Unstract stack it listens on port 6333." + }, + { + "id": "pgvector", + "text": "pgvector adds vector similarity search to PostgreSQL. Unstract's local PostgreSQL vector adapter uses a separate pgvector-enabled service." + }, + { + "id": "weaviate", + "text": "Weaviate is a vector database with an HTTP API. The local Unstract development stack runs it with anonymous access on a loopback port." + }, + { + "id": "milvus", + "text": "Milvus is a vector database for similarity search. A standalone local deployment uses Milvus together with etcd metadata and MinIO object storage." + }, + { + "id": "openai-compatible", + "text": "An OpenAI-compatible embedding endpoint accepts a model name and one or more input texts, then returns vectors in the standard embeddings response shape." + }, + { + "id": "qwen3", + "text": "Qwen3 embedding models are multilingual text encoders. For retrieval, Qwen recommends an instruction and Query prefix for queries, while passages are encoded without that query instruction." + }, + { + "id": "vector-dimension", + "text": "Every vector collection must use the embedding dimension produced by its model. Changing from one embedding model to another normally requires a separate collection and a full reindex." + }, + { + "id": "cpu", + "text": "CPU inference does not require an NVIDIA device, but it generally has lower embedding throughput and higher latency than GPU inference for the same model." + }, + { + "id": "gpu", + "text": "GPU inference requires an NVIDIA driver, the NVIDIA Container Toolkit, and enough device memory for the model and its inference batches." + }, + { + "id": "benchmark", + "text": "A useful embedding benchmark records vector dimension, warm and cold startup behavior, single-request latency, batch throughput, and retrieval quality on representative documents and queries." + } + ], + "queries": [ + { + "id": "which-service-stores-vectors", + "text": "Which local service stores and searches dense embeddings?", + "relevant": ["qdrant", "pgvector", "weaviate", "milvus"] + }, + { + "id": "which-service-uses-postgres", + "text": "Which vector option adds similarity search to PostgreSQL?", + "relevant": ["pgvector"] + }, + { + "id": "how-connect-embedding", + "text": "What API shape should a local embedding server expose to connect to Unstract?", + "relevant": ["openai-compatible"] + }, + { + "id": "how-handle-qwen-query", + "text": "How should a Qwen3 retrieval query be formatted?", + "relevant": ["qwen3"] + }, + { + "id": "model-switch", + "text": "What must happen when the embedding model or vector dimension changes?", + "relevant": ["vector-dimension"] + }, + { + "id": "compare-hardware", + "text": "What should be compared between CPU and GPU embedding runs?", + "relevant": ["cpu", "gpu", "benchmark"] + } + ] +} diff --git a/docker/docker-compose-dev-essentials.yaml b/docker/docker-compose-dev-essentials.yaml index 605781be15..bc7c71fd2d 100644 --- a/docker/docker-compose-dev-essentials.yaml +++ b/docker/docker-compose-dev-essentials.yaml @@ -124,13 +124,144 @@ services: container_name: unstract-vector-db restart: unless-stopped ports: - - "6333:6333" + - "127.0.0.1:6333:6333" volumes: - qdrant_data:/var/lib/qdrant/data/ labels: - traefik.enable=false + + # Separate pgvector instance for the Postgres adapter. The platform's `db` + # service is also pgvector-enabled, but keeping this volume/container + # independent prevents adapter connection tests from sharing platform data. + postgres-vector: + image: "pgvector/pgvector:pg15" + container_name: unstract-postgres-vector + restart: unless-stopped + shm_size: 128mb + ports: + - "127.0.0.1:5433:5432" + volumes: + - postgres_vector_data:/var/lib/postgresql/data/ + - ./scripts/db-setup/vector_db_setup.sh:/docker-entrypoint-initdb.d/vector_db_setup.sh:ro env_file: + # Reuse the local platform database credentials from essentials.env while + # retaining a separate data directory and container. - ./essentials.env + healthcheck: + test: ["CMD-SHELL", "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}"] + interval: 10s + timeout: 5s + retries: 5 + labels: + - traefik.enable=false + + # Weaviate uses anonymous access for this loopback-only development service. + # The adapter uses the local connector for this HTTP endpoint; no cloud + # account or API key is required. + weaviate: + image: "docker.io/semitechnologies/weaviate:1.39.2" + container_name: unstract-weaviate + restart: unless-stopped + command: + - --host + - 0.0.0.0 + - --port + - "8080" + - --scheme + - http + ports: + - "127.0.0.1:8084:8080" + - "127.0.0.1:50051:50051" + volumes: + - weaviate_data:/var/lib/weaviate + environment: + QUERY_DEFAULTS_LIMIT: "25" + AUTHENTICATION_ANONYMOUS_ACCESS_ENABLED: "true" + PERSISTENCE_DATA_PATH: "/var/lib/weaviate" + DEFAULT_VECTORIZER_MODULE: "none" + CLUSTER_HOSTNAME: "node1" + labels: + - traefik.enable=false + + # Milvus standalone needs etcd for metadata and MinIO for object storage. + # Keep those dependencies private to this Compose network; only Milvus's + # client and WebUI ports are published for local development. + milvus-etcd: + image: "quay.io/coreos/etcd:v3.5.18" + container_name: unstract-milvus-etcd + restart: unless-stopped + environment: + ETCD_AUTO_COMPACTION_MODE: revision + ETCD_AUTO_COMPACTION_RETENTION: "1000" + ETCD_QUOTA_BACKEND_BYTES: "4294967296" + ETCD_SNAPSHOT_COUNT: "50000" + volumes: + - milvus_etcd_data:/etcd + command: >- + etcd -advertise-client-urls=http://milvus-etcd:2379 + -listen-client-urls http://0.0.0.0:2379 --data-dir /etcd + healthcheck: + test: ["CMD", "etcdctl", "endpoint", "health"] + interval: 30s + timeout: 20s + retries: 3 + labels: + - traefik.enable=false + + milvus-minio: + image: "minio/minio:RELEASE.2024-05-28T17-19-04Z" + container_name: unstract-milvus-minio + restart: unless-stopped + environment: + MINIO_ROOT_USER: minioadmin + MINIO_ROOT_PASSWORD: minioadmin + # Keep the legacy names aligned with the credentials expected by the + # Milvus 2.5 standalone configuration. + MINIO_ACCESS_KEY: minioadmin + MINIO_SECRET_KEY: minioadmin + volumes: + - milvus_minio_data:/minio_data + command: minio server /minio_data --console-address ":9001" + healthcheck: + test: ["CMD", "curl", "-f", "http://localhost:9000/minio/health/live"] + interval: 30s + timeout: 20s + retries: 3 + labels: + - traefik.enable=false + + milvus: + # Keep this aligned with the pymilvus 2.5.x dependency used by sdk1. + image: "milvusdb/milvus:v2.5.17" + container_name: unstract-milvus + restart: unless-stopped + command: ["milvus", "run", "standalone"] + security_opt: + - seccomp:unconfined + environment: + MINIO_REGION: us-east-1 + ETCD_ENDPOINTS: milvus-etcd:2379 + MINIO_ADDRESS: milvus-minio:9000 + MINIO_ACCESS_KEY: minioadmin + MINIO_SECRET_KEY: minioadmin + volumes: + - milvus_data:/var/lib/milvus + healthcheck: + test: ["CMD", "curl", "-f", "http://localhost:9091/healthz"] + interval: 30s + start_period: 90s + timeout: 20s + retries: 3 + ports: + - "127.0.0.1:19530:19530" + - "127.0.0.1:9091:9091" + depends_on: + milvus-etcd: + condition: service_healthy + milvus-minio: + condition: service_healthy + labels: + - traefik.enable=false rabbitmq: image: rabbitmq:4.1.0-management @@ -149,7 +280,12 @@ volumes: flipt_data: minio_data: postgres_data: + postgres_vector_data: qdrant_data: redis_data: prompt_studio_data: rabbitmq_data: + weaviate_data: + milvus_etcd_data: + milvus_minio_data: + milvus_data: diff --git a/docker/docker-compose-local-embeddings.yaml b/docker/docker-compose-local-embeddings.yaml new file mode 100644 index 0000000000..f870d0e10d --- /dev/null +++ b/docker/docker-compose-local-embeddings.yaml @@ -0,0 +1,253 @@ +# Optional local Qwen3 embedding services. +# +# The services are profile-gated so the normal development stack does not +# download multi-gigabyte model weights. CPU and GPU services deliberately have +# separate names, ports, caches, and containers so their benchmark results can +# be compared without changing an endpoint in place. + +x-qwen3-healthcheck: &qwen3_healthcheck + test: + - CMD-SHELL + - >- + curl --fail --silent http://localhost:80/health >/dev/null || + exit 1 + interval: 15s + timeout: 10s + retries: 40 + start_period: 10m + +x-qwen3-common: &qwen3_common + restart: unless-stopped + shm_size: 1gb + expose: + - "80" + environment: + HF_HOME: /data + HF_HUB_DISABLE_TELEMETRY: "1" + DO_NOT_TRACK: "1" + HF_TOKEN: ${HF_TOKEN:-} + TOKENIZERS_PARALLELISM: "false" + labels: + traefik.enable: "false" + stop_grace_period: 30s + healthcheck: + <<: *qwen3_healthcheck + +x-qwen3-cpu-image: &qwen3_cpu_image + ${QWEN3_TEI_CPU_IMAGE:-ghcr.io/huggingface/text-embeddings-inference:cpu-1.9} + +x-qwen3-gpu-image: &qwen3_gpu_image + ${QWEN3_TEI_GPU_IMAGE:-ghcr.io/huggingface/text-embeddings-inference:cuda-1.9} + +services: + qwen3-embedding-06b-cpu: + <<: *qwen3_common + image: *qwen3_cpu_image + container_name: unstract-qwen3-embedding-06b-cpu + profiles: + - embeddings-cpu + - embeddings-both + ports: + - "127.0.0.1:${QWEN3_06B_CPU_PORT:-8101}:80" + volumes: + - qwen3_embedding_06b_cpu_cache:/data + command: + - --model-id + - Qwen/Qwen3-Embedding-0.6B + - --revision + - ${QWEN3_06B_REVISION:-97b0c614be4d77ee51c0cef4e5f07c00f9eb65b3} + - --served-model-name + - qwen3-embedding-06b + - --pooling + - last-token + - --max-batch-tokens + - ${QWEN3_MAX_BATCH_TOKENS:-8192} + - --max-client-batch-size + - ${QWEN3_MAX_CLIENT_BATCH_SIZE:-32} + - --max-concurrent-requests + - ${QWEN3_MAX_CONCURRENT_REQUESTS:-4} + - --tokenization-workers + - ${QWEN3_TOKENIZATION_WORKERS:-4} + + qwen3-embedding-4b-cpu: + <<: *qwen3_common + image: *qwen3_cpu_image + container_name: unstract-qwen3-embedding-4b-cpu + profiles: + - embeddings-cpu + - embeddings-both + ports: + - "127.0.0.1:${QWEN3_4B_CPU_PORT:-8102}:80" + volumes: + - qwen3_embedding_4b_cpu_cache:/data + command: + - --model-id + - Qwen/Qwen3-Embedding-4B + - --revision + - ${QWEN3_4B_REVISION:-5cf2132abc99cad020ac570b19d031efec650f2b} + - --served-model-name + - qwen3-embedding-4b + - --pooling + - last-token + - --max-batch-tokens + - ${QWEN3_MAX_BATCH_TOKENS:-8192} + - --max-client-batch-size + - ${QWEN3_MAX_CLIENT_BATCH_SIZE:-32} + - --max-concurrent-requests + - ${QWEN3_MAX_CONCURRENT_REQUESTS:-4} + - --tokenization-workers + - ${QWEN3_TOKENIZATION_WORKERS:-4} + + qwen3-embedding-8b-cpu: + <<: *qwen3_common + image: *qwen3_cpu_image + container_name: unstract-qwen3-embedding-8b-cpu + profiles: + - embeddings-cpu + - embeddings-both + ports: + - "127.0.0.1:${QWEN3_8B_CPU_PORT:-8103}:80" + volumes: + - qwen3_embedding_8b_cpu_cache:/data + command: + - --model-id + - Qwen/Qwen3-Embedding-8B + - --revision + - ${QWEN3_8B_REVISION:-1d8ad4ca9b3dd8059ad90a75d4983776a23d44af} + - --served-model-name + - qwen3-embedding-8b + - --pooling + - last-token + - --max-batch-tokens + - ${QWEN3_MAX_BATCH_TOKENS:-8192} + - --max-client-batch-size + - ${QWEN3_MAX_CLIENT_BATCH_SIZE:-32} + - --max-concurrent-requests + - ${QWEN3_MAX_CONCURRENT_REQUESTS:-4} + - --tokenization-workers + - ${QWEN3_TOKENIZATION_WORKERS:-4} + + qwen3-embedding-06b-gpu: + <<: *qwen3_common + image: *qwen3_gpu_image + container_name: unstract-qwen3-embedding-06b-gpu + profiles: + - embeddings-gpu + - embeddings-both + ports: + - "127.0.0.1:${QWEN3_06B_GPU_PORT:-8201}:80" + volumes: + - qwen3_embedding_06b_gpu_cache:/data + deploy: + resources: + reservations: + devices: + - driver: nvidia + count: 1 + capabilities: + - gpu + command: + - --model-id + - Qwen/Qwen3-Embedding-0.6B + - --revision + - ${QWEN3_06B_REVISION:-97b0c614be4d77ee51c0cef4e5f07c00f9eb65b3} + - --served-model-name + - qwen3-embedding-06b + - --dtype + - float16 + - --pooling + - last-token + - --max-batch-tokens + - ${QWEN3_MAX_BATCH_TOKENS:-8192} + - --max-client-batch-size + - ${QWEN3_MAX_CLIENT_BATCH_SIZE:-32} + - --max-concurrent-requests + - ${QWEN3_MAX_CONCURRENT_REQUESTS:-4} + - --tokenization-workers + - ${QWEN3_TOKENIZATION_WORKERS:-4} + + qwen3-embedding-4b-gpu: + <<: *qwen3_common + image: *qwen3_gpu_image + container_name: unstract-qwen3-embedding-4b-gpu + profiles: + - embeddings-gpu + - embeddings-both + ports: + - "127.0.0.1:${QWEN3_4B_GPU_PORT:-8202}:80" + volumes: + - qwen3_embedding_4b_gpu_cache:/data + deploy: + resources: + reservations: + devices: + - driver: nvidia + count: 1 + capabilities: + - gpu + command: + - --model-id + - Qwen/Qwen3-Embedding-4B + - --revision + - ${QWEN3_4B_REVISION:-5cf2132abc99cad020ac570b19d031efec650f2b} + - --served-model-name + - qwen3-embedding-4b + - --dtype + - float16 + - --pooling + - last-token + - --max-batch-tokens + - ${QWEN3_MAX_BATCH_TOKENS:-8192} + - --max-client-batch-size + - ${QWEN3_MAX_CLIENT_BATCH_SIZE:-32} + - --max-concurrent-requests + - ${QWEN3_MAX_CONCURRENT_REQUESTS:-4} + - --tokenization-workers + - ${QWEN3_TOKENIZATION_WORKERS:-4} + + qwen3-embedding-8b-gpu: + <<: *qwen3_common + image: *qwen3_gpu_image + container_name: unstract-qwen3-embedding-8b-gpu + profiles: + - embeddings-gpu + - embeddings-both + ports: + - "127.0.0.1:${QWEN3_8B_GPU_PORT:-8203}:80" + volumes: + - qwen3_embedding_8b_gpu_cache:/data + deploy: + resources: + reservations: + devices: + - driver: nvidia + count: 1 + capabilities: + - gpu + command: + - --model-id + - Qwen/Qwen3-Embedding-8B + - --revision + - ${QWEN3_8B_REVISION:-1d8ad4ca9b3dd8059ad90a75d4983776a23d44af} + - --served-model-name + - qwen3-embedding-8b + - --dtype + - float16 + - --pooling + - last-token + - --max-batch-tokens + - ${QWEN3_MAX_BATCH_TOKENS:-8192} + - --max-client-batch-size + - ${QWEN3_MAX_CLIENT_BATCH_SIZE:-32} + - --max-concurrent-requests + - ${QWEN3_MAX_CONCURRENT_REQUESTS:-4} + - --tokenization-workers + - ${QWEN3_TOKENIZATION_WORKERS:-4} + +volumes: + qwen3_embedding_06b_cpu_cache: + qwen3_embedding_4b_cpu_cache: + qwen3_embedding_8b_cpu_cache: + qwen3_embedding_06b_gpu_cache: + qwen3_embedding_4b_gpu_cache: + qwen3_embedding_8b_gpu_cache: diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index 9f6ca7789d..6bea2ea1d0 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -1,6 +1,7 @@ name: ${COMPOSE_PROJECT_NAME:-docker} include: - docker-compose-dev-essentials.yaml + - docker-compose-local-embeddings.yaml # Reusable host-gateway mapping so containers can reach services on the host # (e.g. host-installed Ollama at http://host.docker.internal:11434). diff --git a/docker/local-embeddings.config.json b/docker/local-embeddings.config.json new file mode 100644 index 0000000000..f36400772e --- /dev/null +++ b/docker/local-embeddings.config.json @@ -0,0 +1,55 @@ +{ + "description": "Local Qwen3 embedding benchmark matrix for Unstract.", + "query_prefix": "Instruct: Given a web search query, retrieve relevant passages that answer the query\nQuery: ", + "passage_prefix": "", + "models": [ + { + "id": "qwen3-embedding-06b", + "repository": "Qwen/Qwen3-Embedding-0.6B", + "revision": "97b0c614be4d77ee51c0cef4e5f07c00f9eb65b3", + "dimension": 1024, + "cpu": { + "service": "qwen3-embedding-06b-cpu", + "base_url": "http://127.0.0.1:8101/v1", + "internal_base_url": "http://qwen3-embedding-06b-cpu:80/v1" + }, + "gpu": { + "service": "qwen3-embedding-06b-gpu", + "base_url": "http://127.0.0.1:8201/v1", + "internal_base_url": "http://qwen3-embedding-06b-gpu:80/v1" + } + }, + { + "id": "qwen3-embedding-4b", + "repository": "Qwen/Qwen3-Embedding-4B", + "revision": "5cf2132abc99cad020ac570b19d031efec650f2b", + "dimension": 2560, + "cpu": { + "service": "qwen3-embedding-4b-cpu", + "base_url": "http://127.0.0.1:8102/v1", + "internal_base_url": "http://qwen3-embedding-4b-cpu:80/v1" + }, + "gpu": { + "service": "qwen3-embedding-4b-gpu", + "base_url": "http://127.0.0.1:8202/v1", + "internal_base_url": "http://qwen3-embedding-4b-gpu:80/v1" + } + }, + { + "id": "qwen3-embedding-8b", + "repository": "Qwen/Qwen3-Embedding-8B", + "revision": "1d8ad4ca9b3dd8059ad90a75d4983776a23d44af", + "dimension": 4096, + "cpu": { + "service": "qwen3-embedding-8b-cpu", + "base_url": "http://127.0.0.1:8103/v1", + "internal_base_url": "http://qwen3-embedding-8b-cpu:80/v1" + }, + "gpu": { + "service": "qwen3-embedding-8b-gpu", + "base_url": "http://127.0.0.1:8203/v1", + "internal_base_url": "http://qwen3-embedding-8b-gpu:80/v1" + } + } + ] +} diff --git a/docker/sample.essentials.env b/docker/sample.essentials.env index 51876fb8f9..238e61d72a 100644 --- a/docker/sample.essentials.env +++ b/docker/sample.essentials.env @@ -10,10 +10,6 @@ MINIO_ROOT_PASSWORD=minio123 MINIO_ACCESS_KEY=minio MINIO_SECRET_KEY=minio123 -QDRANT_USER=unstract_vector_dev -QDRANT_PASS=unstract_vector_pass -QDRANT_DB=unstract_vector_db - # RabbitMQ related envs RABBITMQ_DEFAULT_USER=admin RABBITMQ_DEFAULT_PASS=password diff --git a/docker/scripts/db-setup/vector_db_setup.sh b/docker/scripts/db-setup/vector_db_setup.sh new file mode 100755 index 0000000000..b45bce4fac --- /dev/null +++ b/docker/scripts/db-setup/vector_db_setup.sh @@ -0,0 +1,18 @@ +#!/usr/bin/env bash + +set -Eeuo pipefail + +: "${POSTGRES_USER:?POSTGRES_USER is required}" +: "${POSTGRES_DB:?POSTGRES_DB is required}" +: "${POSTGRES_SCHEMA:?POSTGRES_SCHEMA is required}" + +echo "Configuring pgvector database '$POSTGRES_DB' and schema '$POSTGRES_SCHEMA'" + +psql \ + --username "$POSTGRES_USER" \ + --dbname "$POSTGRES_DB" \ + --set ON_ERROR_STOP=1 \ + --set schema="$POSTGRES_SCHEMA" <<'SQL' +CREATE EXTENSION IF NOT EXISTS vector; +CREATE SCHEMA IF NOT EXISTS :"schema"; +SQL diff --git a/docs/local-dev-setup-executor-migration.md b/docs/local-dev-setup-executor-migration.md index 8bb6921fee..013e74b8f0 100644 --- a/docs/local-dev-setup-executor-migration.md +++ b/docs/local-dev-setup-executor-migration.md @@ -364,6 +364,11 @@ cd workers && ./run-worker.sh all | MinIO S3 API | 9000 | http://localhost:9000 | | MinIO Console | 9001 | http://localhost:9001 (minio/minio123) | | Qdrant | 6333 | http://localhost:6333 | +| Postgres vector DB | 5433 | `psql -h localhost -p 5433 -U unstract_dev -d unstract_db` | +| Weaviate HTTP API | 8084 | http://localhost:8084 | +| Weaviate gRPC API | 50051 | localhost:50051 | +| Milvus client API | 19530 | http://localhost:19530 | +| Milvus WebUI | 9091 | http://localhost:9091/webui/ | | Traefik Dashboard | 8080 | http://localhost:8080 | ### Application diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/base1.py b/unstract/sdk1/src/unstract/sdk1/adapters/base1.py index 716587a23e..115619e80a 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/base1.py +++ b/unstract/sdk1/src/unstract/sdk1/adapters/base1.py @@ -1733,6 +1733,10 @@ class OpenAICompatibleEmbeddingParameters(OpenAIEmbeddingParameters): # Some gateways are keyless; the endpoint is always required. api_key: str | None = None api_base: str + # Optional asymmetric retrieval controls. They are consumed by the SDK + # wrapper before dispatch and never sent as provider request parameters. + query_prefix: str = "" + passage_prefix: str = "" @staticmethod def validate(adapter_metadata: dict[str, "Any"]) -> dict[str, "Any"]: diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/embedding1/static/custom_openai.json b/unstract/sdk1/src/unstract/sdk1/adapters/embedding1/static/custom_openai.json index bde1665229..e4b1ec36f8 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/embedding1/static/custom_openai.json +++ b/unstract/sdk1/src/unstract/sdk1/adapters/embedding1/static/custom_openai.json @@ -33,6 +33,18 @@ "title": "API Base", "description": "Base URL for the OpenAI-compatible embeddings endpoint. Examples: https://gateway.example.com/v1, https://llm.example.net/openai/v1" }, + "query_prefix": { + "type": "string", + "title": "Query Prefix", + "default": "", + "description": "Optional text prepended to search queries before embedding. Use this for asymmetric models such as Qwen3. Example: Instruct: Given a web search query, retrieve relevant passages that answer the query\\nQuery: " + }, + "passage_prefix": { + "type": "string", + "title": "Passage Prefix", + "default": "", + "description": "Optional text prepended to indexed passages before embedding. Leave empty when the model's retrieval instructions apply only to queries." + }, "max_retries": { "type": "number", "minimum": 0, diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/vectordb/weaviate/src/weaviate.py b/unstract/sdk1/src/unstract/sdk1/adapters/vectordb/weaviate/src/weaviate.py index 1c08892f6b..7861f01c21 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/vectordb/weaviate/src/weaviate.py +++ b/unstract/sdk1/src/unstract/sdk1/adapters/vectordb/weaviate/src/weaviate.py @@ -1,5 +1,6 @@ import logging import os +from urllib.parse import urlparse import weaviate from llama_index.core.vector_stores.types import BasePydanticVectorStore @@ -58,6 +59,39 @@ def get_doc_url() -> str: def get_vector_db_instance(self) -> BasePydanticVectorStore: return self._vector_db_instance + def _connect(self) -> weaviate.Client: + """Connect to either a local Weaviate server or Weaviate Cloud. + + The adapter historically used the cloud connector for every URL. Local + Docker deployments use a different gRPC endpoint and can be reached + without a cloud cluster URL, so route plain HTTP URLs through + ``connect_to_local`` while preserving the existing HTTPS cloud path. + """ + configured_url = str(self._config.get(Constants.URL, "")).strip() + if not configured_url: + raise ValueError("Weaviate URL is required") + + api_key = str(self._config.get(Constants.API_KEY) or "").strip() + auth_credentials = Auth.api_key(api_key) if api_key else None + parsed_url = urlparse( + configured_url if "://" in configured_url else f"https://{configured_url}" + ) + + if parsed_url.scheme == "http": + if parsed_url.hostname is None: + raise ValueError("Weaviate URL must include a hostname") + return weaviate.connect_to_local( + host=parsed_url.hostname, + port=parsed_url.port or 8080, + grpc_port=50051, + auth_credentials=auth_credentials, + ) + + return weaviate.connect_to_weaviate_cloud( + cluster_url=configured_url, + auth_credentials=auth_credentials, + ) + def _get_vector_db_instance(self) -> BasePydanticVectorStore: try: collection_name = VectorDBHelper.get_collection_name( @@ -68,10 +102,7 @@ def _get_vector_db_instance(self) -> BasePydanticVectorStore: # LLama-index throws the error if not capitalised while using # Weaviate self._collection_name = collection_name.capitalize() - self._client = weaviate.connect_to_weaviate_cloud( - cluster_url=str(self._config.get(Constants.URL)), - auth_credentials=Auth.api_key(str(self._config.get(Constants.API_KEY))), - ) + self._client = self._connect() try: # Class definition object. Weaviate's autoschema diff --git a/unstract/sdk1/src/unstract/sdk1/embedding.py b/unstract/sdk1/src/unstract/sdk1/embedding.py index 1e02944f83..97e753aa78 100644 --- a/unstract/sdk1/src/unstract/sdk1/embedding.py +++ b/unstract/sdk1/src/unstract/sdk1/embedding.py @@ -107,6 +107,11 @@ def __init__( self.platform_kwargs: dict[str, object] = kwargs self.kwargs: dict[str, object] = self.adapter.validate(self._adapter_metadata) self._cost_model: str | None = self.kwargs.pop("cost_model", None) + # Some local embedding models use asymmetric query/document + # prefixes. These are adapter controls, not provider request + # parameters, so consume them before calling LiteLLM. + self._query_prefix = str(self.kwargs.pop("query_prefix", "") or "") + self._passage_prefix = str(self.kwargs.pop("passage_prefix", "") or "") # Client-side batching hint, not an API field — keep it off the wire. self.kwargs.pop("embed_batch_size", None) except (ValidationError, ValueError) as e: @@ -134,12 +139,21 @@ def _prepare_call(self, input_type: str | None) -> tuple[str, dict, int | None]: kwargs["input_type"] = input_type return model, kwargs, max_retries + def _prepare_text(self, text: str, input_type: str | None) -> str: + """Apply an optional query or passage prefix before provider dispatch.""" + prefix = self._query_prefix if input_type == "query" else self._passage_prefix + return f"{prefix}{text}" if prefix else text + def get_embedding(self, text: str, input_type: str = "query") -> list[float]: """Return embedding vector for query string.""" try: model, kwargs, max_retries = self._prepare_call(input_type) resp = call_with_retry( - lambda: litellm.embedding(model=model, input=[text], **kwargs), + lambda: litellm.embedding( + model=model, + input=[self._prepare_text(text, input_type)], + **kwargs, + ), max_retries=max_retries, retry_predicate=is_retryable_litellm_error, description=self._get_adapter_info(), @@ -155,7 +169,11 @@ def get_embeddings( try: model, kwargs, max_retries = self._prepare_call(input_type) resp = call_with_retry( - lambda: litellm.embedding(model=model, input=texts, **kwargs), + lambda: litellm.embedding( + model=model, + input=[self._prepare_text(text, input_type) for text in texts], + **kwargs, + ), max_retries=max_retries, retry_predicate=is_retryable_litellm_error, description=self._get_adapter_info(), @@ -169,7 +187,11 @@ async def get_aembedding(self, text: str, input_type: str = "query") -> list[flo try: model, kwargs, max_retries = self._prepare_call(input_type) resp = await acall_with_retry( - lambda: litellm.aembedding(model=model, input=[text], **kwargs), + lambda: litellm.aembedding( + model=model, + input=[self._prepare_text(text, input_type)], + **kwargs, + ), max_retries=max_retries, retry_predicate=is_retryable_litellm_error, description=self._get_adapter_info(), @@ -185,7 +207,11 @@ async def get_aembeddings( try: model, kwargs, max_retries = self._prepare_call(input_type) resp = await acall_with_retry( - lambda: litellm.aembedding(model=model, input=texts, **kwargs), + lambda: litellm.aembedding( + model=model, + input=[self._prepare_text(text, input_type) for text in texts], + **kwargs, + ), max_retries=max_retries, retry_predicate=is_retryable_litellm_error, description=self._get_adapter_info(), diff --git a/unstract/sdk1/tests/test_branded_openai_adapters.py b/unstract/sdk1/tests/test_branded_openai_adapters.py index c1556a3dd7..13f4610c3f 100644 --- a/unstract/sdk1/tests/test_branded_openai_adapters.py +++ b/unstract/sdk1/tests/test_branded_openai_adapters.py @@ -447,6 +447,8 @@ def test_compatible_embedding_schema_loadable() -> None: assert schema["title"] == "OpenAI Compatible Embedding" assert "api_base" in schema["required"] assert "model" in schema["required"] + assert "query_prefix" in schema["properties"] + assert "passage_prefix" in schema["properties"] @pytest.mark.parametrize( @@ -536,3 +538,40 @@ def test_compatible_embedding_omits_input_type( ) assert "input_type" not in captured assert captured["model"] == "openai/BAAI/bge-m3" + + +def test_compatible_embedding_applies_query_and_passage_prefixes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import unstract.sdk1.embedding as emb_mod + + calls: list[list[str]] = [] + provider_kwargs: list[dict[str, object]] = [] + + def fake_embedding(model: str, input: list[str], **kwargs: object) -> dict: # noqa: A002 + del model + provider_kwargs.append(kwargs) + calls.append(input) + return {"data": [{"embedding": [0.0, 1.0]}] * len(input)} + + monkeypatch.setattr(emb_mod.litellm, "embedding", fake_embedding) + emb = emb_mod.Embedding( + adapter_id=OpenAICompatibleEmbeddingAdapter.get_id(), + adapter_metadata={ + "model": "qwen3-embedding-06b", + "api_base": "http://qwen3-embedding-06b-cpu:80/v1", + "api_key": "", + "query_prefix": "Q: ", + "passage_prefix": "P: ", + }, + ) + + # The constructor's connection test uses the query path. + assert calls[0] == ["Q: Hello, I am Unstract"] + emb.get_embedding("question") + emb.get_embeddings(["document 1", "document 2"], input_type="passage") + + assert calls[1] == ["Q: question"] + assert calls[2] == ["P: document 1", "P: document 2"] + assert all("query_prefix" not in kwargs for kwargs in provider_kwargs) + assert all("passage_prefix" not in kwargs for kwargs in provider_kwargs) diff --git a/unstract/sdk1/tests/test_weaviate_adapter.py b/unstract/sdk1/tests/test_weaviate_adapter.py new file mode 100644 index 0000000000..440232b5e7 --- /dev/null +++ b/unstract/sdk1/tests/test_weaviate_adapter.py @@ -0,0 +1,63 @@ +from unittest.mock import MagicMock, patch + +from unstract.sdk1.adapters.vectordb.weaviate.src.weaviate import ( + Constants, + Weaviate, +) + + +def _adapter(url: str, api_key: str = "") -> Weaviate: + adapter = object.__new__(Weaviate) + adapter._config = { + Constants.URL: url, + Constants.API_KEY: api_key, + } + return adapter + + +def test_local_url_uses_local_connector_with_api_key() -> None: + client = MagicMock() + + with patch( + "unstract.sdk1.adapters.vectordb.weaviate.src.weaviate.weaviate.connect_to_local", + return_value=client, + ) as connect_to_local: + result = _adapter("http://weaviate:8080", "local-key")._connect() + + assert result is client + connect_to_local.assert_called_once() + call_kwargs = connect_to_local.call_args.kwargs + assert call_kwargs["host"] == "weaviate" + assert call_kwargs["port"] == 8080 + assert call_kwargs["grpc_port"] == 50051 + assert call_kwargs["auth_credentials"] is not None + + +def test_local_url_allows_anonymous_connection() -> None: + client = MagicMock() + + with patch( + "unstract.sdk1.adapters.vectordb.weaviate.src.weaviate.weaviate.connect_to_local", + return_value=client, + ) as connect_to_local: + _adapter("http://localhost:8084")._connect() + + assert connect_to_local.call_args.kwargs["host"] == "localhost" + assert connect_to_local.call_args.kwargs["port"] == 8084 + assert connect_to_local.call_args.kwargs["auth_credentials"] is None + + +def test_https_url_preserves_cloud_connector() -> None: + client = MagicMock() + + with patch( + "unstract.sdk1.adapters.vectordb.weaviate.src.weaviate.weaviate.connect_to_weaviate_cloud", + return_value=client, + ) as connect_to_cloud: + result = _adapter("https://example.weaviate.cloud", "cloud-key")._connect() + + assert result is client + connect_to_cloud.assert_called_once() + assert connect_to_cloud.call_args.kwargs["cluster_url"] == ( + "https://example.weaviate.cloud" + )