Migrate docsite embed/search to Postgres + pgvector (WIP)

.env.example

@@ -23,11 +23,17 @@
 ANTHROPIC_API_KEY=sk-YOUR-API-KEY-HERE
 
 OPENAI_API_KEY=sk-YOUR-API-KEY-HERE
-PINECONE_KEY=YOUR-API-KEY-HERE
+PINECONE_API_KEY=YOUR-API-KEY-HERE
 POSTGRES_URL=postgresql://localhost:5432/apollo_dev
 SENTRY_DSN=YOUR-API-KEY-HERE
 GITHUB_TOKEN=KEY
 
+# Which backend serves docsite search reads: 'pinecone' (default) or 'postgres'.
+DOCSITE_SEARCH_BACKEND=pinecone
+
+# Database for the Postgres docsite integration suite
+POSTGRES_TEST_URL=postgresql://postgres:postgres@127.0.0.1:5433/postgres
+
 # Langfuse observability
 LANGFUSE_SECRET_KEY=sk-lf-...
 LANGFUSE_PUBLIC_KEY=pk-lf-...

pyproject.toml

@@ -26,6 +26,8 @@ psycopg2-binary = "^2.9.10"
 langfuse = "^4.14.1"
 opentelemetry-instrumentation-anthropic = "^0.62.1"
 opentelemetry-instrumentation-threading = "0.65b0"
+pgvector = "^0.5.0"
+pandas = "^2.2"
 
 [tool.poetry.group.dev]
 optional = false
@@ -50,6 +52,7 @@ testpaths = [
     "services/job_chat/tests",
     "services/latest_adaptors/tests",
     "services/search_docsite/tests",
+    "services/embed_docsite/tests",
     "services/tools",
 ]
 

services/db_migrations.py

@@ -0,0 +1,64 @@
+"""Versioned schema migrations for the Python-owned docs database (POSTGRES_URL).
+
+Applies .sql files in lexical order, records applied filenames so re-runs are a
+no-op, and takes an advisory lock so concurrent starters queue.
+
+The tracking table (_migrations_docs) and lock key are distinct from the
+TypeScript runner's in platform/src/db/migrate.ts, because APOLLO_CLIENTS_DB_URL
+falls back to POSTGRES_URL locally and both runners can target one database.
+"""
+
+from pathlib import Path
+
+from util import create_logger
+
+logger = create_logger("db_migrations")
+
+MIGRATIONS_DIR = Path(__file__).parent / "migrations"
+
+# Distinct from the TypeScript runner's 8314_2025 so the two never block each other.
+MIGRATION_LOCK_KEY = 8314_2026
+
+CREATE_TRACKING_TABLE_SQL = """
+CREATE TABLE IF NOT EXISTS _migrations_docs (
+    filename   TEXT PRIMARY KEY,
+    applied_at TIMESTAMPTZ NOT NULL DEFAULT now()
+)
+"""
+
+
+def _migration_files():
+    """Every .sql file in the migrations directory, in lexical order."""
+    if not MIGRATIONS_DIR.is_dir():
+        return []
+    return sorted(MIGRATIONS_DIR.glob("*.sql"))
+
+
+def run_migrations(conn):
+    """Apply any migrations not yet recorded. Returns the count applied this run.
+
+    Everything happens in one transaction: the advisory lock is held for its
+    duration, so a racing process waits and then sees the migrations already
+    recorded rather than colliding on CREATE TABLE.
+    """
+    files = _migration_files()
+
+    with conn.cursor() as cur:
+        cur.execute("SELECT pg_advisory_xact_lock(%s)", (MIGRATION_LOCK_KEY,))
+        cur.execute(CREATE_TRACKING_TABLE_SQL)
+
+        cur.execute("SELECT filename FROM _migrations_docs")
+        already_applied = {row[0] for row in cur.fetchall()}
+
+        pending = [f for f in files if f.name not in already_applied]
+        for path in pending:
+            logger.info(f"Applying migration {path.name}")
+            cur.execute(path.read_text(encoding="utf-8"))
+            cur.execute("INSERT INTO _migrations_docs (filename) VALUES (%s)", (path.name,))
+
+    conn.commit()
+
+    if pending:
+        logger.info(f"Applied {len(pending)} migration(s)")
+
+    return len(pending)

services/embed_docsite/README.md

@@ -1,10 +1,13 @@
 ## Embed Docsite (RAG)
 
-This service embeds the OpenFn Documentation to a vector database. It downloads, chunks, processes metadata, embeds and uploads the documentation to a vector database (Pinecone). 
+This service embeds the OpenFn Documentation to a vector database. It downloads,
+chunks, processes metadata, embeds and uploads the documentation to a vector
+database (Pinecone).
 
 ## Usage - Embedding OpenFn Documentation
 
-The vector database used here is Pinecone. To obtain the env variables follow these steps:
+The vector database used here is Pinecone. To obtain the env variables follow
+these steps:
 
 1. Create an account on [Pinecone] and set up a free cluster.
 2. Obtain the URL and token for the cluster and add them to the `.env` file.
@@ -15,26 +18,43 @@ The vector database used here is Pinecone. To obtain the env variables follow th
 ```bash
 openfn apollo embed_docsite tmp/payload.json
 ```
+
 To run directly from this repo (note that the server must be started):
 
 ```bash
 bun py embed_docsite tmp/payload.json -O
 ```
 
 ## Implementation
-The service uses the DocsiteProcessor to download the documentation and chunk it into smaller parts. The DocsiteIndexer formats metadata, creates a new collection, embeds the chunked texts (OpenAI) and uploads them into the vector database (Pinecone).
+
+The service uses the DocsiteProcessor to download the documentation and chunk it
+into smaller parts. The DocsiteIndexer formats metadata, creates a new
+collection, embeds the chunked texts (OpenAI) and uploads them into the vector
+database (Pinecone).
 
 The chunked texts can be viewed in `tmp/split_sections`.
 
 ## Payload Reference
+
+The write target is independent of the read backend (`DOCSITE_SEARCH_BACKEND`),
+so a Postgres batch can be built while Pinecone still serves search traffic.
+
 The input payload is a JSON object. All parameters are optional:
 
 ```js
 {
+    "target": "pinecone",             // 'pinecone' | 'postgres'. Defaults to pinecone. Chooses the write destination.
     "docs_to_upload": ["adaptor_docs", "general_docs", "adaptor_functions"], // Select from 3 types of documentation to upload
-    "collection_name": "docsite-20250225", // Name of the collection in the vector database (defaults to the current date)
-    "index_name": "docsite", // Name of the index in the vector database (an index contains collections; defaults to docsite)
     "docs_to_ignore": ["job-examples.md", "release-notes.md"], // Titles of documents that should not be indexed
-    "max_total_collections" : 3 // The max number of collections to keep in the vector database. This will delete older collections by date.
+    "chunk_target_length": 1000,      // Target chunk size in characters
+    "chunk_min_length": 700,          // Minimum chunk size before merging with the next split
+
+    // Pinecone target only:
+    "collection_name": "docsite-20250225",  // Namespace (defaults to the current timestamp)
+    "index_name": "docsite", // Name of the index in the vector database (an index contains collections; defaults to docsite)
+    "max_total_collections": 3, // The max number of collections to keep in the vector database. This will delete older collections by date.
+
+    // Postgres target only:
+    "keep_batches": 2                 // Number of recent complete batches to retain when pruning
 }
 ```