From 76b1a0c253763bfc41b67ea2e2c490d7fa4b4958 Mon Sep 17 00:00:00 2001 From: Daniel Schofield Date: Mon, 6 Jul 2026 11:13:32 -0500 Subject: [PATCH 1/9] Add local RAG MCP server (unstructured_ingest/mcp) A FastMCP STDIO server that lands Unstructured Transform MCP output into a local vector store and serves similarity search. Transform parses/chunks/optionally embeds on the platform; this server fetches the Element JSON out of band, embeds it locally when Transform did not, upserts into a local Chroma collection, and answers queries embedded into that same space. Reuses unstructured-ingest's Chroma stager/uploader and embedding encoders, so the write path matches the ingest pipeline's. A collection is pinned to one embedding space (provider, model, dimension); later loads with a different model are refused, keeping corpus and query vectors comparable. Tools: load_transform_output (embed_policy auto/passthrough/local), search, list_collections. Co-Authored-By: Claude Opus 4.8 (1M context) --- unstructured_ingest/mcp/.env.example | 16 ++ unstructured_ingest/mcp/README.md | 110 +++++++++ unstructured_ingest/mcp/__init__.py | 9 + unstructured_ingest/mcp/__main__.py | 6 + unstructured_ingest/mcp/config.py | 52 ++++ unstructured_ingest/mcp/embedding.py | 52 ++++ .../mcp/example_mcp_config.json | 19 ++ unstructured_ingest/mcp/fetch.py | 66 +++++ unstructured_ingest/mcp/requirements.txt | 8 + unstructured_ingest/mcp/server.py | 226 ++++++++++++++++++ unstructured_ingest/mcp/store.py | 167 +++++++++++++ 11 files changed, 731 insertions(+) create mode 100644 unstructured_ingest/mcp/.env.example create mode 100644 unstructured_ingest/mcp/README.md create mode 100644 unstructured_ingest/mcp/__init__.py create mode 100644 unstructured_ingest/mcp/__main__.py create mode 100644 unstructured_ingest/mcp/config.py create mode 100644 unstructured_ingest/mcp/embedding.py create mode 100644 unstructured_ingest/mcp/example_mcp_config.json create mode 100644 unstructured_ingest/mcp/fetch.py create mode 100644 unstructured_ingest/mcp/requirements.txt create mode 100644 unstructured_ingest/mcp/server.py create mode 100644 unstructured_ingest/mcp/store.py diff --git a/unstructured_ingest/mcp/.env.example b/unstructured_ingest/mcp/.env.example new file mode 100644 index 000000000..26d6c3b96 --- /dev/null +++ b/unstructured_ingest/mcp/.env.example @@ -0,0 +1,16 @@ +# --- required for embedding (ingest-side embed and every query) --- +OPENAI_API_KEY=sk-... + +# --- local vector store --- +# Directory where the Chroma collections persist. This IS your local RAG corpus. +URAG_CHROMA_PATH=~/.unstructured-rag/chroma + +# --- embedding defaults (must match what Transform used, in passthrough mode) --- +URAG_EMBED_PROVIDER=openai +URAG_EMBED_MODEL=text-embedding-3-large + +# --- 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..9887b2fb6 --- /dev/null +++ b/unstructured_ingest/mcp/README.md @@ -0,0 +1,110 @@ +# unstructured-ingest RAG MCP server + +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, and serves similarity search. It +reuses `unstructured-ingest`'s own embedders and Chroma connector, 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 Chroma directory (your corpus) +``` + +## 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]" +pip install -r unstructured_ingest/mcp/requirements.txt # fastmcp +``` + +## 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_CHROMA_PATH` | `~/.unstructured-rag/chroma` | where the corpus persists | +| `URAG_EMBED_MODEL` | `text-embedding-3-large` | must match Transform's model in passthrough | +| `URAG_EMBED_PROVIDER` | `openai` | embed provider | + +## Wire it into your MCP host + +See `example_mcp_config.json`. 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`. + +## 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. + +## Notes & limits + +- **Chroma is the default** because it auto-creates collections and takes rows as + dicts. LanceDB / Qdrant-local connectors exist in `unstructured-ingest` but + need a **pre-provisioned, typed table** (the LanceDB uploader drops columns not + already in the schema), so they aren't wired as the zero-config default here. +- **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. +- **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..dce4b3371 --- /dev/null +++ b/unstructured_ingest/mcp/config.py @@ -0,0 +1,52 @@ +"""Environment-driven configuration for the ingest RAG MCP server. + +Every setting has a default that makes the common case — a local Chroma store +embedding with OpenAI ``text-embedding-3-large`` — work with only an +``OPENAI_API_KEY`` set. The embed defaults intentionally mirror the Transform +MCP's own default embedder: in passthrough mode the corpus vectors were produced +by Transform, so the query side must default to the same provider/model to stay +comparable. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass + +# The Transform MCP's default embedder. Kept in sync so that a passthrough load +# (Transform embedded) and a query embedded by this server land in one space. +DEFAULT_EMBED_PROVIDER = "openai" +DEFAULT_EMBED_MODEL = "text-embedding-3-large" + +# 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 + 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") + ), + 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..7a3936e3a --- /dev/null +++ b/unstructured_ingest/mcp/embedding.py @@ -0,0 +1,52 @@ +"""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/example_mcp_config.json b/unstructured_ingest/mcp/example_mcp_config.json new file mode 100644 index 000000000..d8bae299d --- /dev/null +++ b/unstructured_ingest/mcp/example_mcp_config.json @@ -0,0 +1,19 @@ +{ + "//": "Claude Code .mcp.json / Claude Desktop claude_desktop_config.json. Registers both servers in one host: the remote Transform MCP (parse/chunk/embed) and this local STDIO server (land JSON into local storage + search).", + "mcpServers": { + "unstructured-ingest-rag": { + "command": "python", + "args": ["-m", "unstructured_ingest.mcp"], + "env": { + "OPENAI_API_KEY": "sk-...", + "URAG_CHROMA_PATH": "~/.unstructured-rag/chroma", + "URAG_EMBED_MODEL": "text-embedding-3-large" + } + }, + "unstructured-transform": { + "//": "Remote Transform MCP. Point at your host (SaaS / SND / local platform-api). 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..60d44381c --- /dev/null +++ b/unstructured_ingest/mcp/fetch.py @@ -0,0 +1,66 @@ +"""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..bf9022ceb --- /dev/null +++ b/unstructured_ingest/mcp/server.py @@ -0,0 +1,226 @@ +"""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 Chroma collection. ``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.store import ChromaStore, EmbeddingSpace, SpaceMismatch + +mcp = FastMCP("unstructured-ingest-rag") + +_VALID_POLICIES = ("auto", "passthrough", "local") + + +def _store(cfg: Config) -> ChromaStore: + if cfg.store_backend != "chroma": + raise ValueError( + f"store backend {cfg.store_backend!r} is not wired in this server yet; " + "use 'chroma' (LanceDB/Qdrant-local need a pre-provisioned table)." + ) + return ChromaStore(cfg.chroma_path) + + +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" + + +@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"``) or local file paths. + 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 + files: list[dict] = [] + + for source in sources: + 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/store.py b/unstructured_ingest/mcp/store.py new file mode 100644 index 000000000..bf51001ed --- /dev/null +++ b/unstructured_ingest/mcp/store.py @@ -0,0 +1,167 @@ +"""Chroma-backed vector store for the ingest RAG MCP server. + +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. + +A collection's embedding space (provider, model, dimension) is pinned in its +metadata on first write and verified on every subsequent write. Vectors from +different models are not comparable, so mixing two models into one collection +silently ruins retrieval; refusing the second model is the guard that prevents +it. Chroma is the default because it auto-creates collections and accepts rows +as plain dicts; LanceDB/Qdrant-local need a pre-provisioned, typed table. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from unstructured_ingest.data_types.file_data import FileData +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 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"]), + ) + + +class ChromaStore: + 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: + """Create ``collection`` pinned to ``space``, or verify an existing one. + + Raises :class:`SpaceMismatch` if the collection already holds vectors + from a different model/dimension. + """ + 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.model, existing.dimension) != ( + space.model, + space.dimension, + ): + raise SpaceMismatch( + f"collection {collection!r} was built with {existing.model} " + f"({existing.dimension}-dim); refusing to load {space.model} " + f"({space.dimension}-dim) into it. Use a different collection name." + ) + + 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 distance metric. + """ + 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] From efd4f11ce53450763d0823b848014363c89260d9 Mon Sep 17 00:00:00 2001 From: Daniel Schofield Date: Mon, 6 Jul 2026 16:56:32 -0500 Subject: [PATCH 2/9] Support directory sources in load_transform_output A source may now be a directory; every *.json file under it is ingested as its own document. --- unstructured_ingest/mcp/server.py | 34 ++++++++++++++++++++++++++++--- 1 file changed, 31 insertions(+), 3 deletions(-) diff --git a/unstructured_ingest/mcp/server.py b/unstructured_ingest/mcp/server.py index bf9022ceb..d2575779a 100644 --- a/unstructured_ingest/mcp/server.py +++ b/unstructured_ingest/mcp/server.py @@ -50,6 +50,31 @@ def _source_identity(source: str) -> str: 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], @@ -64,7 +89,9 @@ def load_transform_output( Args: sources: One or more ``download_url``s from ``get_transform_results`` - (call it with ``output_format="json"``) or local file paths. + (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 @@ -91,9 +118,10 @@ def load_transform_output( model = embed_model or cfg.embed_model encoder = None # built lazily and reused across sources that embed locally pinned_space: EmbeddingSpace | None = None - files: list[dict] = [] - for source in sources: + 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: From 457aceca2faa8897ca98b4538776f2a146d5590c Mon Sep 17 00:00:00 2001 From: Daniel Schofield Date: Mon, 6 Jul 2026 16:57:48 -0500 Subject: [PATCH 3/9] Rename server to rag-ingest-mcp; default embed model to text-embedding-3-small The 1536-dim default fits every candidate backend's ANN-index limits (pgvector caps indexable vectors at 2000 dims). Passthrough loads must name the model Transform actually used via embed_model; documented. --- unstructured_ingest/mcp/.env.example | 2 +- unstructured_ingest/mcp/README.md | 9 +++++++-- unstructured_ingest/mcp/config.py | 19 ++++++++++--------- .../mcp/example_mcp_config.json | 4 ++-- unstructured_ingest/mcp/server.py | 2 +- 5 files changed, 21 insertions(+), 15 deletions(-) diff --git a/unstructured_ingest/mcp/.env.example b/unstructured_ingest/mcp/.env.example index 26d6c3b96..240b6358d 100644 --- a/unstructured_ingest/mcp/.env.example +++ b/unstructured_ingest/mcp/.env.example @@ -7,7 +7,7 @@ URAG_CHROMA_PATH=~/.unstructured-rag/chroma # --- embedding defaults (must match what Transform used, in passthrough mode) --- URAG_EMBED_PROVIDER=openai -URAG_EMBED_MODEL=text-embedding-3-large +URAG_EMBED_MODEL=text-embedding-3-small # --- optional knobs (defaults shown) --- # URAG_STORE_BACKEND=chroma diff --git a/unstructured_ingest/mcp/README.md b/unstructured_ingest/mcp/README.md index 9887b2fb6..ce4329b37 100644 --- a/unstructured_ingest/mcp/README.md +++ b/unstructured_ingest/mcp/README.md @@ -1,4 +1,4 @@ -# unstructured-ingest RAG MCP server +# 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. @@ -64,7 +64,7 @@ Copy `.env.example` and set at least `OPENAI_API_KEY`. Key settings: |---|---|---| | `OPENAI_API_KEY` | — | required; embeds queries (and corpus in local mode) | | `URAG_CHROMA_PATH` | `~/.unstructured-rag/chroma` | where the corpus persists | -| `URAG_EMBED_MODEL` | `text-embedding-3-large` | must match Transform's model in passthrough | +| `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 @@ -104,6 +104,11 @@ In Claude Desktop use `claude_desktop_config.json`; in Claude Code use `.mcp.jso - **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. diff --git a/unstructured_ingest/mcp/config.py b/unstructured_ingest/mcp/config.py index dce4b3371..6356232cc 100644 --- a/unstructured_ingest/mcp/config.py +++ b/unstructured_ingest/mcp/config.py @@ -1,11 +1,11 @@ -"""Environment-driven configuration for the ingest RAG MCP server. +"""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-large`` — work with only an -``OPENAI_API_KEY`` set. The embed defaults intentionally mirror the Transform -MCP's own default embedder: in passthrough mode the corpus vectors were produced -by Transform, so the query side must default to the same provider/model to stay -comparable. +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 @@ -13,10 +13,11 @@ import os from dataclasses import dataclass -# The Transform MCP's default embedder. Kept in sync so that a passthrough load -# (Transform embedded) and a query embedded by this server land in one space. +# 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-large" +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. diff --git a/unstructured_ingest/mcp/example_mcp_config.json b/unstructured_ingest/mcp/example_mcp_config.json index d8bae299d..8e2afcc72 100644 --- a/unstructured_ingest/mcp/example_mcp_config.json +++ b/unstructured_ingest/mcp/example_mcp_config.json @@ -1,13 +1,13 @@ { "//": "Claude Code .mcp.json / Claude Desktop claude_desktop_config.json. Registers both servers in one host: the remote Transform MCP (parse/chunk/embed) and this local STDIO server (land JSON into local storage + search).", "mcpServers": { - "unstructured-ingest-rag": { + "rag-ingest-mcp": { "command": "python", "args": ["-m", "unstructured_ingest.mcp"], "env": { "OPENAI_API_KEY": "sk-...", "URAG_CHROMA_PATH": "~/.unstructured-rag/chroma", - "URAG_EMBED_MODEL": "text-embedding-3-large" + "URAG_EMBED_MODEL": "text-embedding-3-small" } }, "unstructured-transform": { diff --git a/unstructured_ingest/mcp/server.py b/unstructured_ingest/mcp/server.py index d2575779a..07ce5df3c 100644 --- a/unstructured_ingest/mcp/server.py +++ b/unstructured_ingest/mcp/server.py @@ -26,7 +26,7 @@ from unstructured_ingest.mcp.fetch import fetch_elements from unstructured_ingest.mcp.store import ChromaStore, EmbeddingSpace, SpaceMismatch -mcp = FastMCP("unstructured-ingest-rag") +mcp = FastMCP("rag-ingest-mcp") _VALID_POLICIES = ("auto", "passthrough", "local") From 82b42494101668a7385079193ae2cb80053c3427 Mon Sep 17 00:00:00 2001 From: Daniel Schofield Date: Mon, 6 Jul 2026 17:43:35 -0500 Subject: [PATCH 4/9] Add Qdrant and pgvector backends behind a VectorStore interface The store is now a package: a VectorStore contract (space pin, auto-provision, idempotent upsert, cosine scores, connector-shaped writes) with one implementation per backend, selected by URAG_STORE_BACKEND. - chroma (default): unchanged behavior, moved to stores/chroma.py - qdrant: embedded local mode or a server; space pins live in a reserved meta collection since Qdrant collections carry no metadata - pgvector: psycopg 3; a table per collection with a JSONB metadata column and a cosine HNSW index when the dimension is indexable; pins live in a rag_ingest_spaces table A conformance pytest suite (test/unit/mcp) runs identically against every backend and defines what supported means; a server-level smoke suite drives both embedded backends through FastMCP in-memory. --- test/unit/mcp/__init__.py | 0 test/unit/mcp/test_server.py | 104 +++++++++ test/unit/mcp/test_stores.py | 131 ++++++++++++ unstructured_ingest/mcp/.env.example | 13 +- unstructured_ingest/mcp/README.md | 50 ++++- unstructured_ingest/mcp/config.py | 11 + unstructured_ingest/mcp/embedding.py | 4 +- unstructured_ingest/mcp/fetch.py | 4 +- unstructured_ingest/mcp/server.py | 25 ++- unstructured_ingest/mcp/stores/__init__.py | 55 +++++ unstructured_ingest/mcp/stores/base.py | 99 +++++++++ .../mcp/{store.py => stores/chroma.py} | 82 ++----- unstructured_ingest/mcp/stores/pgvector.py | 202 ++++++++++++++++++ unstructured_ingest/mcp/stores/qdrant.py | 159 ++++++++++++++ 14 files changed, 840 insertions(+), 99 deletions(-) create mode 100644 test/unit/mcp/__init__.py create mode 100644 test/unit/mcp/test_server.py create mode 100644 test/unit/mcp/test_stores.py create mode 100644 unstructured_ingest/mcp/stores/__init__.py create mode 100644 unstructured_ingest/mcp/stores/base.py rename unstructured_ingest/mcp/{store.py => stores/chroma.py} (57%) create mode 100644 unstructured_ingest/mcp/stores/pgvector.py create mode 100644 unstructured_ingest/mcp/stores/qdrant.py 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..51f7cd00d --- /dev/null +++ b/test/unit/mcp/test_server.py @@ -0,0 +1,104 @@ +"""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 +from fastmcp import Client + +from unstructured_ingest.mcp.server import mcp + +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: + 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..ab6d19290 --- /dev/null +++ b/test/unit/mcp/test_stores.py @@ -0,0 +1,131 @@ +"""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 +from unstructured_ingest.mcp.stores.chroma import ChromaStore +from unstructured_ingest.mcp.stores.qdrant import QdrantStore + +# 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": + yield ChromaStore(str(tmp_path / "chroma")) + return + if request.param == "qdrant": + 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") + 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/mcp/.env.example b/unstructured_ingest/mcp/.env.example index 240b6358d..834328155 100644 --- a/unstructured_ingest/mcp/.env.example +++ b/unstructured_ingest/mcp/.env.example @@ -2,9 +2,20 @@ OPENAI_API_KEY=sk-... # --- local vector store --- -# Directory where the Chroma collections persist. This IS your local RAG corpus. +# 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 diff --git a/unstructured_ingest/mcp/README.md b/unstructured_ingest/mcp/README.md index ce4329b37..1d8c6cf45 100644 --- a/unstructured_ingest/mcp/README.md +++ b/unstructured_ingest/mcp/README.md @@ -6,9 +6,10 @@ 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, and serves similarity search. It -reuses `unstructured-ingest`'s own embedders and Chroma connector, so the write -path is identical to the ingest pipeline's. +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) ────────────────────────────────┐ @@ -23,7 +24,7 @@ path is identical to the ingest pipeline's. │ search(query) → nearest chunks │ └───────────────────────────────────────────────────────────────────┘ │ - ▼ local Chroma directory (your corpus) + ▼ local vector store (your corpus): chroma | qdrant | pgvector ``` ## The one design idea: embedding is the only knob @@ -52,10 +53,17 @@ keeps retrieval correct — mixing models silently returns nonsense. From the `unstructured-ingest` repo root (this module lives inside the package): ```bash -pip install -e ".[chroma,openai]" +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: @@ -63,7 +71,12 @@ 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_CHROMA_PATH` | `~/.unstructured-rag/chroma` | where the corpus persists | +| `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 | @@ -95,12 +108,27 @@ In Claude Desktop use `claude_desktop_config.json`; in Claude Code use `.mcp.jso pinned space and return the nearest chunks. - **`list_collections()`** — collections and the model each is pinned to. -## Notes & limits +## Backends -- **Chroma is the default** because it auto-creates collections and takes rows as - dicts. LanceDB / Qdrant-local connectors exist in `unstructured-ingest` but - need a **pre-provisioned, typed table** (the LanceDB uploader drops columns not - already in the schema), so they aren't wired as the zero-config default here. +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. diff --git a/unstructured_ingest/mcp/config.py b/unstructured_ingest/mcp/config.py index 6356232cc..5f9ebd868 100644 --- a/unstructured_ingest/mcp/config.py +++ b/unstructured_ingest/mcp/config.py @@ -28,6 +28,10 @@ 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 @@ -42,6 +46,13 @@ def from_env(cls) -> "Config": 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"), diff --git a/unstructured_ingest/mcp/embedding.py b/unstructured_ingest/mcp/embedding.py index 7a3936e3a..fe19ba43d 100644 --- a/unstructured_ingest/mcp/embedding.py +++ b/unstructured_ingest/mcp/embedding.py @@ -46,7 +46,5 @@ def build_encoder( ) 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 - ) + 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/fetch.py b/unstructured_ingest/mcp/fetch.py index 60d44381c..83515e115 100644 --- a/unstructured_ingest/mcp/fetch.py +++ b/unstructured_ingest/mcp/fetch.py @@ -53,9 +53,7 @@ def _stream_json(url: str, *, max_bytes: int) -> Any: # 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: + 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) diff --git a/unstructured_ingest/mcp/server.py b/unstructured_ingest/mcp/server.py index 07ce5df3c..c63e9cd68 100644 --- a/unstructured_ingest/mcp/server.py +++ b/unstructured_ingest/mcp/server.py @@ -4,8 +4,9 @@ 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 Chroma collection. ``search`` then embeds a -query into that collection's own space and returns the nearest chunks. +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 @@ -24,20 +25,20 @@ 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.store import ChromaStore, EmbeddingSpace, SpaceMismatch +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) -> ChromaStore: - if cfg.store_backend != "chroma": - raise ValueError( - f"store backend {cfg.store_backend!r} is not wired in this server yet; " - "use 'chroma' (LanceDB/Qdrant-local need a pre-provisioned table)." - ) - return ChromaStore(cfg.chroma_path) +def _store(cfg: Config) -> VectorStore: + return build_store(cfg) def _source_identity(source: str) -> str: @@ -107,7 +108,9 @@ def load_transform_output( 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}"} + return { + "error": f"embed_policy must be one of {list(_VALID_POLICIES)}; got {embed_policy!r}" + } cfg = Config.from_env() try: 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/store.py b/unstructured_ingest/mcp/stores/chroma.py similarity index 57% rename from unstructured_ingest/mcp/store.py rename to unstructured_ingest/mcp/stores/chroma.py index bf51001ed..74331c2d5 100644 --- a/unstructured_ingest/mcp/store.py +++ b/unstructured_ingest/mcp/stores/chroma.py @@ -1,4 +1,4 @@ -"""Chroma-backed vector store for the ingest RAG MCP server. +"""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 @@ -6,20 +6,17 @@ directly, because ingest connectors are write-only destinations — retrieval is this server's own concern. -A collection's embedding space (provider, model, dimension) is pinned in its -metadata on first write and verified on every subsequent write. Vectors from -different models are not comparable, so mixing two models into one collection -silently ruins retrieval; refusing the second model is the guard that prevents -it. Chroma is the default because it auto-creates collections and accepts rows -as plain dicts; LanceDB/Qdrant-local need a pre-provisioned, typed table. +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 dataclasses import dataclass 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, @@ -30,37 +27,7 @@ _DISTANCE = "cosine" -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"]), - ) - - -class ChromaStore: +class ChromaStore(VectorStore): def __init__(self, path: str) -> None: self._path = path self._client = None @@ -85,11 +52,6 @@ def _get_collection_or_none(self, name: str): return None 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. - """ coll = self._get_collection_or_none(collection) if coll is None: self._get_client().create_collection( @@ -98,27 +60,13 @@ def ensure_space(self, collection: str, space: EmbeddingSpace) -> None: ) return existing = EmbeddingSpace.from_metadata(coll.metadata) - if existing is not None and (existing.model, existing.dimension) != ( - space.model, - space.dimension, - ): - raise SpaceMismatch( - f"collection {collection!r} was built with {existing.model} " - f"({existing.dimension}-dim); refusing to load {space.model} " - f"({space.dimension}-dim) into it. Use a different collection name." - ) + 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: - """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 distance metric. - """ stager = ChromaUploadStager() rows = [ - stager.conform_dict(element_dict=element, file_data=file_data) - for element in elements + 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: @@ -139,14 +87,10 @@ 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" - ) + 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]]: + 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], @@ -154,9 +98,7 @@ def search( include=["documents", "metadatas", "distances"], ) matches: list[dict[str, Any]] = [] - for doc, meta, dist in zip( - res["documents"][0], res["metadatas"][0], res["distances"][0] - ): + 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 diff --git a/unstructured_ingest/mcp/stores/pgvector.py b/unstructured_ingest/mcp/stores/pgvector.py new file mode 100644 index 000000000..68deba63d --- /dev/null +++ b/unstructured_ingest/mcp/stores/pgvector.py @@ -0,0 +1,202 @@ +"""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 landed in a +schema this store owns: ``id`` (deterministic, the upsert key), ``record_id`` +(the per-source identity), ``text``, ``metadata`` (the stager's row as JSONB), +and ``embedding vector(dim)``. The ingest SQL connector fits rows to whatever +table you provide; here the pinned space is known at first load, so the table — +and its cosine HNSW index — are created from it instead. + +The space pin lives in one small table (``rag_ingest_spaces``), which is also +the registry of collections this store owns. pgvector caps HNSW-indexable +vectors at 2,000 dims; wider spaces still work but search is a sequential scan +(or re-create the index over ``halfvec`` yourself). + +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 +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 +from unstructured_ingest.utils.constants import RECORD_ID_LABEL + +_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 + +# 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") + + +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 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 + conn.execute( + sql.SQL( + """ + CREATE TABLE IF NOT EXISTS {table} ( + id TEXT PRIMARY KEY, + record_id TEXT, + text TEXT, + metadata JSONB, + embedding vector({dim}) + ) + """ + ).format(table=sql.Identifier(collection), 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 (embedding vector_cosine_ops)" + ).format( + index=sql.Identifier(f"{collection}_embedding_hnsw"), + table=sql.Identifier(collection), + ) + ) + conn.execute( + f"INSERT INTO {_SPACES_TABLE} (collection, provider, model, dimension) " + "VALUES (%s, %s, %s, %s)", + (collection, space.provider, space.model, space.dimension), + ) + + def write(self, collection: str, elements: list[dict], file_data: FileData) -> int: + from psycopg import sql + from psycopg.types.json import Json + + _check_name(collection) + stager = SQLUploadStager() + rows = [ + stager.conform_dict(element_dict=element, file_data=file_data) for element in elements + ] + params = [] + for row in rows: + vector = row.pop("embeddings", None) + if not vector: + continue + params.append( + ( + row.pop("id"), + row.pop(RECORD_ID_LABEL, file_data.identifier), + row.pop("text", None), + # The stager's remaining flat row is the metadata; default=str + # absorbs dates and other non-JSON scalars it leaves in place. + Json(row, dumps=lambda obj: json.dumps(obj, default=str)), + _vector_literal(vector), + ) + ) + if not params: + return 0 + statement = sql.SQL( + """ + INSERT INTO {table} (id, record_id, text, metadata, embedding) + VALUES (%s, %s, %s, %s, %s::vector) + ON CONFLICT (id) DO UPDATE SET + record_id = EXCLUDED.record_id, + text = EXCLUDED.text, + metadata = EXCLUDED.metadata, + embedding = EXCLUDED.embedding + """ + ).format(table=sql.Identifier(collection)) + with self._connect() as conn, 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) + statement = sql.SQL( + """ + SELECT text, metadata, 1 - (embedding <=> %s::vector) AS score + FROM {table} + ORDER BY embedding <=> %s::vector + LIMIT %s + """ + ).format(table=sql.Identifier(collection)) + with self._connect() as conn: + rows = conn.execute(statement, (literal, literal, max(1, k))).fetchall() + return [ + {"score": score, "text": text, "metadata": metadata} for text, metadata, score in rows + ] + + 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] From 5eeacd5f430cdeef7cc618810e3912a9ae67750c Mon Sep 17 00:00:00 2001 From: Daniel Schofield Date: Mon, 6 Jul 2026 22:13:27 -0500 Subject: [PATCH 5/9] Add per-backend getting-started guides under mcp/examples An overview of the local agentic RAG loop (Transform MCP + this server) plus one self-contained guide each for Chroma, Qdrant, and pgvector: setup, docker-compose where a service is needed, the agent workflow, CLI inspection of each database, and when the DB vendor's own MCP server does or does not pair safely with a corpus built here. --- unstructured_ingest/mcp/README.md | 6 + unstructured_ingest/mcp/examples/README.md | 110 +++++++++++++ .../mcp/examples/chroma/README.md | 116 +++++++++++++ .../mcp/examples/chroma/mcp-config.json | 18 +++ .../mcp/examples/pgvector/README.md | 145 +++++++++++++++++ .../mcp/examples/pgvector/docker-compose.yml | 16 ++ .../mcp/examples/pgvector/mcp-config.json | 19 +++ .../mcp/examples/qdrant/README.md | 153 ++++++++++++++++++ .../mcp/examples/qdrant/docker-compose.yml | 13 ++ .../mcp/examples/qdrant/mcp-config.json | 19 +++ 10 files changed, 615 insertions(+) create mode 100644 unstructured_ingest/mcp/examples/README.md create mode 100644 unstructured_ingest/mcp/examples/chroma/README.md create mode 100644 unstructured_ingest/mcp/examples/chroma/mcp-config.json create mode 100644 unstructured_ingest/mcp/examples/pgvector/README.md create mode 100644 unstructured_ingest/mcp/examples/pgvector/docker-compose.yml create mode 100644 unstructured_ingest/mcp/examples/pgvector/mcp-config.json create mode 100644 unstructured_ingest/mcp/examples/qdrant/README.md create mode 100644 unstructured_ingest/mcp/examples/qdrant/docker-compose.yml create mode 100644 unstructured_ingest/mcp/examples/qdrant/mcp-config.json diff --git a/unstructured_ingest/mcp/README.md b/unstructured_ingest/mcp/README.md index 1d8c6cf45..1c5bbd8bc 100644 --- a/unstructured_ingest/mcp/README.md +++ b/unstructured_ingest/mcp/README.md @@ -86,6 +86,12 @@ See `example_mcp_config.json`. 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`. +**New here? Start with [`examples/`](examples/README.md)** — a getting-started +guide per backend (Chroma, Qdrant, pgvector) covering setup, docker-compose +where a service is needed, the full agentic loop with Transform MCP, how to +inspect each database from the CLI, 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 → 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" + } + } +} From f6ba18cbb44b060bd8f65b6981e4b6fb3595e2c0 Mon Sep 17 00:00:00 2001 From: Daniel Schofield Date: Tue, 7 Jul 2026 10:45:42 -0500 Subject: [PATCH 6/9] Point the module README at every backend's example config The wire-it-up section now links each backend's getting-started guide and its ready-to-copy host config instead of a single Chroma-only example; the standalone example_mcp_config.json is superseded by the per-backend examples/*/mcp-config.json files. --- unstructured_ingest/mcp/README.md | 26 ++++++++++++------- .../mcp/example_mcp_config.json | 19 -------------- 2 files changed, 17 insertions(+), 28 deletions(-) delete mode 100644 unstructured_ingest/mcp/example_mcp_config.json diff --git a/unstructured_ingest/mcp/README.md b/unstructured_ingest/mcp/README.md index 1c5bbd8bc..53301b354 100644 --- a/unstructured_ingest/mcp/README.md +++ b/unstructured_ingest/mcp/README.md @@ -82,15 +82,23 @@ Copy `.env.example` and set at least `OPENAI_API_KEY`. Key settings: ## Wire it into your MCP host -See `example_mcp_config.json`. 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`. - -**New here? Start with [`examples/`](examples/README.md)** — a getting-started -guide per backend (Chroma, Qdrant, pgvector) covering setup, docker-compose -where a service is needed, the full agentic loop with Transform MCP, how to -inspect each database from the CLI, and when (not) to pair the DB vendor's own -MCP server. +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 diff --git a/unstructured_ingest/mcp/example_mcp_config.json b/unstructured_ingest/mcp/example_mcp_config.json deleted file mode 100644 index 8e2afcc72..000000000 --- a/unstructured_ingest/mcp/example_mcp_config.json +++ /dev/null @@ -1,19 +0,0 @@ -{ - "//": "Claude Code .mcp.json / Claude Desktop claude_desktop_config.json. Registers both servers in one host: the remote Transform MCP (parse/chunk/embed) and this local STDIO server (land JSON into local storage + search).", - "mcpServers": { - "rag-ingest-mcp": { - "command": "python", - "args": ["-m", "unstructured_ingest.mcp"], - "env": { - "OPENAI_API_KEY": "sk-...", - "URAG_CHROMA_PATH": "~/.unstructured-rag/chroma", - "URAG_EMBED_MODEL": "text-embedding-3-small" - } - }, - "unstructured-transform": { - "//": "Remote Transform MCP. Point at your host (SaaS / SND / local platform-api). OAuth is handled by the MCP host on first use.", - "type": "http", - "url": "https:///transform-mcp/mcp" - } - } -} From 18ad412458e057aa058da94930f14419e904e65b Mon Sep 17 00:00:00 2001 From: Daniel Schofield Date: Tue, 7 Jul 2026 11:01:28 -0500 Subject: [PATCH 7/9] Skip rag-ingest-mcp tests when backend extras are absent MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The unit-test CI environment installs only locked base dependencies, so fastmcp, chromadb, and qdrant_client are not importable there. Gate the server suite and each conformance backend with pytest.importorskip — the repo's existing pattern for optional connector dependencies. --- test/unit/mcp/test_server.py | 8 ++++++-- test/unit/mcp/test_stores.py | 9 +++++++-- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/test/unit/mcp/test_server.py b/test/unit/mcp/test_server.py index 51f7cd00d..65ca9b30d 100644 --- a/test/unit/mcp/test_server.py +++ b/test/unit/mcp/test_server.py @@ -12,9 +12,12 @@ import anyio import pytest -from fastmcp import Client -from unstructured_ingest.mcp.server import mcp +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], @@ -25,6 +28,7 @@ @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")) diff --git a/test/unit/mcp/test_stores.py b/test/unit/mcp/test_stores.py index ab6d19290..cbc5ca483 100644 --- a/test/unit/mcp/test_stores.py +++ b/test/unit/mcp/test_stores.py @@ -16,8 +16,6 @@ from unstructured_ingest.data_types.file_data import FileData, SourceIdentifiers from unstructured_ingest.mcp.stores import EmbeddingSpace, SpaceMismatch, VectorStore -from unstructured_ingest.mcp.stores.chroma import ChromaStore -from unstructured_ingest.mcp.stores.qdrant import QdrantStore # Orthogonal 4-dim vectors: cosine ranking is exact, no embedding model needed. TEXTS_TO_VECTORS = { @@ -52,14 +50,21 @@ def make_file_data(identifier: str = "doc-1") -> FileData: @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) From 718bcc9f28090e1e0eccbb8cbd68f14da9271318 Mon Sep 17 00:00:00 2001 From: Daniel Schofield Date: Tue, 7 Jul 2026 11:02:20 -0500 Subject: [PATCH 8/9] Bump version to 1.6.29 with changelog entry for rag-ingest-mcp --- CHANGELOG.md | 6 ++++++ unstructured_ingest/__version__.py | 2 +- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 95b316dea..e2d5b0634 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 (auto-provisions the table and a cosine HNSW index). 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/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 From 46dbd8d61fbea02553e4b151c68f8112fc234d13 Mon Sep 17 00:00:00 2001 From: Daniel Schofield Date: Fri, 10 Jul 2026 18:48:21 -0500 Subject: [PATCH 9/9] fix(rag-ingest-mcp): pgvector store writes the Unstructured schema MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The pgvector VectorStore hand-rolled a narrow table of its own (id text, record_id, text, metadata jsonb, embedding) instead of the Unstructured-compliant schema the SQL connector and downstream readers (e.g. Dragon Lab's pgvector index) expect. A reader pointed at a rag-ingest-loaded table therefore found none of its expected columns (embeddings, type, element_id, …), so the two could not share a table. Rewrite ensure_space/write/search to: - adopt a pre-created table — verify the vector dimension and record the space pin without recreating it — or, standalone, auto-provision an equivalent id/type/record_id/element_id/text/embeddings schema with a cosine HNSW index; - fit staged rows to the table's actual columns (union across rows); - cast the vector to the column's own vector/halfvec type, which the generic SQLUploader does not (it adapts the list to a Postgres array). Chroma and Qdrant stores already reuse their connector stagers and build connector-shaped writes, so they are unchanged. The backend conformance suite (test/unit/mcp/test_stores.py) passes for all three. Co-Authored-By: Claude Opus 4.8 --- CHANGELOG.md | 2 +- unstructured_ingest/mcp/stores/pgvector.py | 272 +++++++++++++++------ 2 files changed, 204 insertions(+), 70 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e2d5b0634..d007830d4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,7 @@ ### 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 (auto-provisions the table and a cosine HNSW index). 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/`. +- **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] diff --git a/unstructured_ingest/mcp/stores/pgvector.py b/unstructured_ingest/mcp/stores/pgvector.py index 68deba63d..f00f3faca 100644 --- a/unstructured_ingest/mcp/stores/pgvector.py +++ b/unstructured_ingest/mcp/stores/pgvector.py @@ -1,17 +1,26 @@ """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 landed in a -schema this store owns: ``id`` (deterministic, the upsert key), ``record_id`` -(the per-source identity), ``text``, ``metadata`` (the stager's row as JSONB), -and ``embedding vector(dim)``. The ingest SQL connector fits rows to whatever -table you provide; here the pinned space is known at first load, so the table — -and its cosine HNSW index — are created from it instead. +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. pgvector caps HNSW-indexable -vectors at 2,000 dims; wider spaces still work but search is a sequential scan -(or re-create the index over ``halfvec`` yourself). +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). @@ -21,12 +30,12 @@ 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 -from unstructured_ingest.utils.constants import RECORD_ID_LABEL _SPACES_TABLE = "rag_ingest_spaces" @@ -34,10 +43,30 @@ # 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) @@ -83,6 +112,46 @@ def _read_pin(self, conn, collection: str) -> EmbeddingSpace | 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 @@ -94,75 +163,121 @@ def ensure_space(self, collection: str, space: EmbeddingSpace) -> None: if existing.conflicts_with(space): raise existing.mismatch_error(collection, space) return - conn.execute( - sql.SQL( - """ - CREATE TABLE IF NOT EXISTS {table} ( - id TEXT PRIMARY KEY, - record_id TEXT, - text TEXT, - metadata JSONB, - embedding vector({dim}) + + 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), ) - """ - ).format(table=sql.Identifier(collection), dim=sql.Literal(space.dimension)) - ) - if space.dimension <= _MAX_INDEXABLE_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 INDEX IF NOT EXISTS {index} ON {table} " - "USING hnsw (embedding vector_cosine_ops)" + "CREATE TABLE IF NOT EXISTS {table} " + "(id UUID PRIMARY KEY, " + columns + ", " + "{vector_column} vector({dim}))" ).format( - index=sql.Identifier(f"{collection}_embedding_hnsw"), 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 - from psycopg.types.json import Json _check_name(collection) stager = SQLUploadStager() rows = [ stager.conform_dict(element_dict=element, file_data=file_data) for element in elements ] - params = [] - for row in rows: - vector = row.pop("embeddings", None) - if not vector: - continue - params.append( - ( - row.pop("id"), - row.pop(RECORD_ID_LABEL, file_data.identifier), - row.pop("text", None), - # The stager's remaining flat row is the metadata; default=str - # absorbs dates and other non-JSON scalars it leaves in place. - Json(row, dumps=lambda obj: json.dumps(obj, default=str)), - _vector_literal(vector), - ) + + 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 + ), ) - if not params: - return 0 - statement = sql.SQL( - """ - INSERT INTO {table} (id, record_id, text, metadata, embedding) - VALUES (%s, %s, %s, %s, %s::vector) - ON CONFLICT (id) DO UPDATE SET - record_id = EXCLUDED.record_id, - text = EXCLUDED.text, - metadata = EXCLUDED.metadata, - embedding = EXCLUDED.embedding - """ - ).format(table=sql.Identifier(collection)) - with self._connect() as conn, conn.cursor() as cursor: - cursor.executemany(statement, params) + with conn.cursor() as cursor: + cursor.executemany(statement, params) return len(params) def space_of(self, collection: str) -> EmbeddingSpace: @@ -179,19 +294,38 @@ def search(self, collection: str, query_vector: list[float], k: int) -> list[dic _check_name(collection) literal = _vector_literal(query_vector) - statement = sql.SQL( - """ - SELECT text, metadata, 1 - (embedding <=> %s::vector) AS score - FROM {table} - ORDER BY embedding <=> %s::vector - LIMIT %s - """ - ).format(table=sql.Identifier(collection)) with self._connect() as conn: - rows = conn.execute(statement, (literal, literal, max(1, k))).fetchall() - return [ - {"score": score, "text": text, "metadata": metadata} for text, metadata, score in rows - ] + 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: