diff --git a/CHANGELOG.md b/CHANGELOG.md index 95b316dea..d007830d4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,9 @@ +## [1.6.29] + +### Enhancements + +- **feat: add `rag-ingest-mcp`, a local RAG MCP server (`unstructured_ingest/mcp/`).** A FastMCP STDIO server that lands Transform MCP Element JSON into a local vector store and serves matched-space similarity search, with three tools (`load_transform_output`, `search`, `list_collections`) and three backends behind one `VectorStore` contract selected by `URAG_STORE_BACKEND`: Chroma (default, embedded), Qdrant (embedded local mode or a server), and pgvector (writes the Unstructured-compliant schema — `id, type, record_id, element_id, text, embeddings` — fitting rows to a pre-created table or auto-provisioning an equivalent one with a cosine HNSW index, and casting the vector to the column's own `vector`/`halfvec` type). A collection is pinned to one embedding space (provider, model, dimension) on first load; loads with a different model are refused and every query is embedded into the pinned space, so corpus and query vectors are always comparable. Writes are conformed by the corresponding ingest connector stagers with deterministic element ids, so re-loading a source upserts instead of duplicating. Includes per-backend getting-started guides under `unstructured_ingest/mcp/examples/`. + ## [1.6.28] ### Fixes diff --git a/test/unit/mcp/__init__.py b/test/unit/mcp/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/test/unit/mcp/test_server.py b/test/unit/mcp/test_server.py new file mode 100644 index 000000000..65ca9b30d --- /dev/null +++ b/test/unit/mcp/test_server.py @@ -0,0 +1,108 @@ +"""Server-level smoke tests for rag-ingest-mcp, offline. + +Drives the FastMCP server in-memory with passthrough (pre-embedded) elements, +so no embedding provider key or running service is needed. Runs against both +embedded backends; pgvector's server path is covered by the store conformance +suite plus an integration environment. +""" + +from __future__ import annotations + +import json + +import anyio +import pytest + +pytest.importorskip("fastmcp", reason="rag-ingest-mcp requires fastmcp") + +from fastmcp import Client # noqa: E402 + +from unstructured_ingest.mcp.server import mcp # noqa: E402 + +VECTORS = { + "cats": [1.0, 0.0, 0.0, 0.0], + "dogs": [0.0, 1.0, 0.0, 0.0], + "stocks": [0.0, 0.0, 1.0, 0.0], +} + + +@pytest.fixture(params=["chroma", "qdrant"]) +def backend_env(request, monkeypatch, tmp_path) -> str: + pytest.importorskip("chromadb" if request.param == "chroma" else "qdrant_client") + monkeypatch.setenv("URAG_STORE_BACKEND", request.param) + monkeypatch.setenv("URAG_CHROMA_PATH", str(tmp_path / "chroma")) + monkeypatch.setenv("URAG_QDRANT_PATH", str(tmp_path / "qdrant")) + monkeypatch.setenv("URAG_EMBED_MODEL", "smoke-model") + return request.param + + +def write_elements_file(path, texts_to_vectors=VECTORS) -> str: + elements = [ + { + "element_id": f"element-{index}", + "text": text, + "embeddings": vector, + "metadata": {"filename": path.name}, + } + for index, (text, vector) in enumerate(texts_to_vectors.items()) + ] + path.write_text(json.dumps(elements)) + return str(path) + + +def call(tool: str, arguments: dict) -> dict: + async def run(): + async with Client(mcp) as client: + return (await client.call_tool(tool, arguments)).data + + return anyio.run(run) + + +def test_load_search_and_list(backend_env, tmp_path): + source = write_elements_file(tmp_path / "doc.json") + result = call( + "load_transform_output", + {"sources": [source], "collection": "smoke", "embed_policy": "passthrough"}, + ) + assert result["files"][0]["loaded"] == 3, result + assert result["space"]["embed_dim"] == 4 + + listed = call("list_collections", {}) + entry = next(c for c in listed["collections"] if c["collection"] == "smoke") + assert entry["model"] == "smoke-model" + assert entry["dimension"] == 4 + + +def test_directory_source_expands_to_all_json(backend_env, tmp_path): + docs = tmp_path / "transform-output" + docs.mkdir() + write_elements_file(docs / "a.json") + write_elements_file(docs / "b.json") + result = call( + "load_transform_output", + {"sources": [str(docs)], "collection": "smokedir", "embed_policy": "passthrough"}, + ) + loaded = [f for f in result["files"] if "loaded" in f] + assert len(loaded) == 2 and all(f["loaded"] == 3 for f in loaded), result + + +def test_space_guard_reports_per_source_error(backend_env, tmp_path): + call( + "load_transform_output", + { + "sources": [write_elements_file(tmp_path / "a.json")], + "collection": "smokeguard", + "embed_policy": "passthrough", + }, + ) + five_dim = {text: vector + [0.0] for text, vector in VECTORS.items()} + result = call( + "load_transform_output", + { + "sources": [write_elements_file(tmp_path / "b.json", five_dim)], + "collection": "smokeguard", + "embed_policy": "passthrough", + "embed_model": "other-model", + }, + ) + assert "refusing" in result["files"][0]["error"], result diff --git a/test/unit/mcp/test_stores.py b/test/unit/mcp/test_stores.py new file mode 100644 index 000000000..cbc5ca483 --- /dev/null +++ b/test/unit/mcp/test_stores.py @@ -0,0 +1,136 @@ +"""Backend conformance suite for rag-ingest-mcp vector stores. + +Every backend must pass the same tests; passing this suite is what "supported +backend" means. Chroma and Qdrant run embedded (a tmp directory, no service). +pgvector needs a reachable Postgres with the vector extension available — set +``URAG_TEST_PG_DSN`` (e.g. a throwaway ``pgvector/pgvector`` container) or its +cases skip. +""" + +from __future__ import annotations + +import os +import uuid + +import pytest + +from unstructured_ingest.data_types.file_data import FileData, SourceIdentifiers +from unstructured_ingest.mcp.stores import EmbeddingSpace, SpaceMismatch, VectorStore + +# Orthogonal 4-dim vectors: cosine ranking is exact, no embedding model needed. +TEXTS_TO_VECTORS = { + "cats": [1.0, 0.0, 0.0, 0.0], + "dogs": [0.0, 1.0, 0.0, 0.0], + "stocks": [0.0, 0.0, 1.0, 0.0], +} +SPACE = EmbeddingSpace(provider="test", model="conformance-model", dimension=4) +OTHER_SPACE = EmbeddingSpace(provider="test", model="other-model", dimension=5) + + +def make_elements() -> list[dict]: + return [ + { + "element_id": f"element-{index}", + "text": text, + "embeddings": vector, + "metadata": {"filename": "doc.txt"}, + } + for index, (text, vector) in enumerate(TEXTS_TO_VECTORS.items()) + ] + + +def make_file_data(identifier: str = "doc-1") -> FileData: + return FileData( + identifier=identifier, + connector_type="transform_mcp", + source_identifiers=SourceIdentifiers(filename=identifier, fullpath=identifier), + ) + + +@pytest.fixture(params=["chroma", "qdrant", "pgvector"]) +def store(request, tmp_path) -> VectorStore: + if request.param == "chroma": + pytest.importorskip("chromadb", reason="chroma backend requires the chroma extra") + from unstructured_ingest.mcp.stores.chroma import ChromaStore + + yield ChromaStore(str(tmp_path / "chroma")) + return + if request.param == "qdrant": + pytest.importorskip("qdrant_client", reason="qdrant backend requires the qdrant extra") + from unstructured_ingest.mcp.stores.qdrant import QdrantStore + + yield QdrantStore(path=str(tmp_path / "qdrant")) + return + dsn = os.environ.get("URAG_TEST_PG_DSN") + if not dsn: + pytest.skip("URAG_TEST_PG_DSN not set; skipping pgvector conformance") + pytest.importorskip("psycopg", reason="pgvector backend requires psycopg") + from unstructured_ingest.mcp.stores.pgvector import PgvectorStore + + pg_store = PgvectorStore(dsn) + yield pg_store + # The database outlives the test run, so drop everything this test created. + import psycopg + + with psycopg.connect(dsn) as conn: + for name in pg_store.collections(): + if name.startswith("conf_"): + conn.execute(f'DROP TABLE IF EXISTS "{name}"') + conn.execute("DELETE FROM rag_ingest_spaces WHERE collection = %s", (name,)) + + +@pytest.fixture +def collection() -> str: + # Unique per test so a shared backend (pgvector) never sees collisions. + return f"conf_{uuid.uuid4().hex[:12]}" + + +def load(store: VectorStore, collection: str, identifier: str = "doc-1") -> int: + store.ensure_space(collection, SPACE) + return store.write(collection, make_elements(), make_file_data(identifier)) + + +def test_load_then_search_ranks_by_cosine(store, collection): + assert load(store, collection) == 3 + matches = store.search(collection, [0.9, 0.1, 0.0, 0.0], k=3) + assert [m["text"] for m in matches] == ["cats", "dogs", "stocks"] + scores = [m["score"] for m in matches] + assert scores == sorted(scores, reverse=True) + assert scores[0] > 0.9 + + +def test_search_returns_metadata(store, collection): + load(store, collection) + top = store.search(collection, [1.0, 0.0, 0.0, 0.0], k=1)[0] + assert top["metadata"], "flattened element metadata should be returned" + + +def test_space_is_pinned_and_guarded(store, collection): + load(store, collection) + space = store.space_of(collection) + assert (space.model, space.dimension) == (SPACE.model, SPACE.dimension) + with pytest.raises(SpaceMismatch): + store.ensure_space(collection, OTHER_SPACE) + + +def test_reload_upserts_instead_of_duplicating(store, collection): + load(store, collection) + load(store, collection) + matches = store.search(collection, [1.0, 0.0, 0.0, 0.0], k=10) + assert len(matches) == 3 + + +def test_unembedded_elements_are_dropped(store, collection): + store.ensure_space(collection, SPACE) + elements = make_elements() + [{"element_id": "bare", "text": "no vector"}] + assert store.write(collection, elements, make_file_data()) == 3 + + +def test_collections_lists_what_was_loaded(store, collection): + load(store, collection) + assert collection in store.collections() + + +def test_space_of_unknown_collection_raises(store): + with pytest.raises(Exception): + store.space_of("conf_never_loaded") diff --git a/unstructured_ingest/__version__.py b/unstructured_ingest/__version__.py index 5a23218dd..d2e431858 100644 --- a/unstructured_ingest/__version__.py +++ b/unstructured_ingest/__version__.py @@ -1 +1 @@ -__version__ = "1.6.28" # pragma: no cover +__version__ = "1.6.29" # pragma: no cover diff --git a/unstructured_ingest/mcp/.env.example b/unstructured_ingest/mcp/.env.example new file mode 100644 index 000000000..834328155 --- /dev/null +++ b/unstructured_ingest/mcp/.env.example @@ -0,0 +1,27 @@ +# --- required for embedding (ingest-side embed and every query) --- +OPENAI_API_KEY=sk-... + +# --- local vector store --- +# Backend: chroma (default, embedded) | qdrant (embedded or server) | pgvector +# URAG_STORE_BACKEND=chroma + +# chroma: directory where the collections persist. This IS your local RAG corpus. +URAG_CHROMA_PATH=~/.unstructured-rag/chroma + +# qdrant: embedded local mode by default; set URAG_QDRANT_URL to use a server. +# URAG_QDRANT_PATH=~/.unstructured-rag/qdrant +# URAG_QDRANT_URL= +# URAG_QDRANT_API_KEY= + +# pgvector: required when URAG_STORE_BACKEND=pgvector. +# URAG_PG_DSN=postgresql://user:pass@localhost:5432/db + +# --- embedding defaults (must match what Transform used, in passthrough mode) --- +URAG_EMBED_PROVIDER=openai +URAG_EMBED_MODEL=text-embedding-3-small + +# --- optional knobs (defaults shown) --- +# URAG_STORE_BACKEND=chroma +# URAG_EMBED_BATCH_SIZE=64 +# URAG_MAX_FETCH_BYTES=536870912 +# OPENAI_BASE_URL= diff --git a/unstructured_ingest/mcp/README.md b/unstructured_ingest/mcp/README.md new file mode 100644 index 000000000..53301b354 --- /dev/null +++ b/unstructured_ingest/mcp/README.md @@ -0,0 +1,157 @@ +# rag-ingest-mcp + +A small **local STDIO MCP server** that turns [Unstructured Transform +MCP](https://docs.unstructured.io) output into a queryable **local** RAG corpus. + +Transform MCP does the heavy lifting on the platform — parse → chunk → *(optionally)* +embed — and hands back Element JSON. This server, running as a subprocess of your +MCP host, fetches that JSON **out of band**, embeds it locally if it wasn't +already, upserts it into a local vector store — **Chroma** (default), **Qdrant**, +or **pgvector** — and serves similarity search. It reuses `unstructured-ingest`'s +own embedders and connector stagers, so the write path is identical to the +ingest pipeline's. + +``` +┌─ MCP Host (Claude Code / Desktop) ────────────────────────────────┐ +│ the agent orchestrates both servers; large JSON never enters its │ +│ context — it only passes a download_url string │ +│ │ +│ ① Transform MCP (remote, OAuth) │ +│ parse → chunk → [embed] ⇒ Element JSON + download_url │ +│ │ +│ ② this server (local, STDIO) │ +│ load_transform_output(download_url) → local store │ +│ search(query) → nearest chunks │ +└───────────────────────────────────────────────────────────────────┘ + │ + ▼ local vector store (your corpus): chroma | qdrant | pgvector +``` + +## The one design idea: embedding is the only knob + +Everything is fixed except **where the embedding happens**, and this server +absorbs that choice per load via `embed_policy`: + +| | **Transform embeds** (passthrough) | **This server embeds** (local) | +|---|---|---| +| Transform `stages` | `{chunk, embed}` | `{chunk}` only | +| JSON returned | text **+ vectors** | text, **no vectors** | +| This server does | stage → upsert as-is | `embed_documents()` → stage → upsert | +| Query vector | this server rebuilds the same model (needs the provider key here) | the *same encoder object* — symmetric by construction | +| Embedding billed | on the platform | on your local OpenAI key | + +`embed_policy="auto"` (default) picks per source by sniffing whether vectors are +present, so the agent chooses simply by toggling Transform's `embed` stage. + +**The rule this enforces:** a collection is pinned to one embedding space +(provider, model, dimension) on first load; a later load with a different model +is **refused**, and every query is embedded into that same space. That's what +keeps retrieval correct — mixing models silently returns nonsense. + +## Install + +From the `unstructured-ingest` repo root (this module lives inside the package): + +```bash +pip install -e ".[chroma,openai]" # default backend +pip install -r unstructured_ingest/mcp/requirements.txt # fastmcp +``` + +For the other backends, add what they need: + +```bash +pip install -e ".[qdrant]" # qdrant backend (embedded local mode or server) +pip install "psycopg[binary]" # pgvector backend (psycopg 3) +``` + +## Configure + +Copy `.env.example` and set at least `OPENAI_API_KEY`. Key settings: + +| Env var | Default | Meaning | +|---|---|---| +| `OPENAI_API_KEY` | — | required; embeds queries (and corpus in local mode) | +| `URAG_STORE_BACKEND` | `chroma` | `chroma`, `qdrant`, or `pgvector` | +| `URAG_CHROMA_PATH` | `~/.unstructured-rag/chroma` | chroma: where the corpus persists | +| `URAG_QDRANT_PATH` | `~/.unstructured-rag/qdrant` | qdrant: embedded local-mode directory | +| `URAG_QDRANT_URL` | — | qdrant: server url (overrides local mode) | +| `URAG_QDRANT_API_KEY` | — | qdrant: server api key, if any | +| `URAG_PG_DSN` | — | pgvector: `postgresql://user:pass@host:5432/db` | +| `URAG_EMBED_MODEL` | `text-embedding-3-small` | must match Transform's model in passthrough | +| `URAG_EMBED_PROVIDER` | `openai` | embed provider | + +## Wire it into your MCP host + +Register **two** servers: this one (local, STDIO, `python -m +unstructured_ingest.mcp`) and the remote Transform MCP (HTTP/OAuth). In Claude +Desktop use `claude_desktop_config.json`; in Claude Code use `.mcp.json`. + +Each supported backend has a getting-started guide with a ready-to-copy host +config — start with the [examples overview](examples/README.md), or jump +straight to your backend: + +| Backend | Getting-started guide | Host config | +|---|---|---| +| **Chroma** (default, zero-infra) | [examples/chroma/](examples/chroma/README.md) | [mcp-config.json](examples/chroma/mcp-config.json) | +| **Qdrant** (embedded or docker server) | [examples/qdrant/](examples/qdrant/README.md) | [mcp-config.json](examples/qdrant/mcp-config.json) | +| **pgvector** (your Postgres) | [examples/pgvector/](examples/pgvector/README.md) | [mcp-config.json](examples/pgvector/mcp-config.json) | + +Every guide covers setup, docker-compose where a service is needed, the full +agentic loop with Transform MCP, CLI inspection of the database, and when (not) +to pair the DB vendor's own MCP server. + +## The flow the agent follows + +1. **Transform MCP** — `request_file_upload_url` → `curl -X PUT` the bytes → + `transform_files(file_refs, stages={"chunk": {...}, "embed": {}})` + *(drop `embed` to embed locally instead)* → poll `check_transform_status` → + `get_transform_results(job_id, output_format="json", image_base64="none")` + → yields a `download_url` per file. +2. **This server** — `load_transform_output(sources=[download_url, ...], + collection="my_docs")`. It fetches, embeds if needed, and upserts. Returns + `{loaded, space, files}`. +3. **Query** — `search("your question", collection="my_docs", k=5)` → nearest + chunks with scores and metadata; the agent answers from them. + +## Tools + +- **`load_transform_output(sources, collection, embed_policy="auto", embed_model=None)`** + — fetch Element JSON (URLs or paths), embed per policy, upsert. Per-source + errors don't sink the batch. +- **`search(query, collection, k=5)`** — embed the query in the collection's + pinned space and return the nearest chunks. +- **`list_collections()`** — collections and the model each is pinned to. + +## Backends + +Three backends, all holding the same contract (space pinned per collection on +first load, mismatched loads refused, deterministic-id upsert, cosine scores, +rows conformed by the corresponding `unstructured-ingest` stager). The +conformance suite in `test/unit/mcp/test_stores.py` runs identically against +every backend and is the definition of "supported". + +| Backend | Runs as | Rows go through | Space pin lives in | +|---|---|---|---| +| `chroma` (default) | embedded, a local directory | `ChromaUploadStager` | collection metadata | +| `qdrant` | embedded local mode, or a server via `URAG_QDRANT_URL` | `QdrantUploadStager` | a reserved meta collection | +| `pgvector` | your Postgres (`URAG_PG_DSN`) | `SQLUploadStager` → `{id, record_id, text, metadata JSONB, embedding vector(dim)}` | `rag_ingest_spaces` table | + +The first load auto-provisions the collection/table from the pinned space — no +schema setup is asked of you. For pgvector that includes a cosine **HNSW index** +when the dimension allows it (pgvector caps indexable vectors at 2,000 dims; +wider spaces work but search sequentially — the default +`text-embedding-3-small`/1536 stays comfortably inside the cap). + +## Notes & limits +- **Passthrough still needs a local provider key.** Transform embedded the + corpus, but *this* server embeds each query, so `OPENAI_API_KEY` (matching + Transform's model) must be set even when you never embed a corpus locally. +- **Passthrough records the model you tell it.** The space is pinned from + `embed_model` (or the server default, `text-embedding-3-small`). If Transform + embedded with a different model — its own default is `text-embedding-3-large` — + pass `embed_model` on the load, or queries will be embedded in the wrong space + and rejected on dimension mismatch. +- **Re-loading a source upserts** rather than duplicating: element ids are + derived deterministically from the source identity, so re-running a load + overwrites the same rows. +- Distance is **cosine**; `search` scores are `1 − cosine_distance`. diff --git a/unstructured_ingest/mcp/__init__.py b/unstructured_ingest/mcp/__init__.py new file mode 100644 index 000000000..5be7f5a2c --- /dev/null +++ b/unstructured_ingest/mcp/__init__.py @@ -0,0 +1,9 @@ +"""Local STDIO MCP server that lands Unstructured Transform MCP output into a +local vector store for retrieval. + +This module pairs with the remote Transform MCP. Transform parses, chunks, and +(optionally) embeds documents on the platform and hands back Element JSON; this +server fetches that JSON out of band, embeds it locally when Transform did not, +upserts it into a local Chroma collection, and serves matched-space similarity +search. See ``README.md`` for the end-to-end flow. +""" diff --git a/unstructured_ingest/mcp/__main__.py b/unstructured_ingest/mcp/__main__.py new file mode 100644 index 000000000..35f58dbaf --- /dev/null +++ b/unstructured_ingest/mcp/__main__.py @@ -0,0 +1,6 @@ +"""Entry point so ``python -m unstructured_ingest.mcp`` starts the server.""" + +from unstructured_ingest.mcp.server import main + +if __name__ == "__main__": + main() diff --git a/unstructured_ingest/mcp/config.py b/unstructured_ingest/mcp/config.py new file mode 100644 index 000000000..5f9ebd868 --- /dev/null +++ b/unstructured_ingest/mcp/config.py @@ -0,0 +1,64 @@ +"""Environment-driven configuration for the rag-ingest-mcp server. + +Every setting has a default that makes the common case — a local Chroma store +embedding with OpenAI ``text-embedding-3-small`` — work with only an +``OPENAI_API_KEY`` set. In passthrough mode the corpus vectors were produced by +Transform, so the space recorded (and used for every query) must name the model +Transform actually embedded with — pass ``embed_model`` on load when it differs +from this default. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass + +# text-embedding-3-small (1536 dims): cheap, strong default for a local corpus, +# and its dimension fits every candidate backend's ANN-index limits (pgvector +# caps HNSW-indexable vectors at 2000 dims, which 3-large's 3072 exceeds). +DEFAULT_EMBED_PROVIDER = "openai" +DEFAULT_EMBED_MODEL = "text-embedding-3-small" + +# A JSON render of a text corpus (image base64 stripped) is large but bounded; +# this cap is a backstop against a runaway download, not a tuning knob. +_DEFAULT_MAX_FETCH_BYTES = 512 * 1024 * 1024 + + +@dataclass(frozen=True) +class Config: + store_backend: str + chroma_path: str + qdrant_path: str + qdrant_url: str | None + qdrant_api_key: str | None + pg_dsn: str | None + embed_provider: str + embed_model: str + openai_api_key: str | None + openai_base_url: str | None + embed_batch_size: int + max_fetch_bytes: int + + @classmethod + def from_env(cls) -> "Config": + return cls( + store_backend=os.environ.get("URAG_STORE_BACKEND", "chroma"), + chroma_path=os.path.expanduser( + os.environ.get("URAG_CHROMA_PATH", "~/.unstructured-rag/chroma") + ), + # Qdrant runs embedded on a local path unless a server url is given. + qdrant_path=os.path.expanduser( + os.environ.get("URAG_QDRANT_PATH", "~/.unstructured-rag/qdrant") + ), + qdrant_url=os.environ.get("URAG_QDRANT_URL"), + qdrant_api_key=os.environ.get("URAG_QDRANT_API_KEY"), + pg_dsn=os.environ.get("URAG_PG_DSN"), + embed_provider=os.environ.get("URAG_EMBED_PROVIDER", DEFAULT_EMBED_PROVIDER), + embed_model=os.environ.get("URAG_EMBED_MODEL", DEFAULT_EMBED_MODEL), + openai_api_key=os.environ.get("OPENAI_API_KEY"), + openai_base_url=os.environ.get("OPENAI_BASE_URL"), + embed_batch_size=int(os.environ.get("URAG_EMBED_BATCH_SIZE", "64")), + max_fetch_bytes=int( + os.environ.get("URAG_MAX_FETCH_BYTES", str(_DEFAULT_MAX_FETCH_BYTES)) + ), + ) diff --git a/unstructured_ingest/mcp/embedding.py b/unstructured_ingest/mcp/embedding.py new file mode 100644 index 000000000..fe19ba43d --- /dev/null +++ b/unstructured_ingest/mcp/embedding.py @@ -0,0 +1,50 @@ +"""Embedding encoders for the ingest RAG MCP server. + +Both sides of the system embed through this one builder: the ingest side (when +Transform did not embed) and the query side (always). Routing both through the +same provider/model is what guarantees a query vector lands in the same space as +the corpus it searches — the single property retrieval correctness depends on. + +Only providers whose key this server can supply locally belong in the registry. +Adding one is a single row plus its ``unstructured-ingest`` extra. +""" + +from __future__ import annotations + +from unstructured_ingest.embed.interfaces import BaseEmbeddingEncoder +from unstructured_ingest.embed.openai import OpenAIEmbeddingConfig, OpenAIEmbeddingEncoder + +# provider -> (config class, encoder class). The config's ``model_name`` alias +# and ``api_key`` field are shared across the ingest encoders, so one call shape +# fits every provider added here. +_PROVIDERS: dict[str, tuple[type, type]] = { + "openai": (OpenAIEmbeddingConfig, OpenAIEmbeddingEncoder), +} + + +def build_encoder( + *, + provider: str, + model: str, + api_key: str | None, + base_url: str | None = None, + batch_size: int = 64, +) -> BaseEmbeddingEncoder: + """Construct an ingest embedding encoder for ``provider``/``model``. + + Raises ``ValueError`` — rather than failing deep inside a request — when the + provider is unsupported or its key is absent, so the tool can report a clear + remedy. + """ + if provider not in _PROVIDERS: + raise ValueError( + f"unsupported embed provider {provider!r}; supported: {sorted(_PROVIDERS)}" + ) + if not api_key: + raise ValueError( + f"no API key configured for embed provider {provider!r}; set OPENAI_API_KEY" + ) + config_cls, encoder_cls = _PROVIDERS[provider] + # ``model_name`` is the public alias of the encoder config's model field. + config = config_cls(api_key=api_key, model_name=model, base_url=base_url, batch_size=batch_size) + return encoder_cls(config=config) diff --git a/unstructured_ingest/mcp/examples/README.md b/unstructured_ingest/mcp/examples/README.md new file mode 100644 index 000000000..66794bb19 --- /dev/null +++ b/unstructured_ingest/mcp/examples/README.md @@ -0,0 +1,110 @@ +# Local agentic RAG with rag-ingest-mcp — examples + +These guides set up a complete **local agentic RAG loop**: your agent (Claude +Code, Claude Desktop, or any MCP host) parses documents with the hosted +**Transform MCP**, lands the chunks in a **local vector database** through +**rag-ingest-mcp**, and then answers questions by searching that corpus — all +orchestrated by the agent itself, tool call by tool call. + +Pick a backend and follow its guide end to end: + +| Guide | Backend runs as | Reach for it when | +|---|---|---| +| [chroma/](chroma/README.md) | embedded — a local directory, no service | you want zero infrastructure; the default | +| [qdrant/](qdrant/README.md) | embedded local mode **or** a docker server | you want a dedicated vector DB with a web dashboard, or already run Qdrant | +| [pgvector/](pgvector/README.md) | your Postgres (docker or existing) | you already run Postgres, or want SQL access to the same corpus | + +Every backend behaves identically through the MCP tools — same three tools, +same guarantees. Switching is one env var (`URAG_STORE_BACKEND`); nothing about +the Transform side or the agent's workflow changes. + +## The loop, once + +``` +┌─ MCP host (your agent) ────────────────────────────────────────────┐ +│ │ +│ ① Transform MCP (remote, OAuth) │ +│ request_file_upload_url → PUT bytes → transform_files │ +│ → check_transform_status → get_transform_results │ +│ ⇒ a signed download_url per file (Element JSON) │ +│ │ +│ ② rag-ingest-mcp (local, STDIO — this package) │ +│ load_transform_output(download_url, collection) │ +│ fetches out of band — the JSON never enters agent context │ +│ embeds locally iff Transform didn't │ +│ search(query, collection) ⇒ nearest chunks + scores │ +└────────────────────────────────────────────────────────────────────┘ + │ + ▼ your local vector DB (the corpus) +``` + +The agent drives both servers with ordinary tool calls. A typical session: + +> **You:** Load `~/Contracts/msa-2026.pdf` into a collection called `contracts`, +> then tell me what the termination clause says. +> +> **Agent:** *(Transform: upload → chunk job → download_url; rag-ingest-mcp: +> load_transform_output → search("termination clause", "contracts") → answers +> from the returned chunks, citing them.)* + +## Where embedding happens — the one decision + +Transform always parses and chunks. Embedding can happen in either of two +places, and `load_transform_output`'s `embed_policy="auto"` (the default) +detects which by sniffing for vectors: + +| | **Local embed** (recommended start) | **Transform embeds** (passthrough) | +|---|---|---| +| Transform `stages` | `{"chunk": {...}}` only | `{"chunk": {...}, "embed": {...}}` | +| Who embeds the corpus | rag-ingest-mcp, with your `OPENAI_API_KEY` | the platform, with its cloud model | +| Model used | `URAG_EMBED_MODEL` (default `text-embedding-3-small`) | whatever the embed stage ran — **pass it as `embed_model` on load** | +| Who embeds each query | rag-ingest-mcp — always, in the collection's pinned space | same | + +Either way, **every query needs a provider key locally** — `search` embeds the +query text with the same model the corpus used. That symmetry is enforced, not +suggested: a collection is pinned to one embedding space `(provider, model, +dimension)` on first load, later loads with a different model are refused, and +queries are always embedded into the pinned space. Mixing models silently +returns garbage rankings; the pin is what makes that impossible. + +## Prerequisites (all guides) + +From the `unstructured-ingest` repo root: + +```bash +pip install -e ".[chroma,openai]" # base + default backend +pip install -r unstructured_ingest/mcp/requirements.txt # fastmcp +export OPENAI_API_KEY=sk-... # query + local-corpus embedding +``` + +You also need access to a **Transform MCP** endpoint (the hosted Unstructured +Platform server, or an SND's). Your MCP host handles its OAuth on first use. + +## Other MCP servers in the mix + +- **Embeddings.** You don't need a separate embedding MCP: Transform's `embed` + stage *is* the cloud-embedding service (passthrough mode), and rag-ingest-mcp + embeds locally otherwise. Today the local path supports `openai`; more + `unstructured_ingest.embed` providers are planned. +- **The DB vendor's own MCP.** Each guide has a "pairing with the vendor MCP" + section. Short version: vendor MCPs are great for *admin and inspection*, but + for chunk-level semantic search over a Transform-built corpus, query through + rag-ingest-mcp — most vendor MCPs either can't embed the query at all + (they expect a precomputed vector) or would embed it with a different model + than the corpus (wrong space, garbage results). rag-ingest-mcp exists + precisely to be the matched-space query front-end. + +## Troubleshooting (all backends) + +- **`SpaceMismatch: collection X was built with model-A … refusing model-B`** — + working as designed. Load into a new collection, or pass the matching + `embed_model`. +- **`could not embed query in the collection's space`** — `OPENAI_API_KEY` is + missing/invalid, or the collection was passthrough-loaded under a model name + your key can't run. Check `list_collections` for the pinned model. +- **Passthrough load pinned the wrong model** — in passthrough, the space is + recorded from `embed_model` (or the server default). If Transform embedded + with something else, dimensions won't match at query time. Re-load with the + correct `embed_model`. +- **Elements loaded = 0** — the source had no vectors and `embed_policy` was + `"passthrough"`. Use `"auto"` or `"local"`. diff --git a/unstructured_ingest/mcp/examples/chroma/README.md b/unstructured_ingest/mcp/examples/chroma/README.md new file mode 100644 index 000000000..03f45085a --- /dev/null +++ b/unstructured_ingest/mcp/examples/chroma/README.md @@ -0,0 +1,116 @@ +# Getting started: agentic RAG with Chroma (the zero-infra default) + +Chroma is rag-ingest-mcp's default backend: an embedded vector store that lives +in a local directory. **Nothing to install beyond the Python extras, nothing to +run, nothing to provision** — the first document you load creates the +collection. Start here unless you have a reason not to. + +## 1. Install + +From the `unstructured-ingest` repo root: + +```bash +pip install -e ".[chroma,openai]" +pip install -r unstructured_ingest/mcp/requirements.txt # fastmcp +``` + +## 2. Register both MCP servers with your agent + +Copy [`mcp-config.json`](mcp-config.json) into your host's config — Claude Code +reads `.mcp.json` in the project root; Claude Desktop reads +`claude_desktop_config.json`. Fill in your OpenAI key and your Transform MCP +URL: + +```json +{ + "mcpServers": { + "rag-ingest-mcp": { + "command": "python", + "args": ["-m", "unstructured_ingest.mcp"], + "env": { + "OPENAI_API_KEY": "sk-...", + "URAG_CHROMA_PATH": "~/.unstructured-rag/chroma" + } + }, + "unstructured-transform": { + "type": "http", + "url": "https:///transform-mcp/mcp" + } + } +} +``` + +That's the entire setup. `URAG_STORE_BACKEND` defaults to `chroma`; +`URAG_CHROMA_PATH` is where your corpus persists (this directory **is** your +RAG database — back it up, or delete it to start over). + +## 3. Run the loop + +Ask your agent, in plain language: + +> Parse `~/Docs/quarterly-report.pdf` with Transform (chunk by title), load it +> into a collection called `reports`, then answer: what were the main revenue +> drivers? + +Behind that one request the agent will: + +1. `request_file_upload_url` → `PUT` the bytes → `transform_files` with + `{"stages": {"chunk": {"strategy": "chunk_by_title"}}}` +2. poll `check_transform_status` until `COMPLETED` +3. `get_transform_results(job_id, output_format="json")` → a `download_url` +4. `load_transform_output(sources=[download_url], collection="reports")` — + rag-ingest-mcp fetches the JSON out of band and, seeing no vectors + (chunk-only), embeds locally with `text-embedding-3-small` +5. `search("main revenue drivers", collection="reports")` → answers from the + top chunks + +Loading the same file again **upserts** (element ids are deterministic), so +re-running a pipeline never duplicates rows. To wipe a corpus, delete +`URAG_CHROMA_PATH` — there is no service to reset. + +## 4. Inspect what landed + +Chroma is embedded, so inspection is a Python one-liner against the same path +(stop your MCP host first, or copy the directory — Chroma allows one writer): + +```bash +python -c " +import chromadb +c = chromadb.PersistentClient(path='$HOME/.unstructured-rag/chroma') +for item in c.list_collections(): + coll = c.get_collection(item.name if hasattr(item, 'name') else item) + print(coll.name, coll.count(), coll.metadata) # metadata = the pinned embedding space + rows = coll.peek(3) + for id, doc in zip(rows['ids'], rows['documents']): + print(' ', id, '|', doc[:70]) +" +``` + +The storage is SQLite underneath, so raw poking works too: + +```bash +sqlite3 ~/.unstructured-rag/chroma/chroma.sqlite3 \ + "SELECT name FROM sqlite_master WHERE type='table';" +``` + +## Pairing with Chroma's own MCP server + +The official [`chroma-mcp`](https://github.com/chroma-core/chroma-mcp) can point +at the same directory, and it's handy for admin-style operations (listing, +counting, deleting collections). **Don't use it for querying this corpus, +though:** its `query` embeds with the collection's persisted embedding function, +and collections written by rag-ingest-mcp (like those written by the ingest +connector) have none — chroma-mcp would fall back to its default local model, +which is a different embedding space than your corpus. Same collection, wrong +vectors, quietly wrong results. `search` on rag-ingest-mcp reads the pinned +space and embeds the query to match — that's the reliable query path. + +## Chroma-specific notes + +- The embedding-space pin lives in the collection's metadata (visible in the + inspection snippet above, alongside `hnsw:space: cosine`). +- One process at a time: Chroma's persistent client locks the directory. If an + inspection script hangs, your MCP host's rag-ingest-mcp subprocess probably + has it open. +- There is no dimension ceiling to worry about — any embedding model's width + works. diff --git a/unstructured_ingest/mcp/examples/chroma/mcp-config.json b/unstructured_ingest/mcp/examples/chroma/mcp-config.json new file mode 100644 index 000000000..be1b09f2c --- /dev/null +++ b/unstructured_ingest/mcp/examples/chroma/mcp-config.json @@ -0,0 +1,18 @@ +{ + "//": "Claude Code: save as .mcp.json in your project root. Claude Desktop: merge into claude_desktop_config.json.", + "mcpServers": { + "rag-ingest-mcp": { + "command": "python", + "args": ["-m", "unstructured_ingest.mcp"], + "env": { + "OPENAI_API_KEY": "sk-...", + "URAG_CHROMA_PATH": "~/.unstructured-rag/chroma" + } + }, + "unstructured-transform": { + "//": "Hosted Transform MCP (or your SND's). OAuth is handled by the MCP host on first use.", + "type": "http", + "url": "https:///transform-mcp/mcp" + } + } +} diff --git a/unstructured_ingest/mcp/examples/pgvector/README.md b/unstructured_ingest/mcp/examples/pgvector/README.md new file mode 100644 index 000000000..b347e61b4 --- /dev/null +++ b/unstructured_ingest/mcp/examples/pgvector/README.md @@ -0,0 +1,145 @@ +# Getting started: agentic RAG with pgvector (Postgres) + +pgvector puts your RAG corpus in Postgres: every collection is a plain table +with a `vector` column, so the same data your agent searches semantically is +also one `psql` away — and one *SQL MCP server* away, which makes this the most +composable backend for agentic workflows (see the pairing section below). + +rag-ingest-mcp auto-provisions everything on first load: the `vector` +extension, the collection's table, a cosine HNSW index, and a small +`rag_ingest_spaces` registry that pins each collection's embedding model. You +bring a running Postgres; no schema work. + +## 1. Install + +From the `unstructured-ingest` repo root: + +```bash +pip install -e ".[openai]" +pip install -r unstructured_ingest/mcp/requirements.txt # fastmcp +pip install "psycopg[binary]" # Postgres driver (psycopg 3) +``` + +## 2. Start Postgres (or use one you have) + +```bash +docker compose up -d # uses this directory's docker-compose.yml +``` + +which runs the official pgvector image: + +```yaml +services: + postgres: + image: pgvector/pgvector:pg17 + environment: { POSTGRES_PASSWORD: postgres } + ports: ["5432:5432"] + volumes: ["pg-data:/var/lib/postgresql/data"] +``` + +Using an existing Postgres instead? Any 13+ instance with the +[pgvector extension](https://github.com/pgvector/pgvector) available works. +The first load runs `CREATE EXTENSION IF NOT EXISTS vector`, which needs a +sufficiently privileged role once; after that an ordinary role that can create +tables is enough. + +## 3. Register both MCP servers with your agent + +Copy [`mcp-config.json`](mcp-config.json) — Claude Code reads `.mcp.json` in +your project root; Claude Desktop reads `claude_desktop_config.json`: + +```json +{ + "mcpServers": { + "rag-ingest-mcp": { + "command": "python", + "args": ["-m", "unstructured_ingest.mcp"], + "env": { + "OPENAI_API_KEY": "sk-...", + "URAG_STORE_BACKEND": "pgvector", + "URAG_PG_DSN": "postgresql://postgres:postgres@localhost:5432/postgres" + } + }, + "unstructured-transform": { + "type": "http", + "url": "https:///transform-mcp/mcp" + } + } +} +``` + +One pgvector-specific rule: **collection names must be valid Postgres +identifiers** (letters, digits, underscores; start with a letter) because the +collection *is* the table — `contracts_2026` works, `contracts-2026` is +refused. + +## 4. Run the loop + +Ask your agent: + +> Parse `~/Contracts/msa-2026.pdf` and `~/Contracts/sow-2026.pdf` with +> Transform, chunk by title, load them into `contracts`, then: what are our +> termination and liability terms? + +The agent uploads both files through Transform, runs the chunk job, hands the +resulting `download_url`s to `load_transform_output(collection="contracts")` +(no vectors in chunk-only output, so rag-ingest-mcp embeds locally), and +answers from `search("termination and liability terms", "contracts")`. + +Re-loading a source upserts via `ON CONFLICT` on the deterministic element id — +re-running a pipeline updates rows in place instead of duplicating them. + +## 5. Inspect what landed — it's just Postgres + +```bash +docker compose exec postgres psql -U postgres # or: psql "$URAG_PG_DSN" +``` + +```sql +\dt -- one table per collection + rag_ingest_spaces +\d contracts -- id, record_id, text, metadata jsonb, embedding vector(1536) + -- plus the cosine HNSW index +SELECT * FROM rag_ingest_spaces; -- which embedding model each collection is pinned to + +-- chunks per source document +SELECT metadata->>'filename' AS file, count(*) FROM contracts GROUP BY 1; + +-- full-text-ish triage without touching vectors +SELECT left(text, 80) FROM contracts WHERE text ILIKE '%termination%'; +``` + +The `metadata` column holds the element's flattened metadata (filename, page +number, element type, …) as JSONB, so all of Postgres's JSON operators apply. + +## Pairing with a Postgres MCP server — the agentic-SQL combo + +This is where pgvector earns its keep. Register a general-purpose Postgres MCP +(e.g. [`postgres-mcp`](https://github.com/crystaldba/postgres-mcp)) alongside, +pointed at the same database, and your agent gets **two complementary views of +one corpus**: + +- **rag-ingest-mcp** answers *"what does the corpus say about X?"* — semantic + search, with the query embedded in the collection's pinned space. +- **the Postgres MCP** answers *"what's in the corpus?"* — counts, filters, + grouping, joins against your other tables: *"how many chunks per document?"*, + *"which contracts mention indemnification on pages 1–3?"* (JSONB filters), + *"join chunk counts against my `vendors` table"*. + +The division of labor matters: a generic SQL MCP can't do the semantic side +correctly on its own — pgvector search needs a query *vector*, and the SQL +server has no way to produce one in your corpus's embedding space. That +cloud-embed front-end is exactly the gap rag-ingest-mcp fills. Let each do what +it's for: don't hand-write `ORDER BY embedding <=> ...` through the SQL MCP, +and don't ask rag-ingest-mcp for analytics. + +## pgvector-specific notes + +- **Dimension cap for indexing:** pgvector's HNSW index tops out at 2,000 + dimensions. The default model (`text-embedding-3-small`, 1536) fits. Wider + spaces (e.g. `text-embedding-3-large`, 3072) still load and search correctly + but scan sequentially — fine for small corpora; for large ones, prefer a + ≤2,000-dim model or convert the column to `halfvec` yourself. +- Scores are cosine similarity (`1 - (embedding <=> query)`), comparable with + the other backends. +- `rag_ingest_spaces` is the source of truth for `list_collections`; dropping a + collection means dropping its table *and* its row there. diff --git a/unstructured_ingest/mcp/examples/pgvector/docker-compose.yml b/unstructured_ingest/mcp/examples/pgvector/docker-compose.yml new file mode 100644 index 000000000..3a66729d1 --- /dev/null +++ b/unstructured_ingest/mcp/examples/pgvector/docker-compose.yml @@ -0,0 +1,16 @@ +# Postgres with the pgvector extension for rag-ingest-mcp +# (URAG_PG_DSN=postgresql://postgres:postgres@localhost:5432/postgres). +# Data persists in the named volume across container restarts. +services: + postgres: + image: pgvector/pgvector:pg17 + container_name: rag-ingest-postgres + environment: + POSTGRES_PASSWORD: postgres + ports: + - "5432:5432" + volumes: + - pg-data:/var/lib/postgresql/data + +volumes: + pg-data: diff --git a/unstructured_ingest/mcp/examples/pgvector/mcp-config.json b/unstructured_ingest/mcp/examples/pgvector/mcp-config.json new file mode 100644 index 000000000..fbc4fa19d --- /dev/null +++ b/unstructured_ingest/mcp/examples/pgvector/mcp-config.json @@ -0,0 +1,19 @@ +{ + "//": "Claude Code: save as .mcp.json in your project root. Claude Desktop: merge into claude_desktop_config.json. Optionally add a Postgres MCP (e.g. postgres-mcp) on the same DSN for SQL-side analytics over the corpus.", + "mcpServers": { + "rag-ingest-mcp": { + "command": "python", + "args": ["-m", "unstructured_ingest.mcp"], + "env": { + "OPENAI_API_KEY": "sk-...", + "URAG_STORE_BACKEND": "pgvector", + "URAG_PG_DSN": "postgresql://postgres:postgres@localhost:5432/postgres" + } + }, + "unstructured-transform": { + "//": "Hosted Transform MCP (or your SND's). OAuth is handled by the MCP host on first use.", + "type": "http", + "url": "https:///transform-mcp/mcp" + } + } +} diff --git a/unstructured_ingest/mcp/examples/qdrant/README.md b/unstructured_ingest/mcp/examples/qdrant/README.md new file mode 100644 index 000000000..54c5b77a6 --- /dev/null +++ b/unstructured_ingest/mcp/examples/qdrant/README.md @@ -0,0 +1,153 @@ +# Getting started: agentic RAG with Qdrant + +Qdrant gives you a dedicated vector database with a web dashboard for browsing +your corpus. rag-ingest-mcp supports it in two modes — pick one: + +- **Embedded local mode** — Qdrant runs *inside* the rag-ingest-mcp process + against a local directory. Zero infrastructure, like Chroma. Good for + getting started. +- **Server mode** — a Qdrant instance (docker below, or one you already run). + Survives independent of the MCP process, and you get the dashboard. + +## 1. Install + +From the `unstructured-ingest` repo root: + +```bash +pip install -e ".[qdrant,openai]" +pip install -r unstructured_ingest/mcp/requirements.txt # fastmcp +``` + +## 2a. Embedded mode (no service) + +Nothing to start. In the MCP config below, set only: + +```json +"URAG_STORE_BACKEND": "qdrant", +"URAG_QDRANT_PATH": "~/.unstructured-rag/qdrant" +``` + +Note: embedded mode locks its directory to one process and has no dashboard. +When you outgrow it, switch to server mode — collections don't migrate +automatically, but re-running your load pipeline rebuilds them. + +## 2b. Server mode (docker) + +```bash +docker compose up -d # uses this directory's docker-compose.yml +``` + +which runs: + +```yaml +services: + qdrant: + image: qdrant/qdrant:latest + ports: ["6333:6333"] + volumes: ["qdrant-data:/qdrant/storage"] +``` + +Then in the MCP config: + +```json +"URAG_STORE_BACKEND": "qdrant", +"URAG_QDRANT_URL": "http://localhost:6333" +``` + +(`URAG_QDRANT_URL` takes precedence over the path; add `URAG_QDRANT_API_KEY` +for a secured instance.) + +## 3. Register both MCP servers with your agent + +Copy [`mcp-config.json`](mcp-config.json) — Claude Code reads `.mcp.json` in +your project root; Claude Desktop reads `claude_desktop_config.json`: + +```json +{ + "mcpServers": { + "rag-ingest-mcp": { + "command": "python", + "args": ["-m", "unstructured_ingest.mcp"], + "env": { + "OPENAI_API_KEY": "sk-...", + "URAG_STORE_BACKEND": "qdrant", + "URAG_QDRANT_URL": "http://localhost:6333" + } + }, + "unstructured-transform": { + "type": "http", + "url": "https:///transform-mcp/mcp" + } + } +} +``` + +## 4. Run the loop + +Ask your agent: + +> Parse everything in `~/Research/papers/` with Transform (chunk by title), +> load it into a collection called `papers`, then: what methods do these papers +> use for evaluation? + +The agent will drive Transform (upload → chunk job → `download_url` per file), +hand every URL to `load_transform_output(sources=[...], collection="papers")` — +directories of Element JSON also work as a source — and then call +`search("evaluation methods", collection="papers")` and synthesize an answer +from the returned chunks. + +Re-loading the same sources upserts (deterministic ids); a load with a +different embedding model than the collection was built with is refused. + +## 5. Inspect what landed + +Server mode has the nicest inspection story of the three backends — a built-in +web UI: + +```bash +open http://localhost:6333/dashboard +``` + +Or the REST API: + +```bash +# collections (you'll also see __rag_ingest_meta__ — see below) +curl -s http://localhost:6333/collections | python -m json.tool + +# config + point count: 1536-dim cosine for the default embed model +curl -s http://localhost:6333/collections/papers | python -m json.tool + +# browse points with payloads (text, flattened element metadata) +curl -s -X POST http://localhost:6333/collections/papers/points/scroll \ + -H 'content-type: application/json' \ + -d '{"limit": 3, "with_payload": true}' | python -m json.tool +``` + +**About `__rag_ingest_meta__`:** Qdrant collections carry no free-form +metadata, so rag-ingest-mcp pins each collection's embedding space as one point +in this reserved meta collection. Inspect the pins with a scroll on it; don't +write to it or delete it (deleting orphans the space guard — collections would +re-pin on their next load). + +Also normal: `indexed_vectors_count: 0` on small collections. Qdrant serves +search from raw segments until a collection is large enough (~10k+ vectors) to +be worth building the HNSW graph. + +## Pairing with Qdrant's own MCP server + +The official [`mcp-server-qdrant`](https://github.com/qdrant/mcp-server-qdrant) +embeds queries with **FastEmbed** (local ONNX models). It *cannot* reproduce +the OpenAI vectors your corpus was embedded with, so pointing it at a +rag-ingest-mcp collection gives confidently wrong results — the query lands in +a different embedding space than the corpus. Use it only for corpora it built +itself. For this corpus, `search` on rag-ingest-mcp is the correct query path +(it re-embeds each query in the collection's pinned space); the dashboard and +REST API cover inspection. + +## Qdrant-specific notes + +- Distance is cosine, configured at collection creation from the pinned space; + Qdrant natively rejects wrong-dimension upserts as a second line of defense. +- Scores from `search` are cosine similarity as Qdrant reports it — directly + comparable with the other backends' scores. +- No practical dimension ceiling for common embedding models (65k max). diff --git a/unstructured_ingest/mcp/examples/qdrant/docker-compose.yml b/unstructured_ingest/mcp/examples/qdrant/docker-compose.yml new file mode 100644 index 000000000..4995ef479 --- /dev/null +++ b/unstructured_ingest/mcp/examples/qdrant/docker-compose.yml @@ -0,0 +1,13 @@ +# Qdrant server for rag-ingest-mcp's server mode (URAG_QDRANT_URL=http://localhost:6333). +# Storage persists in the named volume across container restarts. +services: + qdrant: + image: qdrant/qdrant:latest + container_name: rag-ingest-qdrant + ports: + - "6333:6333" + volumes: + - qdrant-data:/qdrant/storage + +volumes: + qdrant-data: diff --git a/unstructured_ingest/mcp/examples/qdrant/mcp-config.json b/unstructured_ingest/mcp/examples/qdrant/mcp-config.json new file mode 100644 index 000000000..757e3c5e8 --- /dev/null +++ b/unstructured_ingest/mcp/examples/qdrant/mcp-config.json @@ -0,0 +1,19 @@ +{ + "//": "Claude Code: save as .mcp.json in your project root. Claude Desktop: merge into claude_desktop_config.json. For embedded mode, drop URAG_QDRANT_URL and set URAG_QDRANT_PATH instead.", + "mcpServers": { + "rag-ingest-mcp": { + "command": "python", + "args": ["-m", "unstructured_ingest.mcp"], + "env": { + "OPENAI_API_KEY": "sk-...", + "URAG_STORE_BACKEND": "qdrant", + "URAG_QDRANT_URL": "http://localhost:6333" + } + }, + "unstructured-transform": { + "//": "Hosted Transform MCP (or your SND's). OAuth is handled by the MCP host on first use.", + "type": "http", + "url": "https:///transform-mcp/mcp" + } + } +} diff --git a/unstructured_ingest/mcp/fetch.py b/unstructured_ingest/mcp/fetch.py new file mode 100644 index 000000000..83515e115 --- /dev/null +++ b/unstructured_ingest/mcp/fetch.py @@ -0,0 +1,64 @@ +"""Fetch Transform MCP output (Element JSON) for local ingestion. + +The agent hands this server a ``download_url`` from ``get_transform_results`` (or +a local path it already saved). The server pulls the bytes itself so the — +potentially large — JSON never passes through the agent's context, which is the +whole point of doing the load out of band. The download is streamed with a byte +cap so a runaway or wrong URL can't exhaust memory. +""" + +from __future__ import annotations + +import json +import tempfile +from pathlib import Path +from typing import Any +from urllib.parse import urlparse + +import httpx + +# Generous enough for a large multi-thousand-page render on a slow link, tight +# enough that an unreachable host fails rather than hanging the tool call. +_FETCH_TIMEOUT_S = 120.0 + +# Spill to disk past this while streaming, so a large body is bounded in RSS. +_SPOOL_THRESHOLD_BYTES = 8 * 1024 * 1024 + + +def fetch_elements(source: str, *, max_bytes: int) -> list[dict[str, Any]]: + """Return the Element JSON array at ``source`` (http(s) URL or local path).""" + parsed = urlparse(source) + scheme = parsed.scheme.lower() + if scheme in ("http", "https"): + elements = _stream_json(source, max_bytes=max_bytes) + elif scheme in ("", "file"): + path = Path(parsed.path if scheme == "file" else source) + raw = path.read_bytes() + if len(raw) > max_bytes: + raise ValueError(f"{source}: {len(raw)} bytes exceeds cap of {max_bytes}") + elements = json.loads(raw) + else: + raise ValueError(f"unsupported source scheme {scheme!r}: {source}") + + if not isinstance(elements, list): + raise ValueError( + f"{source}: expected a JSON array of elements, got {type(elements).__name__}. " + "Fetch get_transform_results with output_format='json'." + ) + return elements + + +def _stream_json(url: str, *, max_bytes: int) -> Any: + total = 0 + # SpooledTemporaryFile keeps a small body in memory and spills a large one to + # disk; either way the parse reads from a single seekable handle. + with tempfile.SpooledTemporaryFile(max_size=_SPOOL_THRESHOLD_BYTES, mode="w+b") as spool: + with httpx.stream("GET", url, timeout=_FETCH_TIMEOUT_S, follow_redirects=True) as resp: + resp.raise_for_status() + for chunk in resp.iter_bytes(): + total += len(chunk) + if total > max_bytes: + raise ValueError(f"{url}: download exceeds cap of {max_bytes} bytes") + spool.write(chunk) + spool.seek(0) + return json.load(spool) diff --git a/unstructured_ingest/mcp/requirements.txt b/unstructured_ingest/mcp/requirements.txt new file mode 100644 index 000000000..2cf4be942 --- /dev/null +++ b/unstructured_ingest/mcp/requirements.txt @@ -0,0 +1,8 @@ +# The only dependency this server adds on top of unstructured-ingest itself. +# The embedders, the Chroma connector, httpx, and pydantic all come from +# unstructured-ingest and its extras — install the package with: +# +# pip install -e ".[chroma,openai]" # from the unstructured-ingest repo root +# pip install -r unstructured_ingest/mcp/requirements.txt +# +fastmcp>=2.0 diff --git a/unstructured_ingest/mcp/server.py b/unstructured_ingest/mcp/server.py new file mode 100644 index 000000000..c63e9cd68 --- /dev/null +++ b/unstructured_ingest/mcp/server.py @@ -0,0 +1,257 @@ +"""FastMCP STDIO server: land Transform MCP output into a local vector store. + +Runs as a local subprocess of the MCP host (Claude Desktop / Claude Code) +alongside the remote Transform MCP. The agent parses/chunks/(optionally)embeds +via Transform, then calls ``load_transform_output`` here with a download_url; +this server fetches the Element JSON out of band, embeds it locally iff Transform +did not, and upserts it into a local vector store — Chroma by default, Qdrant or +pgvector via ``URAG_STORE_BACKEND``. ``search`` then embeds a query into that +collection's own space and returns the nearest chunks. + +The embedding step is the only knob: Transform embeds (passthrough) OR this +server embeds (local), selected per load. Either way a collection records exactly +one embedding space and every query is embedded into it, so corpus and query +vectors are always comparable. +""" + +from __future__ import annotations + +from pathlib import Path +from urllib.parse import urlparse + +from fastmcp import FastMCP + +from unstructured_ingest.data_types.file_data import FileData, SourceIdentifiers +from unstructured_ingest.mcp.config import Config +from unstructured_ingest.mcp.embedding import build_encoder +from unstructured_ingest.mcp.fetch import fetch_elements +from unstructured_ingest.mcp.stores import ( + EmbeddingSpace, + SpaceMismatch, + VectorStore, + build_store, +) + +mcp = FastMCP("rag-ingest-mcp") + +_VALID_POLICIES = ("auto", "passthrough", "local") + + +def _store(cfg: Config) -> VectorStore: + return build_store(cfg) + + +def _source_identity(source: str) -> str: + """A stable per-document id so re-loading the same source upserts. + + A Transform download URL's last path segment is the deterministic artifact id + (job + file + format), which makes a stable key; a local path uses its name. + """ + name = Path(urlparse(source).path).name + return name or "transform-output" + + +def _expand_sources(sources: list[str]) -> tuple[list[str], list[dict]]: + """Expand any local directory into the .json files under it; pass others through. + + Transform MCP output is commonly a directory of per-document Element JSON, so + pointing this tool at that directory ingests every ``*.json`` under it + (recursively). URLs and individual file paths pass through unchanged. Returns + the work list plus per-source errors for directories that held no JSON. + """ + work: list[str] = [] + errors: list[dict] = [] + for source in sources: + parsed = urlparse(source) + if parsed.scheme in ("", "file"): + path = Path(parsed.path if parsed.scheme == "file" else source) + if path.is_dir(): + found = sorted(str(p) for p in path.rglob("*.json")) + if found: + work.extend(found) + else: + errors.append({"source": source, "error": "no .json files found in directory"}) + continue + work.append(source) + return work, errors + + +@mcp.tool +def load_transform_output( + sources: list[str], + collection: str, + embed_policy: str = "auto", + embed_model: str | None = None, +) -> dict: + """Load Transform MCP Element JSON into a local vector-store collection. + + Fetches each source out of band (the JSON never enters the agent context), + embeds locally when needed, and upserts into ``collection``. + + Args: + sources: One or more ``download_url``s from ``get_transform_results`` + (call it with ``output_format="json"``), local file paths, or a + local directory — a directory is expanded to every ``*.json`` under + it, so you can point this at a whole Transform output folder. + collection: Target collection. Created on first load and pinned to the + embedding space used; later loads with a different model are refused. + embed_policy: ``"auto"`` (default) embeds locally only when elements + arrive without vectors; ``"passthrough"`` requires Transform-supplied + vectors; ``"local"`` always (re-)embeds here. + embed_model: Override the embed model. In ``passthrough`` this records + which model Transform used (defaults to the server's configured + model); in ``local`` it selects the model to embed with. + + Returns: + A summary: the collection, the policy, the pinned embedding space, and a + per-source count of rows loaded. Errors are reported per source, not + raised, so one bad URL doesn't sink a batch. + """ + if embed_policy not in _VALID_POLICIES: + return { + "error": f"embed_policy must be one of {list(_VALID_POLICIES)}; got {embed_policy!r}" + } + + cfg = Config.from_env() + try: + store = _store(cfg) + except ValueError as exc: + return {"error": str(exc)} + + model = embed_model or cfg.embed_model + encoder = None # built lazily and reused across sources that embed locally + pinned_space: EmbeddingSpace | None = None + + work_items, files = _expand_sources(sources) + + for source in work_items: + try: + elements = fetch_elements(source, max_bytes=cfg.max_fetch_bytes) + except Exception as exc: + files.append({"source": source, "error": str(exc)}) + continue + + has_vectors = any(el.get("embeddings") for el in elements) + embed_here = embed_policy == "local" or (embed_policy == "auto" and not has_vectors) + + if embed_here: + try: + if encoder is None: + encoder = build_encoder( + provider=cfg.embed_provider, + model=model, + api_key=cfg.openai_api_key, + base_url=cfg.openai_base_url, + batch_size=cfg.embed_batch_size, + ) + elements = encoder.embed_documents(elements) + except Exception as exc: + files.append({"source": source, "error": f"local embedding failed: {exc}"}) + continue + space = EmbeddingSpace(cfg.embed_provider, model, encoder.dimension) + else: + first_vector = next((el["embeddings"] for el in elements if el.get("embeddings")), None) + if first_vector is None: + files.append( + { + "source": source, + "error": "no vectors present; retry this source with embed_policy='local'", + } + ) + continue + space = EmbeddingSpace(cfg.embed_provider, model, len(first_vector)) + + try: + store.ensure_space(collection, space) + except SpaceMismatch as exc: + files.append({"source": source, "error": str(exc)}) + continue + + file_data = FileData( + identifier=_source_identity(source), + connector_type="transform_mcp", + source_identifiers=SourceIdentifiers( + filename=_source_identity(source), fullpath=source + ), + ) + loaded = store.write(collection, elements, file_data) + pinned_space = space + files.append({"source": source, "loaded": loaded, "elements": len(elements)}) + + return { + "collection": collection, + "embed_policy": embed_policy, + "space": pinned_space.as_metadata() if pinned_space else None, + "files": files, + } + + +@mcp.tool +def search(query: str, collection: str, k: int = 5) -> dict: + """Similarity-search a collection, embedding the query in its own space. + + Reads the collection's pinned (provider, model) and embeds ``query`` with the + same one, so the query and the corpus are always comparable — the reason a + passthrough (Transform-embedded) collection still needs the matching provider + key available to this server. + + Args: + query: The natural-language query text. + collection: A collection previously built by ``load_transform_output``. + k: Number of nearest chunks to return. + + Returns: + The matches (each with a cosine ``score``, ``text``, and ``metadata``), + plus the model they were compared in. + """ + cfg = Config.from_env() + try: + store = _store(cfg) + space = store.space_of(collection) + except Exception as exc: + return {"error": f"unknown or unbuilt collection {collection!r}: {exc}"} + + try: + encoder = build_encoder( + provider=space.provider or cfg.embed_provider, + model=space.model, + api_key=cfg.openai_api_key, + base_url=cfg.openai_base_url, + batch_size=cfg.embed_batch_size, + ) + query_vector = encoder.embed_query(query) + except Exception as exc: + return {"error": f"could not embed query in the collection's space: {exc}"} + + return { + "collection": collection, + "model": space.model, + "matches": store.search(collection, query_vector, k), + } + + +@mcp.tool +def list_collections() -> dict: + """List local collections and the embedding space each is pinned to.""" + cfg = Config.from_env() + try: + store = _store(cfg) + except ValueError as exc: + return {"error": str(exc)} + out = [] + for name in store.collections(): + try: + space = store.space_of(name) + out.append({"collection": name, "model": space.model, "dimension": space.dimension}) + except Exception: + out.append({"collection": name}) + return {"collections": out} + + +def main() -> None: + """Run the server over STDIO (the transport MCP hosts launch it with).""" + mcp.run() + + +if __name__ == "__main__": + main() diff --git a/unstructured_ingest/mcp/stores/__init__.py b/unstructured_ingest/mcp/stores/__init__.py new file mode 100644 index 000000000..bba007764 --- /dev/null +++ b/unstructured_ingest/mcp/stores/__init__.py @@ -0,0 +1,55 @@ +"""rag-ingest-mcp storage backends. + +One :class:`VectorStore` implementation per backend, selected by +``URAG_STORE_BACKEND``. Backend client libraries are imported lazily inside +each store, so an install only needs the extras for the backend it uses. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from unstructured_ingest.mcp.stores.base import ( + EmbeddingSpace, + SpaceMismatch, + VectorStore, +) + +if TYPE_CHECKING: + from unstructured_ingest.mcp.config import Config + +BACKENDS = ("chroma", "qdrant", "pgvector") + + +def build_store(cfg: "Config") -> VectorStore: + """The configured backend's store, or ValueError with the fix spelled out.""" + if cfg.store_backend == "chroma": + from unstructured_ingest.mcp.stores.chroma import ChromaStore + + return ChromaStore(cfg.chroma_path) + if cfg.store_backend == "qdrant": + from unstructured_ingest.mcp.stores.qdrant import QdrantStore + + return QdrantStore(path=cfg.qdrant_path, url=cfg.qdrant_url, api_key=cfg.qdrant_api_key) + if cfg.store_backend == "pgvector": + if not cfg.pg_dsn: + raise ValueError( + "the pgvector backend needs URAG_PG_DSN set " + "(e.g. postgresql://user:pass@localhost:5432/db)" + ) + from unstructured_ingest.mcp.stores.pgvector import PgvectorStore + + return PgvectorStore(cfg.pg_dsn) + raise ValueError( + f"unknown store backend {cfg.store_backend!r}; " + f"URAG_STORE_BACKEND must be one of {list(BACKENDS)}" + ) + + +__all__ = [ + "BACKENDS", + "EmbeddingSpace", + "SpaceMismatch", + "VectorStore", + "build_store", +] diff --git a/unstructured_ingest/mcp/stores/base.py b/unstructured_ingest/mcp/stores/base.py new file mode 100644 index 000000000..5f6c1fa0b --- /dev/null +++ b/unstructured_ingest/mcp/stores/base.py @@ -0,0 +1,99 @@ +"""The backend contract every rag-ingest-mcp vector store implements. + +A backend is "supported" when it upholds five invariants behind this interface: + +1. **Space pin** — a collection records the embedding space (provider, model, + dimension) its vectors were produced by; a load with a different space is + refused (:class:`SpaceMismatch`). Vectors from different models are not + comparable, so mixing them silently ruins retrieval. +2. **Auto-provision** — the first load creates the collection/table from the + pinned space; no pre-provisioning step is asked of the user. +3. **Idempotent upsert** — element ids are deterministic, so re-loading a + source overwrites rows rather than duplicating them. +4. **Cosine scoring** — ``search`` returns ``score`` as cosine similarity + (higher is better), whatever the backend's native distance convention. +5. **Connector-shaped writes** — rows are conformed by the corresponding + ``unstructured-ingest`` stager, so what lands in the store matches what the + ingest pipeline's own connector would produce. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Any + +from unstructured_ingest.data_types.file_data import FileData + + +class SpaceMismatch(Exception): + """Raised when a load would mix a second embedding space into a collection.""" + + +@dataclass +class EmbeddingSpace: + """The (provider, model, dimension) a collection's vectors were produced by.""" + + provider: str + model: str + dimension: int + + def as_metadata(self) -> dict[str, Any]: + return { + "embed_provider": self.provider, + "embed_model": self.model, + "embed_dim": self.dimension, + } + + @classmethod + def from_metadata(cls, meta: dict[str, Any] | None) -> "EmbeddingSpace | None": + if not meta or "embed_model" not in meta: + return None + return cls( + provider=meta.get("embed_provider", ""), + model=meta["embed_model"], + dimension=int(meta["embed_dim"]), + ) + + def mismatch_error(self, collection: str, incoming: "EmbeddingSpace") -> SpaceMismatch: + return SpaceMismatch( + f"collection {collection!r} was built with {self.model} " + f"({self.dimension}-dim); refusing to load {incoming.model} " + f"({incoming.dimension}-dim) into it. Use a different collection name." + ) + + def conflicts_with(self, other: "EmbeddingSpace") -> bool: + return (self.model, self.dimension) != (other.model, other.dimension) + + +class VectorStore(ABC): + """Write, guard, and search one backend's collections.""" + + @abstractmethod + def ensure_space(self, collection: str, space: EmbeddingSpace) -> None: + """Create ``collection`` pinned to ``space``, or verify an existing one. + + Raises :class:`SpaceMismatch` if the collection already holds vectors + from a different model/dimension. + """ + + @abstractmethod + def write(self, collection: str, elements: list[dict], file_data: FileData) -> int: + """Upsert ``elements`` into ``collection``; return the rows written. + + Rows without a vector (an element that arrived un-embedded) are dropped — + they can be neither stored nor searched. ``ensure_space`` must have run + first so the collection exists with the intended dimension and metric. + """ + + @abstractmethod + def space_of(self, collection: str) -> EmbeddingSpace: + """The space ``collection`` is pinned to; raises if it has none.""" + + @abstractmethod + def search(self, collection: str, query_vector: list[float], k: int) -> list[dict[str, Any]]: + """The ``k`` nearest rows as ``{score, text, metadata}`` dicts.""" + + @abstractmethod + def collections(self) -> list[str]: + """Names of the collections this store holds.""" diff --git a/unstructured_ingest/mcp/stores/chroma.py b/unstructured_ingest/mcp/stores/chroma.py new file mode 100644 index 000000000..74331c2d5 --- /dev/null +++ b/unstructured_ingest/mcp/stores/chroma.py @@ -0,0 +1,109 @@ +"""Chroma-backed vector store — the zero-config default backend. + +Writes go through ``unstructured-ingest``'s Chroma stager (Element JSON -> Chroma +row) and its columnar transform, so what lands on disk matches what the ingest +pipeline's own Chroma connector would produce. Reads use the Chroma client +directly, because ingest connectors are write-only destinations — retrieval is +this server's own concern. + +The embedding space is pinned in the collection's own metadata. Chroma is the +default backend because it is embedded (a directory, no service) and both the +collection and its metadata are created in one call. +""" + +from __future__ import annotations + +from typing import Any + +from unstructured_ingest.data_types.file_data import FileData +from unstructured_ingest.mcp.stores.base import EmbeddingSpace, VectorStore +from unstructured_ingest.processes.connectors.chroma import ( + ChromaUploader, + ChromaUploadStager, +) + +# Cosine matches the near-unit-length vectors these embedding models emit; set +# once at collection creation and never changed for the life of the collection. +_DISTANCE = "cosine" + + +class ChromaStore(VectorStore): + def __init__(self, path: str) -> None: + self._path = path + self._client = None + + def _get_client(self): + # One PersistentClient per process/path: Chroma keeps a single + # SQLite-backed handle, so the space guard, the upsert, and search all + # share this one rather than opening competing handles to the same dir. + if self._client is None: + import chromadb + + self._client = chromadb.PersistentClient(path=self._path) + return self._client + + def _get_collection_or_none(self, name: str): + try: + return self._get_client().get_collection(name) + except Exception: + # Chroma raises a version-specific "not found" error; any failure to + # open is treated as absent, and a genuine problem re-surfaces at the + # create/upsert that follows. + return None + + def ensure_space(self, collection: str, space: EmbeddingSpace) -> None: + coll = self._get_collection_or_none(collection) + if coll is None: + self._get_client().create_collection( + collection, + metadata={"hnsw:space": _DISTANCE, **space.as_metadata()}, + ) + return + existing = EmbeddingSpace.from_metadata(coll.metadata) + if existing is not None and existing.conflicts_with(space): + raise existing.mismatch_error(collection, space) + + def write(self, collection: str, elements: list[dict], file_data: FileData) -> int: + stager = ChromaUploadStager() + rows = [ + stager.conform_dict(element_dict=element, file_data=file_data) for element in elements + ] + rows = [row for row in rows if row.get("embedding")] + if not rows: + return 0 + # prepare_chroma_list is ingest's own row -> parallel-lists transform, + # reused so this upsert is identical to the connector's. + payload = ChromaUploader.prepare_chroma_list(tuple(rows)) + coll = self._get_client().get_collection(collection) + coll.upsert( + ids=payload["ids"], + documents=payload["documents"], + embeddings=payload["embeddings"], + metadatas=payload["metadatas"], + ) + return len(rows) + + def space_of(self, collection: str) -> EmbeddingSpace: + coll = self._get_client().get_collection(collection) + space = EmbeddingSpace.from_metadata(coll.metadata) + if space is None: + raise ValueError(f"collection {collection!r} has no recorded embedding space") + return space + + def search(self, collection: str, query_vector: list[float], k: int) -> list[dict[str, Any]]: + coll = self._get_client().get_collection(collection) + res = coll.query( + query_embeddings=[query_vector], + n_results=max(1, k), + include=["documents", "metadatas", "distances"], + ) + matches: list[dict[str, Any]] = [] + for doc, meta, dist in zip(res["documents"][0], res["metadatas"][0], res["distances"][0]): + # Collection is cosine, so similarity = 1 - distance. + matches.append({"score": 1.0 - dist, "text": doc, "metadata": meta}) + return matches + + def collections(self) -> list[str]: + items = self._get_client().list_collections() + # Chroma versions return either Collection objects or bare names here. + return [item if isinstance(item, str) else item.name for item in items] diff --git a/unstructured_ingest/mcp/stores/pgvector.py b/unstructured_ingest/mcp/stores/pgvector.py new file mode 100644 index 000000000..f00f3faca --- /dev/null +++ b/unstructured_ingest/mcp/stores/pgvector.py @@ -0,0 +1,336 @@ +"""pgvector-backed vector store: a collection per table in your Postgres. + +Rows are conformed by ``unstructured-ingest``'s SQL stager (Element JSON -> one +flat row per element, metadata lifted to top-level keys), then **fit to the +target table's columns** and landed in the Unstructured-compliant schema the SQL +pgvector connector uses: ``id`` (a deterministic uuid5, the upsert key), plus the +element columns the table actually has — ``type``, ``record_id``, +``element_id``, ``text``, the ``embeddings`` vector, and any wider metadata +columns (``filename``, ``page_number``, …) present in the table. + +The key property this buys: the table this store writes is the *same* shape a +downstream reader that speaks the Unstructured schema expects (e.g. Dragon Lab's +pgvector index). So the collection can be pre-created by that reader and this +store loads straight into it, or — standalone — this store auto-provisions an +equivalent table from the pinned space. Either way the write path stays +pgvector-aware: the vector is sent as a JSON array literal cast to the column's +own vector type (``vector`` or ``halfvec``), which the generic SQL uploader does +not do. + +The space pin lives in one small table (``rag_ingest_spaces``), which is also +the registry of collections this store owns; the element table carries no model +metadata of its own. pgvector caps HNSW-indexable vectors at 2,000 dims; wider +spaces still work but search is a sequential scan. + +Uses psycopg 3; connect with a DSN like ``postgresql://user:pass@host:5432/db``. +The ``vector`` extension is created if missing (needs a privileged role once). +""" + +from __future__ import annotations + +import json +import re +import uuid +from typing import Any + +from unstructured_ingest.data_types.file_data import FileData +from unstructured_ingest.mcp.stores.base import EmbeddingSpace, VectorStore +from unstructured_ingest.processes.connectors.sql.sql import SQLUploadStager + +_SPACES_TABLE = "rag_ingest_spaces" + +# pgvector refuses HNSW/ivfflat indexes above this dimension; wider vectors are +# stored and searched brute-force instead. +_MAX_INDEXABLE_DIM = 2000 + +# The vector-typed column this store writes/searches, and the Postgres vector +# types it may be declared as. ``embeddings`` (plural) is the Unstructured schema +# column name — the same one Dragon Lab's pgvector index uses. +_VECTOR_COLUMN = "embeddings" +_VECTOR_UDTS = ("vector", "halfvec", "sparsevec") + +# Columns that are never returned as search-result metadata (they are the id, +# the document text, and the vector itself). +_NON_METADATA_COLUMNS = ("id", "text", _VECTOR_COLUMN) + +# The schema auto-provisioned standalone (no pre-created table). Column order is +# the DDL order; ``id`` is the deterministic uuid5 upsert key. Every element +# column is nullable so a sparse element (e.g. one with no page number) still +# loads; a downstream reader that requires a subset simply sees the columns it +# needs. The vector column is appended separately with the pinned dimension. +_AUTOPROVISION_COLUMNS = ("id", "type", "record_id", "element_id", "text", "filename", "page_number") + +# Postgres identifier: quoted names could carry anything, but keeping collection +# names plain identifiers keeps them usable from psql and other tools too. +_IDENTIFIER = re.compile(r"[A-Za-z_][A-Za-z0-9_]{0,62}\Z") + +# Parses a declared vector type like ``halfvec(1536)`` into its dimension. +_VECTOR_TYPE_DIM = re.compile(r"\((\d+)\)") + + +def _vector_literal(vector: list[float]) -> str: + return json.dumps(vector) + + +def _check_name(collection: str) -> str: + if not _IDENTIFIER.match(collection): + raise ValueError( + f"collection {collection!r} is not usable as a Postgres table name; " + "use letters, digits and underscores, starting with a letter (max 63 chars)" + ) + return collection + + +class PgvectorStore(VectorStore): + def __init__(self, dsn: str) -> None: + self._dsn = dsn + + def _connect(self): + import psycopg + + return psycopg.connect(self._dsn) + + def _ensure_spaces_table(self, conn) -> None: + conn.execute("CREATE EXTENSION IF NOT EXISTS vector") + conn.execute( + f""" + CREATE TABLE IF NOT EXISTS {_SPACES_TABLE} ( + collection TEXT PRIMARY KEY, + provider TEXT NOT NULL, + model TEXT NOT NULL, + dimension INTEGER NOT NULL + ) + """ + ) + + def _read_pin(self, conn, collection: str) -> EmbeddingSpace | None: + row = conn.execute( + f"SELECT provider, model, dimension FROM {_SPACES_TABLE} WHERE collection = %s", + (collection,), + ).fetchone() + if row is None: + return None + return EmbeddingSpace(provider=row[0], model=row[1], dimension=row[2]) + + def _table_exists(self, conn, collection: str) -> bool: + return conn.execute("SELECT to_regclass(%s)", (collection,)).fetchone()[0] is not None + + def _column_types(self, conn, collection: str) -> dict[str, str]: + """Map of column name -> udt_name (e.g. ``text``, ``uuid``, ``halfvec``).""" + rows = conn.execute( + """ + SELECT column_name, udt_name + FROM information_schema.columns + WHERE table_name = %s + ORDER BY ordinal_position + """, + (collection,), + ).fetchall() + return {name: udt for name, udt in rows} + + def _vector_column(self, column_types: dict[str, str]) -> tuple[str, str]: + """The vector column and its udt; prefers ``embeddings``, else any vector-typed column.""" + if column_types.get(_VECTOR_COLUMN) in _VECTOR_UDTS: + return _VECTOR_COLUMN, column_types[_VECTOR_COLUMN] + for name, udt in column_types.items(): + if udt in _VECTOR_UDTS: + return name, udt + raise ValueError( + f"table has no vector-typed column (one of {list(_VECTOR_UDTS)}); " + "is this an Unstructured-compliant pgvector table?" + ) + + def _existing_vector_dim(self, conn, collection: str, vector_column: str) -> int | None: + """Declared dimension of an existing vector column, or None if unpardimensioned.""" + row = conn.execute( + "SELECT format_type(atttypid, atttypmod) FROM pg_attribute " + "WHERE attrelid = %s::regclass AND attname = %s", + (collection, vector_column), + ).fetchone() + if not row or not row[0]: + return None + match = _VECTOR_TYPE_DIM.search(row[0]) + return int(match.group(1)) if match else None + + def ensure_space(self, collection: str, space: EmbeddingSpace) -> None: + from psycopg import sql + + _check_name(collection) + with self._connect() as conn: + self._ensure_spaces_table(conn) + existing = self._read_pin(conn, collection) + if existing is not None: + if existing.conflicts_with(space): + raise existing.mismatch_error(collection, space) + return + + if self._table_exists(conn, collection): + # Pre-created by a downstream reader (e.g. Dragon Lab). Adopt it: + # verify the vector dimension agrees, then record the pin. Never + # recreate — the reader owns the table's schema. + column_types = self._column_types(conn, collection) + vector_column, _ = self._vector_column(column_types) + declared_dim = self._existing_vector_dim(conn, collection, vector_column) + if declared_dim is not None and declared_dim != space.dimension: + raise space.mismatch_error( + collection, + EmbeddingSpace(space.provider, space.model, declared_dim), + ) + else: + # Standalone: auto-provision an Unstructured-compliant table. + columns = ", ".join(f"{col} TEXT" for col in _AUTOPROVISION_COLUMNS if col != "id") + conn.execute( + sql.SQL( + "CREATE TABLE IF NOT EXISTS {table} " + "(id UUID PRIMARY KEY, " + columns + ", " + "{vector_column} vector({dim}))" + ).format( + table=sql.Identifier(collection), + vector_column=sql.Identifier(_VECTOR_COLUMN), + dim=sql.Literal(space.dimension), + ) + ) + if space.dimension <= _MAX_INDEXABLE_DIM: + conn.execute( + sql.SQL( + "CREATE INDEX IF NOT EXISTS {index} ON {table} " + "USING hnsw ({vector_column} vector_cosine_ops)" + ).format( + index=sql.Identifier(f"{collection}_embeddings_hnsw"), + table=sql.Identifier(collection), + vector_column=sql.Identifier(_VECTOR_COLUMN), + ) + ) + + conn.execute( + f"INSERT INTO {_SPACES_TABLE} (collection, provider, model, dimension) " + "VALUES (%s, %s, %s, %s)", + (collection, space.provider, space.model, space.dimension), + ) + + def _coerce(self, value: Any, udt: str): + """Adapt a staged value to the target column's type (excluding the vector).""" + from psycopg.types.json import Json + + if value is None: + return None + if udt == "uuid": + return value if isinstance(value, uuid.UUID) else uuid.UUID(str(value)) + if udt == "jsonb" or udt == "json": + return Json(value, dumps=lambda obj: json.dumps(obj, default=str)) + if udt.startswith("_"): # a Postgres array type (e.g. _text); psycopg adapts a list + return value if isinstance(value, list) else [value] + # text / varchar / numeric / … — JSON-encode structured values, else stringify + # non-strings so they land in a TEXT column without a type-mismatch error. + if isinstance(value, (dict, list)): + return json.dumps(value, default=str) + return value if isinstance(value, str) else str(value) + + def write(self, collection: str, elements: list[dict], file_data: FileData) -> int: + from psycopg import sql + + _check_name(collection) + stager = SQLUploadStager() + rows = [ + stager.conform_dict(element_dict=element, file_data=file_data) for element in elements + ] + + with self._connect() as conn: + column_types = self._column_types(conn, collection) + vector_column, vector_udt = self._vector_column(column_types) + + # Fit to the table: keep only staged keys that are real columns, and + # union across rows so every INSERT has a uniform column list. + present = {key for row in rows for key in row} + target_columns = [ + col for col in column_types if col in present and col != vector_column + ] + if "id" not in target_columns and "id" in column_types: + target_columns.insert(0, "id") + + params = [] + for row in rows: + vector = row.get("embeddings") + if not vector: + continue # un-embedded elements can be neither stored nor searched + values = [self._coerce(row.get(col), column_types[col]) for col in target_columns] + values.append(_vector_literal(vector)) + params.append(tuple(values)) + if not params: + return 0 + + insert_columns = target_columns + [vector_column] + placeholders = [sql.SQL("%s")] * len(target_columns) + [ + sql.SQL("%s::") + sql.SQL(vector_udt) + ] + update_columns = [col for col in insert_columns if col != "id"] + statement = sql.SQL( + "INSERT INTO {table} ({columns}) VALUES ({placeholders}) " + "ON CONFLICT (id) DO UPDATE SET {assignments}" + ).format( + table=sql.Identifier(collection), + columns=sql.SQL(", ").join(sql.Identifier(col) for col in insert_columns), + placeholders=sql.SQL(", ").join(placeholders), + assignments=sql.SQL(", ").join( + sql.SQL("{col} = EXCLUDED.{col}").format(col=sql.Identifier(col)) + for col in update_columns + ), + ) + with conn.cursor() as cursor: + cursor.executemany(statement, params) + return len(params) + + def space_of(self, collection: str) -> EmbeddingSpace: + _check_name(collection) + with self._connect() as conn: + self._ensure_spaces_table(conn) + space = self._read_pin(conn, collection) + if space is None: + raise ValueError(f"collection {collection!r} has no recorded embedding space") + return space + + def search(self, collection: str, query_vector: list[float], k: int) -> list[dict[str, Any]]: + from psycopg import sql + + _check_name(collection) + literal = _vector_literal(query_vector) + with self._connect() as conn: + column_types = self._column_types(conn, collection) + vector_column, vector_udt = self._vector_column(column_types) + metadata_columns = [ + col for col in column_types if col not in _NON_METADATA_COLUMNS and col != vector_column + ] + + selected = [sql.Identifier("text")] + [sql.Identifier(c) for c in metadata_columns] + statement = sql.SQL( + "SELECT {selected}, 1 - ({vector} <=> {q}::{udt}) AS score " + "FROM {table} ORDER BY {vector} <=> {q}::{udt} LIMIT {k}" + ).format( + selected=sql.SQL(", ").join(selected), + vector=sql.Identifier(vector_column), + q=sql.Literal(literal), + udt=sql.SQL(vector_udt), + table=sql.Identifier(collection), + k=sql.Literal(max(1, k)), + ) + rows = conn.execute(statement).fetchall() + + results = [] + for row in rows: + text = row[0] + score = row[-1] + metadata = { + col: value + for col, value in zip(metadata_columns, row[1:-1]) + if value is not None + } + results.append({"score": score, "text": text, "metadata": metadata}) + return results + + def collections(self) -> list[str]: + with self._connect() as conn: + self._ensure_spaces_table(conn) + rows = conn.execute( + f"SELECT collection FROM {_SPACES_TABLE} ORDER BY collection" + ).fetchall() + return [row[0] for row in rows] diff --git a/unstructured_ingest/mcp/stores/qdrant.py b/unstructured_ingest/mcp/stores/qdrant.py new file mode 100644 index 000000000..c1499645c --- /dev/null +++ b/unstructured_ingest/mcp/stores/qdrant.py @@ -0,0 +1,159 @@ +"""Qdrant-backed vector store: embedded local mode or a running server. + +Writes are conformed by ``unstructured-ingest``'s Qdrant stager (Element JSON -> +point with ``id``/``vector``/``payload``), so what lands in Qdrant matches what +the ingest pipeline's own connector would produce. With ``path`` set the store +runs Qdrant's embedded local mode (a directory, no service); with ``url`` set it +talks to a server. + +Qdrant collections carry no free-form metadata, so the embedding-space pin lives +in a reserved meta collection: one point per data collection, keyed +deterministically by collection name, its payload holding the space. The data +collection's own vector params (size, cosine) are derived from the pinned space +at creation, so the dimension is also enforced natively on every upsert. +""" + +from __future__ import annotations + +from typing import Any +from uuid import NAMESPACE_DNS, uuid5 + +from unstructured_ingest.data_types.file_data import FileData +from unstructured_ingest.mcp.stores.base import EmbeddingSpace, VectorStore +from unstructured_ingest.processes.connectors.qdrant.local import ( + LocalQdrantUploadStager, + LocalQdrantUploadStagerConfig, +) + +# Holds one space-pin point per data collection; never searched, never listed. +_META_COLLECTION = "__rag_ingest_meta__" + +# Payload keys the stager adds that are storage detail, not search-result metadata. +_PAYLOAD_INTERNAL_KEYS = ("text", "element_serialized") + + +def _meta_point_id(collection: str) -> str: + return str(uuid5(NAMESPACE_DNS, f"rag-ingest-space:{collection}")) + + +class QdrantStore(VectorStore): + def __init__( + self, + path: str | None = None, + url: str | None = None, + api_key: str | None = None, + ) -> None: + if not path and not url: + raise ValueError("QdrantStore needs a local path or a server url") + self._path = path + self._url = url + self._api_key = api_key + self._client = None + + def _get_client(self): + # One client per process: local mode locks its directory, so the space + # guard, the upsert, and search must share a single handle. + if self._client is None: + from qdrant_client import QdrantClient + + if self._url: + self._client = QdrantClient(url=self._url, api_key=self._api_key) + else: + self._client = QdrantClient(path=self._path) + return self._client + + def _read_pin(self, collection: str) -> EmbeddingSpace | None: + client = self._get_client() + if not client.collection_exists(_META_COLLECTION): + return None + points = client.retrieve(_META_COLLECTION, ids=[_meta_point_id(collection)]) + if not points: + return None + return EmbeddingSpace.from_metadata(points[0].payload) + + def _write_pin(self, collection: str, space: EmbeddingSpace) -> None: + from qdrant_client import models + + client = self._get_client() + if not client.collection_exists(_META_COLLECTION): + client.create_collection( + _META_COLLECTION, + vectors_config=models.VectorParams(size=1, distance=models.Distance.COSINE), + ) + client.upsert( + _META_COLLECTION, + points=[ + models.PointStruct( + id=_meta_point_id(collection), + vector=[0.0], + payload={"collection": collection, **space.as_metadata()}, + ) + ], + wait=True, + ) + + def ensure_space(self, collection: str, space: EmbeddingSpace) -> None: + from qdrant_client import models + + existing = self._read_pin(collection) + if existing is not None: + if existing.conflicts_with(space): + raise existing.mismatch_error(collection, space) + return + client = self._get_client() + if not client.collection_exists(collection): + client.create_collection( + collection, + vectors_config=models.VectorParams( + size=space.dimension, distance=models.Distance.COSINE + ), + ) + # A pre-existing collection without a pin (created outside this server) + # adopts the incoming space, mirroring the Chroma backend's behavior. + self._write_pin(collection, space) + + def write(self, collection: str, elements: list[dict], file_data: FileData) -> int: + from qdrant_client import models + + stager = LocalQdrantUploadStager(upload_stager_config=LocalQdrantUploadStagerConfig()) + rows = [ + stager.conform_dict(element_dict=element, file_data=file_data) for element in elements + ] + rows = [row for row in rows if row.get("vector")] + if not rows: + return 0 + self._get_client().upsert( + collection, + points=[models.PointStruct(**row) for row in rows], + wait=True, + ) + return len(rows) + + def space_of(self, collection: str) -> EmbeddingSpace: + space = self._read_pin(collection) + if space is None: + raise ValueError(f"collection {collection!r} has no recorded embedding space") + return space + + def search(self, collection: str, query_vector: list[float], k: int) -> list[dict[str, Any]]: + res = self._get_client().query_points( + collection, + query=query_vector, + limit=max(1, k), + with_payload=True, + ) + matches: list[dict[str, Any]] = [] + for point in res.points: + payload = point.payload or {} + metadata = { + key: value for key, value in payload.items() if key not in _PAYLOAD_INTERNAL_KEYS + } + # Qdrant returns cosine scores as similarity already (higher = closer). + matches.append( + {"score": point.score, "text": payload.get("text"), "metadata": metadata} + ) + return matches + + def collections(self) -> list[str]: + response = self._get_client().get_collections() + return [c.name for c in response.collections if c.name != _META_COLLECTION]