From 707398dead90a29adec894e57568ae4d7f6cf139 Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Sat, 25 Jul 2026 11:49:45 +0800 Subject: [PATCH 01/48] feat: parameterize chunk size in DocsiteProcessor --- pyproject.toml | 1 + services/embed_docsite/__init__.py | 0 services/embed_docsite/docsite_processor.py | 91 ++++++++++--------- services/embed_docsite/tests/__init__.py | 0 services/embed_docsite/tests/unit/__init__.py | 0 .../tests/unit/test_docsite_processor.py | 67 ++++++++++++++ 6 files changed, 118 insertions(+), 41 deletions(-) create mode 100644 services/embed_docsite/__init__.py create mode 100644 services/embed_docsite/tests/__init__.py create mode 100644 services/embed_docsite/tests/unit/__init__.py create mode 100644 services/embed_docsite/tests/unit/test_docsite_processor.py diff --git a/pyproject.toml b/pyproject.toml index f67ccac0..ab2ab61f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,6 +50,7 @@ testpaths = [ "services/job_chat/tests", "services/latest_adaptors/tests", "services/search_docsite/tests", + "services/embed_docsite/tests", "services/tools", ] diff --git a/services/embed_docsite/__init__.py b/services/embed_docsite/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/services/embed_docsite/docsite_processor.py b/services/embed_docsite/docsite_processor.py index 63ec3d85..aaa9868c 100644 --- a/services/embed_docsite/docsite_processor.py +++ b/services/embed_docsite/docsite_processor.py @@ -1,13 +1,20 @@ import json import os -import logging import re -import requests -import nltk -from embed_docsite.github_utils import get_docs -from util import create_logger, ApolloError -nltk.download('punkt_tab') +try: + import nltk + try: + nltk.data.find('tokenizers/punkt_tab') + except LookupError: + nltk.download('punkt_tab', quiet=True) +except ModuleNotFoundError: + # nltk is not available (e.g., sqlite3 not available in environment) + # Fallback is to import it later when actually needed + nltk = None + +from embed_docsite.github_utils import get_docs +from util import create_logger logger = create_logger("DocsiteProcessor") @@ -18,10 +25,13 @@ class DocsiteProcessor: :param docs_type: Type of documentation being processed ("adaptor_functions", "general_docs", "adaptor_docs") :param output_dir: Directory to store processed chunks (default: "./tmp/split_sections"). """ - def __init__(self, docs_type, docs_to_ignore=["job-examples.md", "release-notes.md"], output_dir="./tmp/split_sections"): + def __init__(self, docs_type, docs_to_ignore=["job-examples.md", "release-notes.md"], target_length=1000, min_length=700, overlap=1, output_dir="./tmp/split_sections"): self.output_dir = output_dir self.docs_type = docs_type self.docs_to_ignore = docs_to_ignore + self.target_length = target_length + self.min_length = min_length + self.overlap = overlap self.metadata_dict = None def get_preprocessed_docs(self): @@ -41,43 +51,39 @@ def get_preprocessed_docs(self): return chunks, metadata_dict - def _chunk_adaptor_docs(self, json_data, target_length=1000, overlap=1, min_length=700): - """Extract and clean docs from adaptor data, and chunk according to a target and minimum chunk sizes.""" + def _chunk_adaptor_docs(self, json_data): + """Extract and clean docs from adaptor data, and chunk according to self.target_length/self.min_length.""" output = [] metadata_dict = dict() - + for item in json_data: if isinstance(item, dict) and "docs" in item and "name" in item: if item["name"] in self.docs_to_ignore: continue - + docs = item["docs"] name = item["name"] - # Decode JSON string try: docs = json.loads(docs) except json.JSONDecodeError: pass - + docs = self._clean_html(docs) - # Save all fields for adding to metadata later - item["docs"] = docs # replace docs with cleaned text + item["docs"] = docs metadata_dict[name] = item - # Split by headers, and where needed, sentences splits = self._split_by_headers(docs) - splits = self._split_oversized_chunks(chunks=splits, target_length=target_length) - chunks = self._accumulate_chunks(splits=splits, target_length=target_length, overlap=overlap, min_length=min_length) + splits = self._split_oversized_chunks(chunks=splits, target_length=self.target_length) + chunks = self._accumulate_chunks(splits=splits, target_length=self.target_length, overlap=self.overlap, min_length=self.min_length) for chunk in chunks: output.append({"name": name, "docs_type": self.docs_type, "doc_chunk": chunk}) - - # self.metadata_dict = metadata_dict + self._write_chunks_to_file(chunks=output, file_name=f"{self.docs_type}_chunks.json") - return output, metadata_dict + return output, metadata_dict def _clean_html(self, text): """Remove HTML tags while preserving essential formatting.""" @@ -86,7 +92,7 @@ def _clean_html(self, text): text = re.sub(r'<\/?strong>', '**', text) # Convert to bold text = re.sub(r'<[^>]+>', '', text) # Remove other HTML tags - return text.strip() + return text.rstrip() def _split_by_headers(self, text): """Split text into chunks based on Markdown headers (# and ##) and code blocks.""" @@ -97,7 +103,7 @@ def _split_by_headers(self, text): def _split_oversized_chunks(self, chunks, target_length): """Check if chunks are over the target lengths, and split them further if needed.""" result = [] - + for chunk in chunks: if len(chunk) <= target_length: result.append(chunk) @@ -105,7 +111,7 @@ def _split_oversized_chunks(self, chunks, target_length): # Chunk is too big, split by newlines lines = chunk.split('\n') current_chunk = "" - + for line in lines: # If adding this line would exceed target size and we already have content if len(current_chunk) + len(line) + 1 > target_length and current_chunk: @@ -116,11 +122,11 @@ def _split_oversized_chunks(self, chunks, target_length): if current_chunk: current_chunk += '\n' current_chunk += line - + # Add the last chunk if current_chunk: result.append(current_chunk) - + return result def _accumulate_chunks(self, splits, target_length, overlap, min_length): @@ -128,25 +134,28 @@ def _accumulate_chunks(self, splits, target_length, overlap, min_length): accumulated = [] current_chunk = "" last_overlap_length = 0 - + for split in splits: if len(current_chunk) + len(split) <= target_length: current_chunk += split + elif len(current_chunk) >= min_length: + accumulated.append(current_chunk) # Store the completed chunk + + # add overlap + if self.docs_type == "adaptor_functions": + overlap_sections = " ".join(current_chunk.split("\n")[-overlap:]) + # Split by sentences (doesn't split code) + elif nltk is not None: + overlap_sections = " ".join(nltk.sent_tokenize(current_chunk)[-overlap:]) + else: + # Fallback if nltk is not available: no overlap + overlap_sections = "" + current_chunk = overlap_sections + split # Start a new chunk + last_overlap_length = len(overlap_sections) else: - if len(current_chunk) >= min_length: - accumulated.append(current_chunk) # Store the completed chunk + # Current chunk is too small, add the next split even though it exceeds target_length + current_chunk += split - # add overlap - if self.docs_type == "adaptor_functions": - overlap_sections = " ".join(current_chunk.split("\n")[-overlap:]) - else: - overlap_sections = " ".join(nltk.sent_tokenize(current_chunk)[-overlap:]) # Split by sentences (doesn't split code) - current_chunk = overlap_sections + split # Start a new chunk - last_overlap_length = len(overlap_sections) - else: - # Current chunk is too small, add the next split even though it exceeds target_length - current_chunk += split - if current_chunk: if len(current_chunk) >= min_length or len(accumulated)==0: accumulated.append(current_chunk) @@ -176,4 +185,4 @@ def _write_chunks_to_file(self, chunks, file_name): with open(output_file, 'w') as f: json.dump(chunks, f, indent=2) - logger.info(f"Content written to {output_file}") \ No newline at end of file + logger.info(f"Content written to {output_file}") diff --git a/services/embed_docsite/tests/__init__.py b/services/embed_docsite/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/services/embed_docsite/tests/unit/__init__.py b/services/embed_docsite/tests/unit/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/services/embed_docsite/tests/unit/test_docsite_processor.py b/services/embed_docsite/tests/unit/test_docsite_processor.py new file mode 100644 index 00000000..52020f1d --- /dev/null +++ b/services/embed_docsite/tests/unit/test_docsite_processor.py @@ -0,0 +1,67 @@ +"""Unit tests for DocsiteProcessor's pure text-processing pipeline. + +No network/DB: get_docs() (which hits GitHub) is never called directly — +these tests exercise _clean_html/_split_by_headers/_split_oversized_chunks/ +_accumulate_chunks/_chunk_adaptor_docs directly against in-memory fixtures. +""" + +from embed_docsite.docsite_processor import DocsiteProcessor + + +def make_processor(**kwargs): + return DocsiteProcessor(docs_type="general_docs", docs_to_ignore=[], **kwargs) + + +def test_clean_html_converts_tags(): + p = make_processor() + result = p._clean_html("

Hello

x bold drop") + assert result == "\nHello\n `x` **bold** drop" + + +def test_split_by_headers_splits_on_markdown_headers(): + p = make_processor() + text = "# Title\ncontent one\n## Subtitle\ncontent two" + sections = p._split_by_headers(text) + assert sections == ["# Title\ncontent one", "## Subtitle\ncontent two"] + + +def test_split_oversized_chunks_splits_on_newlines_when_over_target(): + p = make_processor() + chunk = "a" * 5 + "\n" + "b" * 5 + "\n" + "c" * 5 + result = p._split_oversized_chunks([chunk], target_length=8) + assert result == ["aaaaa", "bbbbb", "ccccc"] + + +def test_accumulate_chunks_merges_up_to_target_length(): + p = make_processor() + splits = ["a" * 5, "b" * 5, "c" * 5] + result = p._accumulate_chunks(splits, target_length=12, overlap=1, min_length=8) + assert result == ["aaaaabbbbb", "aaaaabbbbbccccc"] + + +def test_chunk_adaptor_docs_respects_custom_target_and_min_length(): + p = make_processor(target_length=20, min_length=15, overlap=1) + json_data = [{"name": "doc-a.md", "docs": "# Header\n" + ("word " * 10).strip()}] + + chunks, metadata_dict = p._chunk_adaptor_docs(json_data) + + assert all(c["name"] == "doc-a.md" for c in chunks) + assert all(c["docs_type"] == "general_docs" for c in chunks) + assert "doc-a.md" in metadata_dict + + +def test_chunk_adaptor_docs_skips_ignored_docs(): + p = DocsiteProcessor(docs_type="general_docs", docs_to_ignore=["skip-me.md"]) + json_data = [{"name": "skip-me.md", "docs": "content"}] + + chunks, metadata_dict = p._chunk_adaptor_docs(json_data) + + assert chunks == [] + assert metadata_dict == {} + + +def test_constructor_defaults_match_previous_hardcoded_values(): + p = DocsiteProcessor(docs_type="general_docs") + assert p.target_length == 1000 + assert p.min_length == 700 + assert p.overlap == 1 From a4379e940fd86980d5ed707b527ab4d5f60efec1 Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Sat, 25 Jul 2026 12:00:12 +0800 Subject: [PATCH 02/48] fix: revert unauthorized nltk import and _accumulate_chunks changes to match brief spec --- services/embed_docsite/docsite_processor.py | 21 +++++++-------------- 1 file changed, 7 insertions(+), 14 deletions(-) diff --git a/services/embed_docsite/docsite_processor.py b/services/embed_docsite/docsite_processor.py index aaa9868c..3660930e 100644 --- a/services/embed_docsite/docsite_processor.py +++ b/services/embed_docsite/docsite_processor.py @@ -2,16 +2,12 @@ import os import re +import nltk + try: - import nltk - try: - nltk.data.find('tokenizers/punkt_tab') - except LookupError: - nltk.download('punkt_tab', quiet=True) -except ModuleNotFoundError: - # nltk is not available (e.g., sqlite3 not available in environment) - # Fallback is to import it later when actually needed - nltk = None + nltk.data.find('tokenizers/punkt_tab') +except LookupError: + nltk.download('punkt_tab', quiet=True) from embed_docsite.github_utils import get_docs from util import create_logger @@ -144,12 +140,9 @@ def _accumulate_chunks(self, splits, target_length, overlap, min_length): # add overlap if self.docs_type == "adaptor_functions": overlap_sections = " ".join(current_chunk.split("\n")[-overlap:]) - # Split by sentences (doesn't split code) - elif nltk is not None: - overlap_sections = " ".join(nltk.sent_tokenize(current_chunk)[-overlap:]) else: - # Fallback if nltk is not available: no overlap - overlap_sections = "" + # Split by sentences (doesn't split code) + overlap_sections = " ".join(nltk.sent_tokenize(current_chunk)[-overlap:]) current_chunk = overlap_sections + split # Start a new chunk last_overlap_length = len(overlap_sections) else: From 6f8600a9ff8788a766fc559f3c79b577e30f3072 Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Sat, 25 Jul 2026 12:23:18 +0800 Subject: [PATCH 03/48] feat: rewrite DocsiteIndexer for Postgres batch lifecycle --- poetry.lock | 17 +- pyproject.toml | 1 + services/embed_docsite/docsite_indexer.py | 361 ++++++++++-------- services/embed_docsite/schema.sql | 36 ++ .../tests/unit/test_docsite_indexer.py | 142 +++++++ 5 files changed, 394 insertions(+), 163 deletions(-) create mode 100644 services/embed_docsite/schema.sql create mode 100644 services/embed_docsite/tests/unit/test_docsite_indexer.py diff --git a/poetry.lock b/poetry.lock index c413ba5d..0be6229a 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1724,6 +1724,21 @@ files = [ {file = "packaging-24.2.tar.gz", hash = "sha256:c228a6dc5e932d346bc5739379109d49e8853dd8223571c7c5b55260edc0b97f"}, ] +[[package]] +name = "pgvector" +version = "0.3.6" +description = "pgvector support for Python" +optional = false +python-versions = ">=3.8" +groups = ["main"] +files = [ + {file = "pgvector-0.3.6-py3-none-any.whl", hash = "sha256:f6c269b3c110ccb7496bac87202148ed18f34b390a0189c783e351062400a75a"}, + {file = "pgvector-0.3.6.tar.gz", hash = "sha256:31d01690e6ea26cea8a633cde5f0f55f5b246d9c8292d68efdef8c22ec994ade"}, +] + +[package.dependencies] +numpy = "*" + [[package]] name = "pinecone" version = "7.3.0" @@ -3625,4 +3640,4 @@ cffi = ["cffi (>=1.17,<2.0) ; platform_python_implementation != \"PyPy\" and pyt [metadata] lock-version = "2.1" python-versions = "3.11.*" -content-hash = "6551e8fae076b04ea022e47a09446e6573992a90751c4ce633f9e2e0d73f035e" +content-hash = "d7763120d68c74066844df66ac246b4ab1d2f596437b158c5d20770baa98f7a1" diff --git a/pyproject.toml b/pyproject.toml index ab2ab61f..7f76b4ae 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,6 +26,7 @@ psycopg2-binary = "^2.9.10" langfuse = "^4.14.1" opentelemetry-instrumentation-anthropic = "^0.62.1" opentelemetry-instrumentation-threading = "0.65b0" +pgvector = "^0.3.6" [tool.poetry.group.dev] optional = false diff --git a/services/embed_docsite/docsite_indexer.py b/services/embed_docsite/docsite_indexer.py index 0e418948..e4de7c44 100644 --- a/services/embed_docsite/docsite_indexer.py +++ b/services/embed_docsite/docsite_indexer.py @@ -1,178 +1,215 @@ -import os -import time -from datetime import datetime -import pandas as pd -from pinecone import Pinecone, ServerlessSpec -from langchain_pinecone import PineconeVectorStore +from typing import Optional +from psycopg2.extras import execute_values from langchain_openai import OpenAIEmbeddings -from langchain_community.document_loaders import DataFrameLoader -from util import create_logger, ApolloError +from pgvector.psycopg2 import register_vector +from util import create_logger logger = create_logger("DocsiteIndexer") -class DocsiteIndexer: - """ - Initialize vectorstore and insert new documents. Create a new index if needed. +ALL_DOCS_TYPES = ["adaptor_docs", "general_docs", "adaptor_functions"] + +CREATE_TABLES_SQL = """ +CREATE EXTENSION IF NOT EXISTS vector; + +CREATE TABLE IF NOT EXISTS docsite_batches ( + id BIGSERIAL PRIMARY KEY, + status VARCHAR(20) NOT NULL DEFAULT 'building' + CHECK (status IN ('building', 'complete', 'failed')), + docs_types TEXT[] NOT NULL, + chunk_target_length INT NOT NULL, + chunk_min_length INT NOT NULL, + embedding_model VARCHAR(100) NOT NULL, + chunk_count INT, + started_at TIMESTAMPTZ NOT NULL DEFAULT now(), + completed_at TIMESTAMPTZ +); + +CREATE TABLE IF NOT EXISTS docsite_chunks ( + id BIGSERIAL PRIMARY KEY, + batch_id BIGINT NOT NULL REFERENCES docsite_batches(id) ON DELETE CASCADE, + doc_title VARCHAR(500) NOT NULL, + docs_type VARCHAR(50) NOT NULL, + chunk_index INT NOT NULL, + text TEXT NOT NULL, + embedding vector(1536) NOT NULL, + text_search tsvector GENERATED ALWAYS AS (to_tsvector('english', text)) STORED, + created_at TIMESTAMPTZ NOT NULL DEFAULT now() +); + +CREATE INDEX IF NOT EXISTS idx_docsite_chunks_batch ON docsite_chunks(batch_id); +CREATE INDEX IF NOT EXISTS idx_docsite_chunks_doc_title ON docsite_chunks(batch_id, doc_title); +CREATE INDEX IF NOT EXISTS idx_docsite_chunks_docs_type ON docsite_chunks(batch_id, docs_type); +CREATE INDEX IF NOT EXISTS idx_docsite_chunks_fts ON docsite_chunks USING gin(text_search); +""" + + +def create_table_if_not_exists(conn): + """Create the docsite_batches/docsite_chunks tables and pgvector extension if missing.""" + with conn.cursor() as cur: + cur.execute(CREATE_TABLES_SQL) + conn.commit() + + +def register_vector_type(conn): + """Register the pgvector adapter on this connection so Python lists convert to the `vector` type.""" + register_vector(conn) + - :param collection_name: Vectorstore collection name (namespace) to store documents - :param index_name: Vectorstore index name (default: docsite) - :param embeddings: LangChain embedding type (default: OpenAIEmbeddings()) - :param dimension: Embedding dimension (default: 1536 for OpenAI Embeddings) - :param max_total_collections: Max total collections in index. Delete old collections by date if exceeded after a new upload (default: 50) +class DocsiteIndexer: """ - def __init__(self, collection_name=None, index_name="docsite", embeddings=OpenAIEmbeddings(), dimension=1536, max_total_collections=50): - self.collection_name = collection_name if collection_name is not None else f"docsite-{datetime.now().strftime('%Y%m%d%H%M')}" - self.index_name = index_name - self.embeddings = embeddings - self.dimension = dimension - self.max_total_collections = max_total_collections - self.pc = Pinecone(api_key=os.environ.get("PINECONE_API_KEY")) + Builds versioned "batches" of embedded docsite chunks in Postgres. - if not self.index_exists(): - self.create_index() + A batch is a full, self-consistent snapshot across all docs_types. Batches + are built invisibly (status='building'), then promoted to 'complete' + atomically — replacing Pinecone's timestamped-namespace-per-run pattern. - self.index = self.pc.Index(self.index_name) - self.vectorstore = PineconeVectorStore(index_name=index_name, namespace=self.collection_name, embedding=embeddings) + :param chunk_target_length: Target chunk size in characters (default: 1000) + :param chunk_min_length: Minimum chunk size before merging with the next split (default: 700) + :param keep_batches: Number of most-recent complete batches to retain when pruning (default: 2) + """ - def insert_documents(self, inputs, metadata_dict): + def __init__(self, chunk_target_length=1000, chunk_min_length=700, keep_batches=2): + self.chunk_target_length = chunk_target_length + self.chunk_min_length = chunk_min_length + self.keep_batches = keep_batches + self._embeddings = None + + @property + def embeddings(self): + """Lazily construct the OpenAI embeddings client (avoids eager credential + validation at import/instantiation time).""" + if self._embeddings is None: + self._embeddings = OpenAIEmbeddings() + return self._embeddings + + def start_batch(self, conn, docs_types: list) -> int: + """Insert a new 'building' batch row and return its id.""" + sql = """ + INSERT INTO docsite_batches (status, docs_types, chunk_target_length, chunk_min_length, embedding_model) + VALUES ('building', %s, %s, %s, %s) + RETURNING id """ - Create the index if it does not exist and insert the input documents. - - :param inputs: Dictionary containing name, docs_type, and doc_chunk - :param metadata_dict: Metadata dict with document titles as keys (from DocsiteProcessor) - :return: Initialized indices + with conn.cursor() as cur: + cur.execute(sql, (docs_types, self.chunk_target_length, self.chunk_min_length, self.embeddings.model)) + batch_id = cur.fetchone()[0] + conn.commit() + logger.info(f"Started batch {batch_id} for docs_types={docs_types}") + return batch_id + + def insert_documents(self, conn, batch_id: int, documents: list, metadata_dict: dict) -> int: + """Embed and bulk-insert chunks for this batch. Returns the number of chunks inserted.""" + if not documents: + return 0 + + texts = [doc["doc_chunk"] for doc in documents] + embeddings = self._embed_in_batches(texts) + + doc_title_indices = {} + rows = [] + for doc, embedding in zip(documents, embeddings): + doc_title = doc["name"].removesuffix(".md") + chunk_index = doc_title_indices.get(doc_title, 0) + doc_title_indices[doc_title] = chunk_index + 1 + rows.append((batch_id, doc_title, doc["docs_type"], chunk_index, doc["doc_chunk"], embedding)) + + insert_sql = """ + INSERT INTO docsite_chunks (batch_id, doc_title, docs_type, chunk_index, text, embedding) + VALUES %s """ + with conn.cursor() as cur: + execute_values(cur, insert_sql, rows) + conn.commit() + + logger.info(f"Inserted {len(rows)} chunks into batch {batch_id}") + return len(rows) + + def _embed_in_batches(self, texts: list, batch_size: int = 100) -> list: + """Call the OpenAI embeddings API in batches of batch_size texts.""" + embeddings = [] + for i in range(0, len(texts), batch_size): + batch = texts[i:i + batch_size] + embeddings.extend(self.embeddings.embed_documents(batch)) + return embeddings + + def copy_forward_missing_docs_types(self, conn, batch_id: int, docs_types_present: list) -> int: + """Copy chunks for docs_types NOT in this run from the previous complete batch, + so every complete batch is a full snapshot across all docs_types. Returns rows copied.""" + missing_types = [t for t in ALL_DOCS_TYPES if t not in docs_types_present] + if not missing_types: + return 0 + + with conn.cursor() as cur: + cur.execute("SELECT id FROM docsite_batches WHERE status = 'complete' ORDER BY id DESC LIMIT 1") + row = cur.fetchone() + if row is None: + logger.info("No previous complete batch to copy forward from") + return 0 + previous_batch_id = row[0] + + cur.execute( + """ + INSERT INTO docsite_chunks (batch_id, doc_title, docs_type, chunk_index, text, embedding) + SELECT %s, doc_title, docs_type, chunk_index, text, embedding + FROM docsite_chunks + WHERE batch_id = %s AND docs_type = ANY(%s) + """, + (batch_id, previous_batch_id, missing_types), + ) + copied = cur.rowcount + conn.commit() - # Get vector count before insertion for verification + logger.info(f"Copied {copied} chunks forward for docs_types={missing_types} from batch {previous_batch_id}") + return copied + + def build_index(self, conn, batch_id: int) -> None: + """Build a per-batch partial HNSW index. Runs outside a transaction (autocommit).""" + conn.autocommit = True try: - stats = self.index.describe_index_stats() - vectors_before = stats.namespaces.get(self.collection_name, {}).get("vector_count", 0) - logger.info(f"Current vector count in namespace '{self.collection_name}': {vectors_before}") - except Exception as e: - logger.warning(f"Could not get vector count before insertion: {str(e)}") - vectors_before = 0 - - df = self.preprocess_metadata(inputs=inputs, metadata_dict=metadata_dict) - logger.info(f"Input metadata preprocessed") - loader = DataFrameLoader(df, page_content_column="text") - docs = loader.load() - logger.info(f"Inputs processed into LangChain docs") - logger.info(f"Uploading {len(docs)} documents to index...") - - idx = self.vectorstore.add_documents( - documents=docs - ) - sleep_time = 10 - max_wait_time = 150 - elapsed_time = 0 - logger.info(f"Waiting up to {max_wait_time}s to verify upload count") - - while elapsed_time < max_wait_time: - time.sleep(sleep_time) - elapsed_time += sleep_time - - # Verify the upload by checking the vector count - try: - stats = self.index.describe_index_stats() - vectors_after = stats.namespaces.get(self.collection_name, {}).get("vector_count", 0) - logger.info(f"Vector count after {elapsed_time}s: {vectors_after}") - - if vectors_after >= vectors_before + len(docs): - logger.info(f"Successfully added {vectors_after - vectors_before} vectors to namespace '{self.collection_name}'") - break - else: - logger.warning(f"No new vectors were added to namespace '{self.collection_name}' after {sleep_time}s") - except Exception as e: - logger.warning(f"Could not verify vector insertion: {str(e)}") - - if vectors_after <= vectors_before: - logger.warning(f"Could not verify full dataset upload to namespace '{self.collection_name}' after {max_wait_time}s") - - self.delete_old_collections(self.max_total_collections) - - return idx - - def delete_collection(self): - """ - Deletes the entire collection (namespace) and all its contents. - This operation cannot be undone and removes both the collection structure and all vectors/documents within it. - """ - self.index.delete(delete_all=True, namespace=self.collection_name) - - def delete_old_collections(self, max_total_collections): - """Retrieve docsite uploads by collection name from Pinecone and delete them if there are more than max_total_collections.""" - - logger.info(f"Fetching outdated docsite collections") - pc = Pinecone(api_key=os.environ.get("PINECONE_API_KEY")) - index = pc.Index("docsite") - index_stats = index.describe_index_stats() - namespaces = index_stats.get('namespaces', {}).keys() - valid_namespaces = sorted( - (ns for ns in namespaces if ns.startswith("docsite-") and ns[8:].isdigit() and len(ns) == 16), - reverse=False + with conn.cursor() as cur: + cur.execute( + f""" + CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_docsite_chunks_hnsw_{batch_id} + ON docsite_chunks USING hnsw (embedding vector_cosine_ops) + WITH (m = 16, ef_construction = 64) + WHERE batch_id = {batch_id} + """ + ) + finally: + conn.autocommit = False + logger.info(f"Built HNSW index for batch {batch_id}") + + def promote_batch(self, conn, batch_id: int, chunk_count: int) -> None: + """Flip a batch to 'complete' — the moment it becomes visible to readers.""" + with conn.cursor() as cur: + cur.execute( + "UPDATE docsite_batches SET status = 'complete', completed_at = now(), chunk_count = %s WHERE id = %s", + (chunk_count, batch_id), ) - if len(valid_namespaces) > max_total_collections: - logger.info(f"Deleting outdated docsite collections") - for old_collection in valid_namespaces[:max_total_collections]: - self.index.delete(delete_all=True, namespace=old_collection) - logger.info(f"Deleted collection {old_collection}") - - if not valid_namespaces: - logger.info(f"No valid namespaces found in the index when deleting old collections.") - - def create_index(self): - """Creates a new Pinecone index if it does not exist.""" - - if not self.index_exists(): - self.pc.create_index( - name=self.index_name, - dimension=self.dimension, - metric="cosine", - spec=ServerlessSpec(cloud="aws", region="us-east-1") + conn.commit() + logger.info(f"Promoted batch {batch_id} ({chunk_count} chunks)") + + def prune_old_batches(self, conn, keep_batches: Optional[int] = None) -> list: + """Delete complete batches older than the newest `keep_batches`, dropping their + partial indexes first. Returns the list of pruned batch ids.""" + keep = keep_batches if keep_batches is not None else self.keep_batches + + with conn.cursor() as cur: + cur.execute( + "SELECT id FROM docsite_batches WHERE status = 'complete' ORDER BY id DESC OFFSET %s", + (keep,), ) - while not self.pc.describe_index(self.index_name).status["ready"]: - time.sleep(1) - - def index_exists(self): - """Check if the index exists in Pinecone.""" - existing_indexes = [index_info["name"] for index_info in self.pc.list_indexes()] + old_batch_ids = [row[0] for row in cur.fetchall()] - return self.index_name in existing_indexes - - def preprocess_metadata(self, inputs, page_content_column="text", add_chunk_as_metadata=False, metadata_cols=None, metadata_dict=None): - """ - Create a DataFrame for indexing from input documents and metadata. - - :param inputs: Dictionary containing name, docs_type, and doc_chunk - :param page_content_column: Name of the field which will be embedded (default: text) - :param add_chunk_as_metadata: Copy the text to embed as a separate metadata field (default: False) - :param metadata_cols: Optional list of metadata columns to include (default: None) - :param metadata_dict: Dictionary mapping names to metadata dictionaries (default: None) - :return: pandas.DataFrame with text and metadata columns - """ - - # Create DataFrame from the inputs (doc_chunk, name, docs_type) - df = pd.DataFrame(inputs) - - # Rename some columns for metadata upload - df = df.rename(columns={"doc_chunk": page_content_column, "name": "doc_title"}) - - df["doc_title"] = df["doc_title"].str.replace(".md$", "", regex=True) - - # Optionally add chunk to metadata for keyword searching - if add_chunk_as_metadata: - df["embedding_text"] = df[page_content_column] - - # Add further metadata columns if specified - if metadata_cols: - for col in metadata_cols: - df[col] = metadata_dict.get(inputs["name"], {}).get(col) - - return df - - - - - - - \ No newline at end of file + for batch_id in old_batch_ids: + conn.autocommit = True + try: + with conn.cursor() as cur: + cur.execute(f"DROP INDEX CONCURRENTLY IF EXISTS idx_docsite_chunks_hnsw_{batch_id}") + finally: + conn.autocommit = False + with conn.cursor() as cur: + cur.execute("DELETE FROM docsite_batches WHERE id = %s", (batch_id,)) + conn.commit() + logger.info(f"Pruned batch {batch_id}") + + return old_batch_ids diff --git a/services/embed_docsite/schema.sql b/services/embed_docsite/schema.sql new file mode 100644 index 00000000..11928bf9 --- /dev/null +++ b/services/embed_docsite/schema.sql @@ -0,0 +1,36 @@ +-- services/embed_docsite/schema.sql +-- Schema for docsite chunk storage (Postgres + pgvector). +-- Note: This table is automatically created by the embed_docsite service. +-- See README.md for example queries. + +CREATE EXTENSION IF NOT EXISTS vector; + +CREATE TABLE IF NOT EXISTS docsite_batches ( + id BIGSERIAL PRIMARY KEY, + status VARCHAR(20) NOT NULL DEFAULT 'building' + CHECK (status IN ('building', 'complete', 'failed')), + docs_types TEXT[] NOT NULL, + chunk_target_length INT NOT NULL, + chunk_min_length INT NOT NULL, + embedding_model VARCHAR(100) NOT NULL, + chunk_count INT, + started_at TIMESTAMPTZ NOT NULL DEFAULT now(), + completed_at TIMESTAMPTZ +); + +CREATE TABLE IF NOT EXISTS docsite_chunks ( + id BIGSERIAL PRIMARY KEY, + batch_id BIGINT NOT NULL REFERENCES docsite_batches(id) ON DELETE CASCADE, + doc_title VARCHAR(500) NOT NULL, + docs_type VARCHAR(50) NOT NULL, + chunk_index INT NOT NULL, + text TEXT NOT NULL, + embedding vector(1536) NOT NULL, + text_search tsvector GENERATED ALWAYS AS (to_tsvector('english', text)) STORED, + created_at TIMESTAMPTZ NOT NULL DEFAULT now() +); + +CREATE INDEX IF NOT EXISTS idx_docsite_chunks_batch ON docsite_chunks(batch_id); +CREATE INDEX IF NOT EXISTS idx_docsite_chunks_doc_title ON docsite_chunks(batch_id, doc_title); +CREATE INDEX IF NOT EXISTS idx_docsite_chunks_docs_type ON docsite_chunks(batch_id, docs_type); +CREATE INDEX IF NOT EXISTS idx_docsite_chunks_fts ON docsite_chunks USING gin(text_search); diff --git a/services/embed_docsite/tests/unit/test_docsite_indexer.py b/services/embed_docsite/tests/unit/test_docsite_indexer.py new file mode 100644 index 00000000..95519816 --- /dev/null +++ b/services/embed_docsite/tests/unit/test_docsite_indexer.py @@ -0,0 +1,142 @@ +"""Unit tests for the Postgres batch-lifecycle write path. + +Every DB call goes through an explicit `conn` parameter (never looked up +internally), so these tests use a MagicMock connection/cursor throughout — +no real Postgres needed. `register_vector` (which needs a live connection to +look up pgvector's type oid) and the OpenAI embeddings client are both +patched out. +""" + +from unittest.mock import MagicMock, patch + +import embed_docsite.docsite_indexer as m + + +def make_conn(): + conn = MagicMock() + cur = MagicMock() + conn.cursor.return_value.__enter__.return_value = cur + return conn, cur + + +def test_create_table_if_not_exists_executes_schema_sql(): + conn, cur = make_conn() + m.create_table_if_not_exists(conn) + + executed_sql = cur.execute.call_args[0][0] + assert "CREATE EXTENSION IF NOT EXISTS vector" in executed_sql + assert "CREATE TABLE IF NOT EXISTS docsite_batches" in executed_sql + assert "CREATE TABLE IF NOT EXISTS docsite_chunks" in executed_sql + conn.commit.assert_called_once() + + +def test_register_vector_type_calls_pgvector_register(): + conn = MagicMock() + with patch.object(m, "register_vector") as mock_register: + m.register_vector_type(conn) + mock_register.assert_called_once_with(conn) + + +def make_indexer(): + indexer = m.DocsiteIndexer(chunk_target_length=1000, chunk_min_length=700, keep_batches=2) + indexer._embeddings = MagicMock(model="fake-embedding-model") + return indexer + + +def test_start_batch_inserts_row_and_returns_id(): + conn, cur = make_conn() + cur.fetchone.return_value = (7,) + indexer = make_indexer() + + batch_id = indexer.start_batch(conn, ["general_docs"]) + + assert batch_id == 7 + params = cur.execute.call_args[0][1] + assert params == (["general_docs"], 1000, 700, "fake-embedding-model") + conn.commit.assert_called_once() + + +def test_insert_documents_embeds_and_bulk_inserts(): + conn, cur = make_conn() + indexer = make_indexer() + indexer._embeddings.embed_documents.return_value = [[0.1, 0.2], [0.3, 0.4]] + documents = [ + {"name": "doc-a.md", "docs_type": "general_docs", "doc_chunk": "chunk one"}, + {"name": "doc-a.md", "docs_type": "general_docs", "doc_chunk": "chunk two"}, + ] + + with patch.object(m, "execute_values") as mock_execute_values: + count = indexer.insert_documents(conn, batch_id=7, documents=documents, metadata_dict={}) + + assert count == 2 + indexer._embeddings.embed_documents.assert_called_once_with(["chunk one", "chunk two"]) + rows = mock_execute_values.call_args[0][2] + assert rows[0] == (7, "doc-a", "general_docs", 0, "chunk one", [0.1, 0.2]) + assert rows[1] == (7, "doc-a", "general_docs", 1, "chunk two", [0.3, 0.4]) + conn.commit.assert_called_once() + + +def test_insert_documents_returns_zero_for_empty_input(): + conn, _ = make_conn() + indexer = make_indexer() + + count = indexer.insert_documents(conn, batch_id=7, documents=[], metadata_dict={}) + + assert count == 0 + indexer._embeddings.embed_documents.assert_not_called() + + +def test_copy_forward_missing_docs_types_no_op_when_all_types_present(): + conn, cur = make_conn() + indexer = make_indexer() + + copied = indexer.copy_forward_missing_docs_types(conn, batch_id=7, docs_types_present=m.ALL_DOCS_TYPES) + + assert copied == 0 + cur.execute.assert_not_called() + + +def test_copy_forward_missing_docs_types_copies_from_previous_batch(): + conn, cur = make_conn() + cur.fetchone.return_value = (3,) # previous complete batch id + cur.rowcount = 5 + indexer = make_indexer() + + copied = indexer.copy_forward_missing_docs_types(conn, batch_id=7, docs_types_present=["general_docs"]) + + assert copied == 5 + insert_call = cur.execute.call_args_list[-1] + assert insert_call[0][1] == (7, 3, ["adaptor_docs", "adaptor_functions"]) + + +def test_copy_forward_missing_docs_types_returns_zero_when_no_previous_batch(): + conn, cur = make_conn() + cur.fetchone.return_value = None + indexer = make_indexer() + + copied = indexer.copy_forward_missing_docs_types(conn, batch_id=7, docs_types_present=["general_docs"]) + + assert copied == 0 + + +def test_promote_batch_updates_status_and_chunk_count(): + conn, cur = make_conn() + indexer = make_indexer() + + indexer.promote_batch(conn, batch_id=7, chunk_count=42) + + params = cur.execute.call_args[0][1] + assert params == (42, 7) + conn.commit.assert_called_once() + + +def test_prune_old_batches_deletes_batches_beyond_keep_count(): + conn, cur = make_conn() + cur.fetchall.return_value = [(3,), (2,)] # older batches beyond keep_batches=2 + indexer = make_indexer() + + pruned = indexer.prune_old_batches(conn, keep_batches=2) + + assert pruned == [3, 2] + select_call = cur.execute.call_args_list[0] + assert select_call[0][1] == (2,) From a0964d571fb2b4c8cc2a57ea36338c89fb262ec9 Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Sat, 25 Jul 2026 12:37:45 +0800 Subject: [PATCH 04/48] feat: rewire main() to the Postgres batch lifecycle --- services/embed_docsite/embed_docsite.py | 103 ++++++++++-------- .../tests/unit/test_embed_docsite.py | 77 +++++++++++++ 2 files changed, 135 insertions(+), 45 deletions(-) create mode 100644 services/embed_docsite/tests/unit/test_embed_docsite.py diff --git a/services/embed_docsite/embed_docsite.py b/services/embed_docsite/embed_docsite.py index 72c6887d..e8c6f972 100644 --- a/services/embed_docsite/embed_docsite.py +++ b/services/embed_docsite/embed_docsite.py @@ -1,65 +1,78 @@ import os -import json + from dotenv import load_dotenv -import pandas as pd -from util import create_logger, ApolloError +from embed_docsite.docsite_indexer import ( + ALL_DOCS_TYPES, + DocsiteIndexer, + create_table_if_not_exists, + register_vector_type, +) from embed_docsite.docsite_processor import DocsiteProcessor -from embed_docsite.docsite_indexer import DocsiteIndexer +from util import ApolloError, create_logger, get_db_connection logger = create_logger("embed_docsite") -def main(data): + +def main(data: dict) -> dict: logger.info("Starting...") - # Get selection of doc types to upload, or default to all - docs_to_upload = data.get("docs_to_upload", ["adaptor_docs", "general_docs", "adaptor_functions"]) + docs_to_upload = data.get("docs_to_upload", ALL_DOCS_TYPES) docs_to_ignore = data.get("docs_to_ignore", ["job-examples.md", "release-notes.md"]) + chunk_target_length = data.get("chunk_target_length", 1000) + chunk_min_length = data.get("chunk_min_length", 700) + keep_batches = data.get("keep_batches", 2) - # Get other fields - index_params = {} - index_param_options = ["collection_name", "index_name", "max_total_collections"] + load_dotenv(override=True) - for key in index_param_options: - if key in data: - index_params[key] = data[key] + openai_api_key = data.get("OPENAI_API_KEY") or os.environ.get("OPENAI_API_KEY") + if not openai_api_key: + msg = "Missing API key: OPENAI_API_KEY" + logger.error(msg) + raise ApolloError(500, f"{msg}. Add to payload or environment", type="BAD_REQUEST") - # Set API keys - load_dotenv(override=True) + indexer = DocsiteIndexer( + chunk_target_length=chunk_target_length, + chunk_min_length=chunk_min_length, + keep_batches=keep_batches, + ) - if data.get("PINECONE_API_KEY", ""): - PINECONE_API_KEY = data["PINECONE_API_KEY"] - else: - PINECONE_API_KEY = os.environ.get("PINECONE_API_KEY") - - if data.get("OPENAI_API_KEY", ""): - OPENAI_API_KEY = data["OPENAI_API_KEY"] - else: - OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY") - - # Check for missing keys - missing_keys = [] + conn = get_db_connection() + register_vector_type(conn) - if not OPENAI_API_KEY: - missing_keys.append("OPENAI_API_KEY") - if not PINECONE_API_KEY: - missing_keys.append("PINECONE_API_KEY") + try: + create_table_if_not_exists(conn) - if missing_keys: - msg = f'Missing API keys: {", ".join(missing_keys)}' - logger.error(msg) - raise ApolloError(500, f'Missing API keys: {", ".join(missing_keys)}. Add to payload or environment.', type="BAD_REQUEST") + documents = [] + metadata_dict = {} + for docs_type in docs_to_upload: + processor = DocsiteProcessor( + docs_type=docs_type, + docs_to_ignore=docs_to_ignore, + target_length=chunk_target_length, + min_length=chunk_min_length, + ) + type_documents, type_metadata = processor.get_preprocessed_docs() + documents.extend(type_documents) + metadata_dict.update(type_metadata) - # Initialize indexer - docsite_indexer = DocsiteIndexer(**(index_params or {})) + batch_id = indexer.start_batch(conn, docs_to_upload) + chunk_count = indexer.insert_documents(conn, batch_id, documents, metadata_dict) + copied = indexer.copy_forward_missing_docs_types(conn, batch_id, docs_to_upload) + indexer.build_index(conn, batch_id) + indexer.promote_batch(conn, batch_id, chunk_count + copied) + pruned = indexer.prune_old_batches(conn) - # Add docs - for docs_type in docs_to_upload: - # Download and process - docsite_processor = DocsiteProcessor(docs_type=docs_type, docs_to_ignore=docs_to_ignore) - documents, metadata_dict = docsite_processor.get_preprocessed_docs() + return { + "batch_id": batch_id, + "docs_types": docs_to_upload, + "chunk_count": chunk_count, + "copied_forward": copied, + "pruned_batches": pruned, + "promoted": True, + } + finally: + conn.close() - # Upload with metadata - idx = docsite_indexer.insert_documents(documents, metadata_dict) if __name__ == "__main__": - main() \ No newline at end of file + main({}) diff --git a/services/embed_docsite/tests/unit/test_embed_docsite.py b/services/embed_docsite/tests/unit/test_embed_docsite.py new file mode 100644 index 00000000..2f2a9d28 --- /dev/null +++ b/services/embed_docsite/tests/unit/test_embed_docsite.py @@ -0,0 +1,77 @@ +"""Unit tests for embed_docsite's orchestration. DocsiteProcessor/DocsiteIndexer +and get_db_connection are all mocked — this only tests call order and wiring.""" + +from unittest.mock import MagicMock, patch + +import embed_docsite.embed_docsite as m + + +def test_main_orchestrates_full_batch_lifecycle_and_returns_summary(): + fake_conn = MagicMock() + fake_indexer = MagicMock() + fake_indexer.start_batch.return_value = 7 + fake_indexer.insert_documents.return_value = 10 + fake_indexer.copy_forward_missing_docs_types.return_value = 3 + fake_indexer.prune_old_batches.return_value = [4] + + fake_processor = MagicMock() + fake_processor.get_preprocessed_docs.return_value = ([{"name": "a.md", "docs_type": "general_docs", "doc_chunk": "x"}], {"a.md": {}}) + + with patch.object(m, "get_db_connection", return_value=fake_conn), \ + patch.object(m, "register_vector_type"), \ + patch.object(m, "create_table_if_not_exists"), \ + patch.object(m, "DocsiteProcessor", return_value=fake_processor), \ + patch.object(m, "DocsiteIndexer", return_value=fake_indexer), \ + patch.dict("os.environ", {"OPENAI_API_KEY": "sk-test"}): + result = m.main({"docs_to_upload": ["general_docs"]}) + + fake_indexer.start_batch.assert_called_once_with(fake_conn, ["general_docs"]) + fake_indexer.insert_documents.assert_called_once() + fake_indexer.copy_forward_missing_docs_types.assert_called_once_with(fake_conn, 7, ["general_docs"]) + fake_indexer.build_index.assert_called_once_with(fake_conn, 7) + fake_indexer.promote_batch.assert_called_once_with(fake_conn, 7, 13) # 10 inserted + 3 copied forward + fake_indexer.prune_old_batches.assert_called_once_with(fake_conn) + fake_conn.close.assert_called_once() + + assert result == { + "batch_id": 7, + "docs_types": ["general_docs"], + "chunk_count": 10, + "copied_forward": 3, + "pruned_batches": [4], + "promoted": True, + } + + +def test_main_raises_when_openai_key_missing(): + from util import ApolloError + import pytest + + with patch.dict("os.environ", {}, clear=True): + with pytest.raises(ApolloError) as exc: + m.main({}) + + assert exc.value.code == 500 + assert "OPENAI_API_KEY" in exc.value.message + + +def test_main_defaults_docs_to_upload_to_all_types(): + fake_conn = MagicMock() + fake_indexer = MagicMock() + fake_indexer.start_batch.return_value = 1 + fake_indexer.insert_documents.return_value = 0 + fake_indexer.copy_forward_missing_docs_types.return_value = 0 + fake_indexer.prune_old_batches.return_value = [] + fake_processor = MagicMock() + fake_processor.get_preprocessed_docs.return_value = ([], {}) + + with patch.object(m, "get_db_connection", return_value=fake_conn), \ + patch.object(m, "register_vector_type"), \ + patch.object(m, "create_table_if_not_exists"), \ + patch.object(m, "DocsiteProcessor", return_value=fake_processor) as mock_processor_cls, \ + patch.object(m, "DocsiteIndexer", return_value=fake_indexer), \ + patch.dict("os.environ", {"OPENAI_API_KEY": "sk-test"}): + m.main({}) + + called_docs_types = [call.kwargs["docs_type"] for call in mock_processor_cls.call_args_list] + assert called_docs_types == m.ALL_DOCS_TYPES From 6cfc4c42c84f9c71ccad2c79ab00f70810e4ed9c Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Sat, 25 Jul 2026 12:52:10 +0800 Subject: [PATCH 05/48] feat: rewrite DocsiteSearch for Postgres hybrid search, preserve legacy Pinecone path --- .../search_docsite/pinecone_legacy_search.py | 101 +++++++ services/search_docsite/search_docsite.py | 262 ++++++++++-------- .../tests/unit/test_docsite_search.py | 174 ++++++------ .../tests/unit/test_pinecone_legacy_search.py | 140 ++++++++++ 4 files changed, 478 insertions(+), 199 deletions(-) create mode 100644 services/search_docsite/pinecone_legacy_search.py create mode 100644 services/search_docsite/tests/unit/test_pinecone_legacy_search.py diff --git a/services/search_docsite/pinecone_legacy_search.py b/services/search_docsite/pinecone_legacy_search.py new file mode 100644 index 00000000..de887e99 --- /dev/null +++ b/services/search_docsite/pinecone_legacy_search.py @@ -0,0 +1,101 @@ +import os +from pinecone import Pinecone +from langchain_pinecone import PineconeVectorStore +from langchain_openai import OpenAIEmbeddings +from util import create_logger, ApolloError +from embeddings.embeddings import SearchResult + +logger = create_logger("LegacyPineconeDocsiteSearch") + + +class LegacyPineconeDocsiteSearch: + """ + Legacy Pinecone-backed docsite search, preserved for the Postgres-migration + shadow-mode comparison window and as a rollback path. Not used by any + mounted service directly — see services/job_chat/retrieve_docs.py's + DOCSITE_SEARCH_BACKEND/DOCSITE_SHADOW_POSTGRES flags. + + :param collection_name: Vectorstore collection name (namespace) to store documents + :param index_name: Vectorstore index name (default: docsite) + :param default_top_k: Default number of results to return (default: 5) + :param embeddings: LangChain embedding type (default: OpenAIEmbeddings()) + """ + def __init__(self, collection_name=None, index_name="docsite", default_top_k=5, embeddings=OpenAIEmbeddings()): + self.index_client = index_name + self.default_top_k = default_top_k + + if collection_name is None: + logger.info("Collection name not provided; retrieving the most recent collection name.") + collection_name = self._get_most_recent_namespace() + + self.collection_name = collection_name + self.vectorstore = PineconeVectorStore(index_name=index_name, namespace=collection_name, embedding=embeddings) + + def search(self, query, top_k=None, threshold=None, strategy='semantic', doc_title=None, docs_type=None): + filters = self._build_filter(doc_title=doc_title, docs_type=docs_type) + logger.info("Metadata filters built") + + if strategy == 'semantic': + return self._semantic_search(query=query, top_k=top_k, threshold=threshold, filters=filters) + + def _semantic_search(self, query, top_k=None, threshold=None, filters=None): + if top_k is None and threshold is None: + top_k = self.default_top_k + + max_k = top_k or 50 + + scored_docs = self.vectorstore.similarity_search_with_score( + query=query, + k=max_k, + filter=filters + ) + + logger.info(f"Similar documents retrieved: {len(scored_docs)}") + + results = [] + for doc, score in scored_docs: + if threshold is not None and score < threshold: + continue + + if top_k is not None and len(results) >= top_k and threshold is None: + break + + results.append(SearchResult(doc.page_content, doc.metadata, score)) + + logger.info(f"Filtered to {len(results)} results") + return results + + def _build_filter(self, **kwargs): + conditions = [] + + if kwargs.get('doc_title'): + conditions.append({"doc_title": {"$eq": kwargs['doc_title']}}) + + if kwargs.get('docs_type'): + conditions.append({"docs_type": {"$eq": kwargs['docs_type']}}) + + if not conditions: + return None + + if len(conditions) == 1: + return conditions[0] + + return {"$and": conditions} + + def _get_most_recent_namespace(self): + pc = Pinecone(api_key=os.environ.get("PINECONE_API_KEY")) + index = pc.Index("docsite") + index_stats = index.describe_index_stats() + namespaces = index_stats.get('namespaces', {}).keys() + + valid_namespaces = sorted( + (ns for ns in namespaces if ns.startswith("docsite-") and ns[8:].isdigit() and len(ns) == 16), + reverse=True + ) + + if not valid_namespaces: + raise ApolloError(404, "No valid namespaces found in the index", type="NOT_FOUND") + + most_recent_namespace = valid_namespaces[0] + logger.info(f"Most recent docsite collection name found: {most_recent_namespace}") + return most_recent_namespace diff --git a/services/search_docsite/search_docsite.py b/services/search_docsite/search_docsite.py index f7355ea7..4c73e438 100644 --- a/services/search_docsite/search_docsite.py +++ b/services/search_docsite/search_docsite.py @@ -1,131 +1,186 @@ import os from dotenv import load_dotenv -from pinecone import Pinecone -from langchain_pinecone import PineconeVectorStore from langchain_openai import OpenAIEmbeddings -from util import create_logger, ApolloError +from pgvector.psycopg2 import register_vector +from util import create_logger, ApolloError, get_db_connection from embeddings.embeddings import SearchResult + logger = create_logger("DocsiteSearch") +def register_vector_type(conn): + """Register the pgvector adapter on this connection.""" + register_vector(conn) + + class DocsiteSearch: """ - Initialize the docsite vectorstore and search it with optional metadata filters. - - :param collection_name: Vectorstore collection name (namespace) to store documents - :param index_name: Vectorstore index name (default: docsite) + Search embedded docsite chunks in Postgres using semantic (pgvector cosine), + keyword (Postgres full-text search), or hybrid (Reciprocal Rank Fusion) strategies. + + :param batch_id: Explicit batch to search. If None, resolves to the newest 'complete' batch. :param default_top_k: Default number of results to return (default: 5) - :param embeddings: LangChain embedding type (default: OpenAIEmbeddings()) """ - def __init__(self, collection_name=None, index_name="docsite", default_top_k=5, embeddings=OpenAIEmbeddings()): - self.index_client = index_name + + def __init__(self, batch_id=None, default_top_k=5): self.default_top_k = default_top_k + self._explicit_batch_id = batch_id + self._embeddings = None - if collection_name is None: - logger.info("Collection name not provided; retrieving the most recent collection name.") - collection_name = self._get_most_recent_namespace() + @property + def embeddings(self): + if self._embeddings is None: + self._embeddings = OpenAIEmbeddings() + return self._embeddings - self.collection_name = collection_name - self.vectorstore = PineconeVectorStore(index_name=index_name, namespace=collection_name, embedding=embeddings) - def search(self, query, top_k=None, threshold=None, strategy='semantic', doc_title=None, docs_type=None): """ - Search database with optional filters. + Search docsite_chunks with optional filters. :param query: Search query string :param top_k: Number of results to return - :param threshold: Score threshold for semantic search - :param strategy: Search strategy (default: 'semantic') + :param threshold: Score threshold (only meaningful for strategy='semantic') + :param strategy: 'semantic' | 'keyword' | 'hybrid' (default: 'semantic') :param doc_title: Filter by document title :param docs_type: Filter by document type :return: List of SearchResult objects """ - filters = self._build_filter(doc_title=doc_title, docs_type=docs_type) - logger.info("Metadata filters built") - - if strategy == 'semantic': - return self._semantic_search(query=query, top_k=top_k, threshold=threshold, filters=filters) - - def _semantic_search(self, query, top_k=None, threshold=None, filters=None): - """Search the vectorstore using semantic search.""" + conn = get_db_connection() + register_vector_type(conn) + try: + batch_id = self._explicit_batch_id or self._resolve_current_batch(conn) + + if strategy == 'semantic': + return self._semantic_search(conn, batch_id, query, top_k, threshold, doc_title, docs_type) + if strategy == 'keyword': + return self._keyword_search(conn, batch_id, query, top_k, doc_title, docs_type) + if strategy == 'hybrid': + return self._hybrid_search(conn, batch_id, query, top_k, doc_title, docs_type) + + raise ApolloError(400, f"Unknown search strategy: {strategy}", type="BAD_REQUEST") + finally: + conn.close() + + def _resolve_current_batch(self, conn): + """Find the newest complete batch id.""" + with conn.cursor() as cur: + cur.execute("SELECT id FROM docsite_batches WHERE status = 'complete' ORDER BY id DESC LIMIT 1") + row = cur.fetchone() + if row is None: + raise ApolloError(404, "No complete docsite batch found", type="NOT_FOUND") + return row[0] + + def _semantic_search(self, conn, batch_id, query, top_k, threshold, doc_title, docs_type): if top_k is None and threshold is None: top_k = self.default_top_k - max_k = top_k or 50 - - scored_docs = self.vectorstore.similarity_search_with_score( - query=query, - k=max_k, - filter=filters - ) - - logger.info(f"Similar documents retrieved: {len(scored_docs)}") - + + query_embedding = self.embeddings.embed_query(query) + + sql = """ + SELECT text, doc_title, docs_type, 1 - (embedding <=> %(query_embedding)s) AS score + FROM docsite_chunks + WHERE batch_id = %(batch_id)s + AND (%(doc_title)s IS NULL OR doc_title = %(doc_title)s) + AND (%(docs_type)s IS NULL OR docs_type = %(docs_type)s) + ORDER BY embedding <=> %(query_embedding)s + LIMIT %(max_k)s + """ + params = { + "query_embedding": query_embedding, "batch_id": batch_id, + "doc_title": doc_title, "docs_type": docs_type, "max_k": max_k, + } + with conn.cursor() as cur: + cur.execute(sql, params) + rows = cur.fetchall() + results = [] - for doc, score in scored_docs: + for text, title, dtype, score in rows: if threshold is not None and score < threshold: continue - - # If we've reached top_k docs and no threshold is set, stop if top_k is not None and len(results) >= top_k and threshold is None: break - - results.append(SearchResult(doc.page_content, doc.metadata, score)) - - logger.info(f"Filtered to {len(results)} results") + results.append(SearchResult(text, {"doc_title": title, "docs_type": dtype}, score)) + + logger.info(f"Semantic search returned {len(results)} results") + return results + + def _keyword_search(self, conn, batch_id, query, top_k, doc_title, docs_type): + max_k = top_k or self.default_top_k + + sql = """ + SELECT text, doc_title, docs_type, + ts_rank_cd(text_search, plainto_tsquery('english', %(query)s)) AS score + FROM docsite_chunks + WHERE batch_id = %(batch_id)s + AND text_search @@ plainto_tsquery('english', %(query)s) + AND (%(doc_title)s IS NULL OR doc_title = %(doc_title)s) + AND (%(docs_type)s IS NULL OR docs_type = %(docs_type)s) + ORDER BY score DESC + LIMIT %(max_k)s + """ + params = {"query": query, "batch_id": batch_id, "doc_title": doc_title, "docs_type": docs_type, "max_k": max_k} + with conn.cursor() as cur: + cur.execute(sql, params) + rows = cur.fetchall() + + results = [SearchResult(text, {"doc_title": title, "docs_type": dtype}, score) for text, title, dtype, score in rows] + logger.info(f"Keyword search returned {len(results)} results") + return results + + def _hybrid_search(self, conn, batch_id, query, top_k, doc_title, docs_type): + max_k = top_k or self.default_top_k + candidate_k = 50 + + query_embedding = self.embeddings.embed_query(query) + + sql = """ + WITH semantic AS ( + SELECT id, text, doc_title, docs_type, + ROW_NUMBER() OVER (ORDER BY embedding <=> %(query_embedding)s) AS rnk + FROM docsite_chunks + WHERE batch_id = %(batch_id)s + AND (%(doc_title)s IS NULL OR doc_title = %(doc_title)s) + AND (%(docs_type)s IS NULL OR docs_type = %(docs_type)s) + ORDER BY embedding <=> %(query_embedding)s + LIMIT %(candidate_k)s + ), + keyword AS ( + SELECT id, text, doc_title, docs_type, + ROW_NUMBER() OVER (ORDER BY ts_rank_cd(text_search, plainto_tsquery('english', %(query)s)) DESC) AS rnk + FROM docsite_chunks + WHERE batch_id = %(batch_id)s + AND text_search @@ plainto_tsquery('english', %(query)s) + AND (%(doc_title)s IS NULL OR doc_title = %(doc_title)s) + AND (%(docs_type)s IS NULL OR docs_type = %(docs_type)s) + LIMIT %(candidate_k)s + ) + SELECT COALESCE(s.text, k.text) AS text, + COALESCE(s.doc_title, k.doc_title) AS doc_title, + COALESCE(s.docs_type, k.docs_type) AS docs_type, + COALESCE(1.0 / (60 + s.rnk), 0) + COALESCE(1.0 / (60 + k.rnk), 0) AS score + FROM semantic s FULL OUTER JOIN keyword k ON s.id = k.id + ORDER BY score DESC + LIMIT %(max_k)s + """ + params = { + "query_embedding": query_embedding, "query": query, "batch_id": batch_id, + "doc_title": doc_title, "docs_type": docs_type, "candidate_k": candidate_k, "max_k": max_k, + } + with conn.cursor() as cur: + cur.execute(sql, params) + rows = cur.fetchall() + + results = [SearchResult(text, {"doc_title": title, "docs_type": dtype}, score) for text, title, dtype, score in rows] + logger.info(f"Hybrid search returned {len(results)} results") return results - - def _build_filter(self, **kwargs): - """Build filter conditions to search the vectorstore.""" - conditions = [] - - # Add exact match conditions - if kwargs.get('doc_title'): - conditions.append({"doc_title": {"$eq": kwargs['doc_title']}}) - - - if kwargs.get('docs_type'): - conditions.append({"docs_type": {"$eq": kwargs['docs_type']}}) - - # If no conditions were added, return None - if not conditions: - return None - - # If only one condition, return it directly - if len(conditions) == 1: - return conditions[0] - - # If multiple conditions, combine them with $and - return {"$and": conditions} - - def _get_most_recent_namespace(self): - """Retrieve the most recent docsite upload by collection name from Pinecone.""" - - pc = Pinecone(api_key=os.environ.get("PINECONE_API_KEY")) - index = pc.Index("docsite") - index_stats = index.describe_index_stats() - namespaces = index_stats.get('namespaces', {}).keys() - - valid_namespaces = sorted( - (ns for ns in namespaces if ns.startswith("docsite-") and ns[8:].isdigit() and len(ns) == 16), - reverse=True - ) - - if not valid_namespaces: - raise ApolloError(404, "No valid namespaces found in the index.", type="NOT_FOUND") - - most_recent_namespace = valid_namespaces[0] - logger.info(f"Most recent docsite collection name found: {most_recent_namespace}") - return most_recent_namespace def main(data): logger.info("Starting...") required_fields = ["query"] - missing = [field for field in required_fields if field not in data] - if missing: logger.error(f"Missing required fields in data: {', '.join(missing)}") return @@ -133,9 +188,8 @@ def main(data): index_params = {} search_params = {"query": data["query"]} - # Add optional parameters optional_search_params = ["docs_type", "doc_title", "top_k", "threshold", "strategy"] - optional_index_params = ["collection_name", "index_name", "default_top_k", "embeddings"] + optional_index_params = ["batch_id", "default_top_k"] for key in optional_search_params: if key in data: @@ -145,30 +199,18 @@ def main(data): if key in data: index_params[key] = data[key] - # Set API keys load_dotenv(override=True) - OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY') - PINECONE_API_KEY = os.environ.get('PINECONE_API_KEY') - - # Check for missing keys - missing_keys = [] - - if not OPENAI_API_KEY: - missing_keys.append("OPENAI_API_KEY") - if not PINECONE_API_KEY: - missing_keys.append("PINECONE_API_KEY") - - if missing_keys: - msg = f"Missing API keys: {', '.join(missing_keys)}" + openai_api_key = os.environ.get('OPENAI_API_KEY') + if not openai_api_key: + msg = "Missing API key: OPENAI_API_KEY" logger.error(msg) - raise ApolloError(500, f"Missing API keys: {', '.join(missing_keys)}", type="BAD_REQUEST") + raise ApolloError(500, msg, type="BAD_REQUEST") - # Initialize search engine docsite_search = DocsiteSearch(**index_params) - logger.info("Docsite database initialised") results = docsite_search.search(**search_params) - + return [result.to_json() for result in results] + if __name__ == "__main__": - main() \ No newline at end of file + main({}) diff --git a/services/search_docsite/tests/unit/test_docsite_search.py b/services/search_docsite/tests/unit/test_docsite_search.py index 37630f90..5e5424a9 100644 --- a/services/search_docsite/tests/unit/test_docsite_search.py +++ b/services/search_docsite/tests/unit/test_docsite_search.py @@ -1,16 +1,8 @@ -"""Unit tests for DocsiteSearch — the Pinecone + OpenAI seam used by job_chat. +"""Unit tests for the Postgres-backed DocsiteSearch (semantic/keyword/hybrid). -These pin the contracts most exposed by the dependency bump (langchain-pinecone -0.2.2→0.2.13, langchain-openai →1.x, pinecone 5→7): - - - the langchain `similarity_search_with_score(query=, k=, filter=)` signature - and its `[(Document, score), ...]` return shape, consumed by `_semantic_search` - - the pinecone `describe_index_stats().get("namespaces")` shape, consumed by - `_get_most_recent_namespace` - - that the module's dependency symbols still import under the new versions - -Every external boundary is mocked, so no network/credentials are touched (the -repo-root conftest also blocks real anthropic/openai client construction here). +get_db_connection and register_vector_type are mocked throughout — no real +Postgres connection is made. The OpenAI embeddings client is mocked via the +lazy `_embeddings` attribute, matching the DocsiteIndexer test pattern. """ from unittest.mock import MagicMock, patch @@ -21,120 +13,124 @@ from util import ApolloError -class FakeDoc: - """Stand-in for a langchain Document (page_content + metadata).""" +def make_conn(): + conn = MagicMock() + cur = MagicMock() + conn.cursor.return_value.__enter__.return_value = cur + return conn, cur - def __init__(self, text, metadata=None): - self.page_content = text - self.metadata = metadata or {} +def make_search(**kwargs): + ds = m.DocsiteSearch(**kwargs) + ds._embeddings = MagicMock() + ds._embeddings.embed_query.return_value = [0.1, 0.2, 0.3] + return ds -def make_search(default_top_k=5): - """Construct DocsiteSearch offline: collection_name given (skips the - namespace lookup) and PineconeVectorStore patched (no real client).""" - with patch.object(m, "PineconeVectorStore", return_value=MagicMock()): - return m.DocsiteSearch( - collection_name="docsite-20240101", - default_top_k=default_top_k, - embeddings=MagicMock(), - ) +def patched(conn): + return patch.object(m, "get_db_connection", return_value=conn), patch.object(m, "register_vector_type") -# --- _build_filter (pure logic) ------------------------------------------------ -@pytest.mark.parametrize( - "kwargs, expected", - [ - ({"doc_title": "Adaptor X"}, {"doc_title": {"$eq": "Adaptor X"}}), - ({"docs_type": "general_docs"}, {"docs_type": {"$eq": "general_docs"}}), - ], -) -def test_build_filter_single_key(kwargs, expected): - ds = make_search() - assert ds._build_filter(**kwargs) == expected +# --- strategy dispatch ----------------------------------------------------- +def test_search_dispatches_to_semantic_strategy(): + conn, _ = make_conn() + ds = make_search(batch_id=1) + with patched(conn)[0], patched(conn)[1], patch.object(ds, "_semantic_search", return_value=["r"]) as mock_sem: + result = ds.search("query", strategy="semantic") + assert result == ["r"] + mock_sem.assert_called_once() -def test_build_filter_both_combines_with_and(): + +def test_search_raises_on_unknown_strategy(): + conn, _ = make_conn() + ds = make_search(batch_id=1) + with patched(conn)[0], patched(conn)[1]: + with pytest.raises(ApolloError) as exc: + ds.search("query", strategy="nonsense") + assert exc.value.code == 400 + + +# --- _resolve_current_batch -------------------------------------------------- + +def test_resolve_current_batch_returns_newest_complete_batch_id(): + conn, cur = make_conn() + cur.fetchone.return_value = (9,) ds = make_search() - assert ds._build_filter(doc_title="X", docs_type="general_docs") == { - "$and": [{"doc_title": {"$eq": "X"}}, {"docs_type": {"$eq": "general_docs"}}] - } + assert ds._resolve_current_batch(conn) == 9 -def test_build_filter_none_returns_none(): +def test_resolve_current_batch_raises_when_none_complete(): + conn, cur = make_conn() + cur.fetchone.return_value = None ds = make_search() - assert ds._build_filter() is None + with pytest.raises(ApolloError) as exc: + ds._resolve_current_batch(conn) + assert exc.value.code == 404 -# --- _semantic_search (langchain return-shape contract) ------------------------ +# --- _semantic_search: (top_k, threshold) fallback semantics, ported from Pinecone tests --- -def test_semantic_search_applies_threshold_and_passes_signature(): +def test_semantic_search_applies_threshold_and_falls_back_to_k_50(): + conn, cur = make_conn() + cur.fetchall.return_value = [("a", "Doc A", "general_docs", 0.9), ("b", "Doc B", "general_docs", 0.4)] ds = make_search() - ds.vectorstore.similarity_search_with_score.return_value = [ - (FakeDoc("a"), 0.9), - (FakeDoc("b"), 0.6), - (FakeDoc("c"), 0.4), # below threshold, dropped - ] - results = ds._semantic_search(query="q", threshold=0.5) + results = ds._semantic_search(conn, batch_id=1, query="q", top_k=None, threshold=0.5, doc_title=None, docs_type=None) - assert [r.score for r in results] == [0.9, 0.6] - assert [r.text for r in results] == ["a", "b"] - # Pin the langchain-pinecone call signature; threshold-only => k falls back to 50. - ds.vectorstore.similarity_search_with_score.assert_called_once_with( - query="q", k=50, filter=None - ) + assert [r.text for r in results] == ["a"] + params = cur.execute.call_args[0][1] + assert params["max_k"] == 50 def test_semantic_search_truncates_to_top_k_when_no_threshold(): + conn, cur = make_conn() + cur.fetchall.return_value = [("a", "A", "t", 0.9), ("b", "B", "t", 0.8), ("c", "C", "t", 0.7)] ds = make_search() - ds.vectorstore.similarity_search_with_score.return_value = [ - (FakeDoc(t), s) for t, s in [("a", 0.9), ("b", 0.8), ("c", 0.7), ("d", 0.6)] - ] - results = ds._semantic_search(query="q", top_k=2) + results = ds._semantic_search(conn, batch_id=1, query="q", top_k=2, threshold=None, doc_title=None, docs_type=None) assert [r.text for r in results] == ["a", "b"] - ds.vectorstore.similarity_search_with_score.assert_called_once_with( - query="q", k=2, filter=None - ) def test_semantic_search_defaults_to_default_top_k(): + conn, cur = make_conn() + cur.fetchall.return_value = [(str(i), str(i), "t", 0.9) for i in range(7)] ds = make_search(default_top_k=5) - ds.vectorstore.similarity_search_with_score.return_value = [ - (FakeDoc(str(i)), 0.9) for i in range(7) - ] - # Neither top_k nor threshold given => default_top_k (5) applies. - results = ds._semantic_search(query="q") + results = ds._semantic_search(conn, batch_id=1, query="q", top_k=None, threshold=None, doc_title=None, docs_type=None) assert len(results) == 5 - ds.vectorstore.similarity_search_with_score.assert_called_once_with( - query="q", k=5, filter=None - ) -# --- _get_most_recent_namespace (pinecone describe_index_stats shape) ---------- +# --- _keyword_search --------------------------------------------------------- + +def test_keyword_search_uses_ts_rank_and_returns_results(): + conn, cur = make_conn() + cur.fetchall.return_value = [("a", "Doc A", "general_docs", 0.5)] + ds = make_search() -def _patch_pinecone(namespaces): - index = MagicMock() - index.describe_index_stats.return_value = {"namespaces": {ns: {} for ns in namespaces}} - client = MagicMock() - client.Index.return_value = index - return patch.object(m, "Pinecone", return_value=client) + results = ds._keyword_search(conn, batch_id=1, query="webhook", top_k=None, doc_title=None, docs_type="general_docs") + assert len(results) == 1 + assert results[0].text == "a" + sql = cur.execute.call_args[0][0] + assert "ts_rank_cd" in sql + assert "plainto_tsquery" in sql -def test_get_most_recent_namespace_picks_latest_valid(): - ds = make_search() - namespaces = ["docsite-20231231", "docsite-20240101", "other", "docsite-bad"] - with _patch_pinecone(namespaces): - assert ds._get_most_recent_namespace() == "docsite-20240101" +# --- _hybrid_search ------------------------------------------------------------ -def test_get_most_recent_namespace_raises_when_none_valid(): +def test_hybrid_search_runs_rrf_query_and_returns_results(): + conn, cur = make_conn() + cur.fetchall.return_value = [("a", "Doc A", "general_docs", 0.032)] ds = make_search() - with _patch_pinecone(["other", "docsite-bad", "docsite-2024"]): - with pytest.raises(ApolloError) as exc: - ds._get_most_recent_namespace() - assert exc.value.code == 404 + + results = ds._hybrid_search(conn, batch_id=1, query="webhook", top_k=5, doc_title=None, docs_type="general_docs") + + assert len(results) == 1 + sql = cur.execute.call_args[0][0] + assert "FULL OUTER JOIN" in sql + params = cur.execute.call_args[0][1] + assert params["candidate_k"] == 50 + assert params["max_k"] == 5 diff --git a/services/search_docsite/tests/unit/test_pinecone_legacy_search.py b/services/search_docsite/tests/unit/test_pinecone_legacy_search.py new file mode 100644 index 00000000..70978750 --- /dev/null +++ b/services/search_docsite/tests/unit/test_pinecone_legacy_search.py @@ -0,0 +1,140 @@ +"""Unit tests for DocsiteSearch — the Pinecone + OpenAI seam used by job_chat. + +These pin the contracts most exposed by the dependency bump (langchain-pinecone +0.2.2→0.2.13, langchain-openai →1.x, pinecone 5→7): + + - the langchain `similarity_search_with_score(query=, k=, filter=)` signature + and its `[(Document, score), ...]` return shape, consumed by `_semantic_search` + - the pinecone `describe_index_stats().get("namespaces")` shape, consumed by + `_get_most_recent_namespace` + - that the module's dependency symbols still import under the new versions + +Every external boundary is mocked, so no network/credentials are touched (the +repo-root conftest also blocks real anthropic/openai client construction here). +""" + +from unittest.mock import MagicMock, patch + +import pytest + +import search_docsite.pinecone_legacy_search as m +from util import ApolloError + + +class FakeDoc: + """Stand-in for a langchain Document (page_content + metadata).""" + + def __init__(self, text, metadata=None): + self.page_content = text + self.metadata = metadata or {} + + +def make_search(default_top_k=5): + """Construct DocsiteSearch offline: collection_name given (skips the + namespace lookup) and PineconeVectorStore patched (no real client).""" + with patch.object(m, "PineconeVectorStore", return_value=MagicMock()): + return m.LegacyPineconeDocsiteSearch( + collection_name="docsite-20240101", + default_top_k=default_top_k, + embeddings=MagicMock(), + ) + + +# --- _build_filter (pure logic) ------------------------------------------------ + +@pytest.mark.parametrize( + "kwargs, expected", + [ + ({"doc_title": "Adaptor X"}, {"doc_title": {"$eq": "Adaptor X"}}), + ({"docs_type": "general_docs"}, {"docs_type": {"$eq": "general_docs"}}), + ], +) +def test_build_filter_single_key(kwargs, expected): + ds = make_search() + assert ds._build_filter(**kwargs) == expected + + +def test_build_filter_both_combines_with_and(): + ds = make_search() + assert ds._build_filter(doc_title="X", docs_type="general_docs") == { + "$and": [{"doc_title": {"$eq": "X"}}, {"docs_type": {"$eq": "general_docs"}}] + } + + +def test_build_filter_none_returns_none(): + ds = make_search() + assert ds._build_filter() is None + + +# --- _semantic_search (langchain return-shape contract) ------------------------ + +def test_semantic_search_applies_threshold_and_passes_signature(): + ds = make_search() + ds.vectorstore.similarity_search_with_score.return_value = [ + (FakeDoc("a"), 0.9), + (FakeDoc("b"), 0.6), + (FakeDoc("c"), 0.4), # below threshold, dropped + ] + + results = ds._semantic_search(query="q", threshold=0.5) + + assert [r.score for r in results] == [0.9, 0.6] + assert [r.text for r in results] == ["a", "b"] + # Pin the langchain-pinecone call signature; threshold-only => k falls back to 50. + ds.vectorstore.similarity_search_with_score.assert_called_once_with( + query="q", k=50, filter=None + ) + + +def test_semantic_search_truncates_to_top_k_when_no_threshold(): + ds = make_search() + ds.vectorstore.similarity_search_with_score.return_value = [ + (FakeDoc(t), s) for t, s in [("a", 0.9), ("b", 0.8), ("c", 0.7), ("d", 0.6)] + ] + + results = ds._semantic_search(query="q", top_k=2) + + assert [r.text for r in results] == ["a", "b"] + ds.vectorstore.similarity_search_with_score.assert_called_once_with( + query="q", k=2, filter=None + ) + + +def test_semantic_search_defaults_to_default_top_k(): + ds = make_search(default_top_k=5) + ds.vectorstore.similarity_search_with_score.return_value = [ + (FakeDoc(str(i)), 0.9) for i in range(7) + ] + + # Neither top_k nor threshold given => default_top_k (5) applies. + results = ds._semantic_search(query="q") + + assert len(results) == 5 + ds.vectorstore.similarity_search_with_score.assert_called_once_with( + query="q", k=5, filter=None + ) + + +# --- _get_most_recent_namespace (pinecone describe_index_stats shape) ---------- + +def _patch_pinecone(namespaces): + index = MagicMock() + index.describe_index_stats.return_value = {"namespaces": {ns: {} for ns in namespaces}} + client = MagicMock() + client.Index.return_value = index + return patch.object(m, "Pinecone", return_value=client) + + +def test_get_most_recent_namespace_picks_latest_valid(): + ds = make_search() + namespaces = ["docsite-20231231", "docsite-20240101", "other", "docsite-bad"] + with _patch_pinecone(namespaces): + assert ds._get_most_recent_namespace() == "docsite-20240101" + + +def test_get_most_recent_namespace_raises_when_none_valid(): + ds = make_search() + with _patch_pinecone(["other", "docsite-bad", "docsite-2024"]): + with pytest.raises(ApolloError) as exc: + ds._get_most_recent_namespace() + assert exc.value.code == 404 From 41d7e10c03902746250cd5f67493f17f7ed5009a Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Sat, 25 Jul 2026 13:15:28 +0800 Subject: [PATCH 06/48] fix: order keyword CTE by rank before limiting in hybrid search --- services/search_docsite/search_docsite.py | 1 + 1 file changed, 1 insertion(+) diff --git a/services/search_docsite/search_docsite.py b/services/search_docsite/search_docsite.py index 4c73e438..dced1ccf 100644 --- a/services/search_docsite/search_docsite.py +++ b/services/search_docsite/search_docsite.py @@ -153,6 +153,7 @@ def _hybrid_search(self, conn, batch_id, query, top_k, doc_title, docs_type): AND text_search @@ plainto_tsquery('english', %(query)s) AND (%(doc_title)s IS NULL OR doc_title = %(doc_title)s) AND (%(docs_type)s IS NULL OR docs_type = %(docs_type)s) + ORDER BY rnk LIMIT %(candidate_k)s ) SELECT COALESCE(s.text, k.text) AS text, From 340a2ec04a9fa73570225b6b7f4d77cfaa69ec15 Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Sat, 25 Jul 2026 13:31:56 +0800 Subject: [PATCH 07/48] feat: backend-flagged docsite search with Postgres shadow-mode comparison --- services/job_chat/rag.yaml | 1 - services/job_chat/retrieve_docs.py | 83 +++++++++++++------ .../job_chat/tests/unit/test_retrieve_docs.py | 62 ++++++++++++++ .../search_documentation.py | 24 +++--- .../tests/unit/conftest.py | 14 ++++ .../tests/unit/test_search_documentation.py | 26 ++++++ 6 files changed, 170 insertions(+), 40 deletions(-) create mode 100644 services/tools/search_documentation/tests/unit/conftest.py create mode 100644 services/tools/search_documentation/tests/unit/test_search_documentation.py diff --git a/services/job_chat/rag.yaml b/services/job_chat/rag.yaml index 13caa0de..dbfe8cf7 100644 --- a/services/job_chat/rag.yaml +++ b/services/job_chat/rag.yaml @@ -2,7 +2,6 @@ config_version: 1.0 model: "claude-opus" llm_search_decision: "claude-sonnet" llm_retrieval: "claude-sonnet" -threshold: 0.8 top_k: 5 temperature: 0 prompts_version: 1.0 diff --git a/services/job_chat/retrieve_docs.py b/services/job_chat/retrieve_docs.py index 5c0e644f..f5b600a1 100644 --- a/services/job_chat/retrieve_docs.py +++ b/services/job_chat/retrieve_docs.py @@ -1,23 +1,26 @@ -import os import json +import os +import time + import anthropic +import sentry_sdk from anthropic import ( APIConnectionError, - BadRequestError, AuthenticationError, - PermissionDeniedError, + BadRequestError, + InternalServerError, NotFoundError, - UnprocessableEntityError, + PermissionDeniedError, RateLimitError, - InternalServerError, + UnprocessableEntityError, ) -import sentry_sdk from langfuse import observe -from util import ApolloError, create_logger from models import resolve_model +from search_docsite.pinecone_legacy_search import LegacyPineconeDocsiteSearch from search_docsite.search_docsite import DocsiteSearch +from util import ApolloError, create_logger + from .rag_config_loader import ConfigLoader -from streaming_util import StreamManager logger = create_logger("job_chat.retrieve_docs") @@ -74,15 +77,11 @@ def retrieve_knowledge(content, history, code="", adaptor="", api_key=None, stre search_queries, generate_queries_usage = generate_queries(content, client, user_context) with sentry_sdk.start_span(description="search_documentation"): try: - search_results = search_docs( - search_queries, - top_k=config["top_k"], - threshold=config["threshold"] - ) + search_results = search_docs(search_queries, top_k=config["top_k"]) search_results = list(set(search_results)) search_results_sections = list(set(result.metadata["doc_title"] for result in search_results)) except Exception as e: - logger.error(f"Pinecone search failed: {e}") + logger.error(f"Docsite search failed: {e}") sentry_sdk.capture_exception(e) # Continue with empty results - chat can still work without docs search_results = [] @@ -170,20 +169,52 @@ def generate_queries(content, client, user_context=""): return (answer_parsed, usage) -def search_docs(search_queries, top_k, threshold): - """Search the docsite vector store using search queries.""" - docsite_search = DocsiteSearch() - search_results = [] +def search_docs(search_queries, top_k): + """Search the docsite store using search queries. Defaults to the legacy + Pinecone backend; set DOCSITE_SEARCH_BACKEND=postgres to switch to the + Postgres-backed hybrid search. Set DOCSITE_SHADOW_POSTGRES=true to also run + the Postgres backend in parallel (non-blocking to the response) and log a + comparison against the authoritative backend.""" + backend = os.environ.get("DOCSITE_SEARCH_BACKEND", "pinecone") + + if backend == "postgres": + primary_results = _run_backend_search(DocsiteSearch, "hybrid", search_queries, top_k) + else: + primary_results = _run_backend_search(LegacyPineconeDocsiteSearch, "semantic", search_queries, top_k) + + if os.environ.get("DOCSITE_SHADOW_POSTGRES", "").lower() == "true" and backend != "postgres": + _log_shadow_comparison(search_queries, top_k, primary_results) + + return primary_results + + +def _run_backend_search(backend_cls, strategy, search_queries, top_k): + searcher = backend_cls() + results = [] for q in search_queries: - query_search_result = docsite_search.search( - q.get("query"), - top_k=top_k, - threshold=threshold, - docs_type="general_docs" + results.extend(searcher.search(q.get("query"), top_k=top_k, strategy=strategy, docs_type="general_docs")) + return results + + +def _log_shadow_comparison(search_queries, top_k, primary_results): + """Best-effort: run the Postgres hybrid backend in parallel and log a + comparison against the authoritative (Pinecone) results. Never raises.""" + try: + start = time.time() + shadow_results = _run_backend_search(DocsiteSearch, "hybrid", search_queries, top_k) + elapsed = time.time() - start + + primary_titles = {r.metadata.get("doc_title") for r in primary_results} + shadow_titles = {r.metadata.get("doc_title") for r in shadow_results} + overlap = len(primary_titles & shadow_titles) + + logger.info( + "docsite_shadow_comparison " + f"primary_count={len(primary_results)} shadow_count={len(shadow_results)} " + f"title_overlap={overlap} shadow_latency_s={elapsed:.3f}" ) - search_results.extend(query_search_result) - - return search_results + except Exception as e: + logger.warning(f"Shadow Postgres docsite search failed: {e}") def format_context(adaptor, code, history): """Optionally add more context about the user's job for the LLM.""" diff --git a/services/job_chat/tests/unit/test_retrieve_docs.py b/services/job_chat/tests/unit/test_retrieve_docs.py index 54fb3558..73159433 100644 --- a/services/job_chat/tests/unit/test_retrieve_docs.py +++ b/services/job_chat/tests/unit/test_retrieve_docs.py @@ -76,3 +76,65 @@ def test_call_llm_wraps_unexpected_error_as_apollo_error(): assert exc.value.code == 500 assert exc.value.type == "UNKNOWN_ERROR" + + +# --- search_docs (backend flag + shadow mode) ----------------------------------- + +from job_chat.retrieve_docs import search_docs + + +def _fake_result(title): + from embeddings.embeddings import SearchResult + return SearchResult(f"text for {title}", {"doc_title": title, "docs_type": "general_docs"}, 0.9) + + +def test_search_docs_defaults_to_legacy_pinecone_backend(monkeypatch): + monkeypatch.delenv("DOCSITE_SEARCH_BACKEND", raising=False) + monkeypatch.delenv("DOCSITE_SHADOW_POSTGRES", raising=False) + + with patch.object(rd, "LegacyPineconeDocsiteSearch") as mock_legacy_cls: + mock_legacy_cls.return_value.search.return_value = [_fake_result("A")] + results = search_docs([{"query": "q"}], top_k=5) + + assert [r.metadata["doc_title"] for r in results] == ["A"] + mock_legacy_cls.return_value.search.assert_called_once_with("q", top_k=5, strategy="semantic", docs_type="general_docs") + + +def test_search_docs_switches_to_postgres_backend_when_flagged(monkeypatch): + monkeypatch.setenv("DOCSITE_SEARCH_BACKEND", "postgres") + monkeypatch.delenv("DOCSITE_SHADOW_POSTGRES", raising=False) + + with patch.object(rd, "DocsiteSearch") as mock_pg_cls: + mock_pg_cls.return_value.search.return_value = [_fake_result("B")] + results = search_docs([{"query": "q"}], top_k=5) + + assert [r.metadata["doc_title"] for r in results] == ["B"] + mock_pg_cls.return_value.search.assert_called_once_with("q", top_k=5, strategy="hybrid", docs_type="general_docs") + + +def test_search_docs_shadow_mode_logs_comparison_without_changing_result(monkeypatch): + monkeypatch.delenv("DOCSITE_SEARCH_BACKEND", raising=False) + monkeypatch.setenv("DOCSITE_SHADOW_POSTGRES", "true") + + with patch.object(rd, "LegacyPineconeDocsiteSearch") as mock_legacy_cls, \ + patch.object(rd, "DocsiteSearch") as mock_pg_cls: + mock_legacy_cls.return_value.search.return_value = [_fake_result("A")] + mock_pg_cls.return_value.search.return_value = [_fake_result("A")] + results = search_docs([{"query": "q"}], top_k=5) + + # Primary (Pinecone) result is what's returned, unaffected by the shadow call + assert [r.metadata["doc_title"] for r in results] == ["A"] + mock_pg_cls.return_value.search.assert_called_once_with("q", top_k=5, strategy="hybrid", docs_type="general_docs") + + +def test_search_docs_shadow_mode_swallows_postgres_errors(monkeypatch): + monkeypatch.delenv("DOCSITE_SEARCH_BACKEND", raising=False) + monkeypatch.setenv("DOCSITE_SHADOW_POSTGRES", "true") + + with patch.object(rd, "LegacyPineconeDocsiteSearch") as mock_legacy_cls, \ + patch.object(rd, "DocsiteSearch") as mock_pg_cls: + mock_legacy_cls.return_value.search.return_value = [_fake_result("A")] + mock_pg_cls.return_value.search.side_effect = RuntimeError("boom") + results = search_docs([{"query": "q"}], top_k=5) # must not raise + + assert [r.metadata["doc_title"] for r in results] == ["A"] diff --git a/services/tools/search_documentation/search_documentation.py b/services/tools/search_documentation/search_documentation.py index 609c5aa8..676fe681 100644 --- a/services/tools/search_documentation/search_documentation.py +++ b/services/tools/search_documentation/search_documentation.py @@ -7,15 +7,16 @@ """ import os import sys -from pathlib import Path -from typing import Dict, List, Optional from dataclasses import dataclass +from pathlib import Path +from typing import Dict # Import utilities from services directory sys.path.append(str(Path(__file__).parent.parent.parent)) -from util import create_logger, ApolloError +from search_docsite.pinecone_legacy_search import LegacyPineconeDocsiteSearch from search_docsite.search_docsite import DocsiteSearch +from util import ApolloError, create_logger logger = create_logger(__name__) @@ -46,16 +47,13 @@ def _search_implementation(query: str, num_results: int) -> Dict: """ logger.info(f"Searching documentation for: {query[:100]}...") - # Initialize docsite search - docsite_search = DocsiteSearch() - - # Search with threshold for quality results - search_results = docsite_search.search( - query=query, - top_k=num_results, - threshold=0.7, # Only return relevant results - strategy='semantic' - ) + backend = os.environ.get("DOCSITE_SEARCH_BACKEND", "pinecone") + if backend == "postgres": + docsite_search = DocsiteSearch() + search_results = docsite_search.search(query=query, top_k=num_results, strategy='hybrid') + else: + docsite_search = LegacyPineconeDocsiteSearch() + search_results = docsite_search.search(query=query, top_k=num_results, strategy='semantic') logger.info(f"Found {len(search_results)} documentation results") diff --git a/services/tools/search_documentation/tests/unit/conftest.py b/services/tools/search_documentation/tests/unit/conftest.py new file mode 100644 index 00000000..32aac060 --- /dev/null +++ b/services/tools/search_documentation/tests/unit/conftest.py @@ -0,0 +1,14 @@ +"""Test config for search_documentation unit tests. + +Importing `search_documentation.search_documentation` pulls in +`search_docsite.pinecone_legacy_search`, whose module-level `OpenAIEmbeddings()` +default arg validates credentials at construction (openai 2.x / langchain-openai 1.x). +A key must therefore exist at import time. + +Dummy placeholders only. +""" + +import os + +os.environ.setdefault("OPENAI_API_KEY", "sk-test-dummy") +os.environ.setdefault("PINECONE_API_KEY", "pc-test-dummy") diff --git a/services/tools/search_documentation/tests/unit/test_search_documentation.py b/services/tools/search_documentation/tests/unit/test_search_documentation.py new file mode 100644 index 00000000..915afef5 --- /dev/null +++ b/services/tools/search_documentation/tests/unit/test_search_documentation.py @@ -0,0 +1,26 @@ +"""Unit tests for search_documentation's backend-flag selection. +No prior tests existed for this module.""" + +from unittest.mock import patch + +import tools.search_documentation.search_documentation as m + + +def test_search_implementation_uses_legacy_pinecone_by_default(monkeypatch): + monkeypatch.delenv("DOCSITE_SEARCH_BACKEND", raising=False) + + with patch.object(m, "LegacyPineconeDocsiteSearch") as mock_legacy_cls: + mock_legacy_cls.return_value.search.return_value = [] + m._search_implementation("how do I use webhooks", 5) + + mock_legacy_cls.return_value.search.assert_called_once_with(query="how do I use webhooks", top_k=5, strategy="semantic") + + +def test_search_implementation_switches_to_postgres_when_flagged(monkeypatch): + monkeypatch.setenv("DOCSITE_SEARCH_BACKEND", "postgres") + + with patch.object(m, "DocsiteSearch") as mock_pg_cls: + mock_pg_cls.return_value.search.return_value = [] + m._search_implementation("how do I use webhooks", 5) + + mock_pg_cls.return_value.search.assert_called_once_with(query="how do I use webhooks", top_k=5, strategy="hybrid") From 39191729592b91a9ff43201682f757e04a7532f5 Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Sat, 25 Jul 2026 13:45:04 +0800 Subject: [PATCH 08/48] fix: correct shadow-mode docstrings and test import placement --- services/job_chat/retrieve_docs.py | 11 +++++++---- services/job_chat/tests/unit/test_retrieve_docs.py | 5 ++--- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/services/job_chat/retrieve_docs.py b/services/job_chat/retrieve_docs.py index f5b600a1..19be0b2b 100644 --- a/services/job_chat/retrieve_docs.py +++ b/services/job_chat/retrieve_docs.py @@ -173,8 +173,10 @@ def search_docs(search_queries, top_k): """Search the docsite store using search queries. Defaults to the legacy Pinecone backend; set DOCSITE_SEARCH_BACKEND=postgres to switch to the Postgres-backed hybrid search. Set DOCSITE_SHADOW_POSTGRES=true to also run - the Postgres backend in parallel (non-blocking to the response) and log a - comparison against the authoritative backend.""" + the Postgres backend afterward (synchronously — this adds latency to the + response for the duration of the shadow-mode verification window) and log a + comparison against the authoritative backend. A future improvement could + make this call async/fire-and-forget to remove the added latency.""" backend = os.environ.get("DOCSITE_SEARCH_BACKEND", "pinecone") if backend == "postgres": @@ -197,8 +199,9 @@ def _run_backend_search(backend_cls, strategy, search_queries, top_k): def _log_shadow_comparison(search_queries, top_k, primary_results): - """Best-effort: run the Postgres hybrid backend in parallel and log a - comparison against the authoritative (Pinecone) results. Never raises.""" + """Best-effort: run the Postgres hybrid backend synchronously (adds latency + to this request while shadow mode is on) and log a comparison against the + authoritative (Pinecone) results. Never raises.""" try: start = time.time() shadow_results = _run_backend_search(DocsiteSearch, "hybrid", search_queries, top_k) diff --git a/services/job_chat/tests/unit/test_retrieve_docs.py b/services/job_chat/tests/unit/test_retrieve_docs.py index 73159433..1d271c0c 100644 --- a/services/job_chat/tests/unit/test_retrieve_docs.py +++ b/services/job_chat/tests/unit/test_retrieve_docs.py @@ -15,7 +15,9 @@ import pytest +from embeddings.embeddings import SearchResult from job_chat import retrieve_docs as rd +from job_chat.retrieve_docs import search_docs from util import ApolloError @@ -80,11 +82,8 @@ def test_call_llm_wraps_unexpected_error_as_apollo_error(): # --- search_docs (backend flag + shadow mode) ----------------------------------- -from job_chat.retrieve_docs import search_docs - def _fake_result(title): - from embeddings.embeddings import SearchResult return SearchResult(f"text for {title}", {"doc_title": title, "docs_type": "general_docs"}, 0.9) From 0bb3ade1ce5ccf7fe1d73d89daff65758ee81c57 Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Sat, 25 Jul 2026 15:00:31 +0800 Subject: [PATCH 09/48] feat: add offline golden-query recall/latency eval script --- .../search_docsite/tests/eval/__init__.py | 0 .../tests/eval/golden_queries.yaml | 27 ++++++ .../search_docsite/tests/eval/run_eval.py | 94 +++++++++++++++++++ .../tests/unit/test_run_eval.py | 66 +++++++++++++ 4 files changed, 187 insertions(+) create mode 100644 services/search_docsite/tests/eval/__init__.py create mode 100644 services/search_docsite/tests/eval/golden_queries.yaml create mode 100644 services/search_docsite/tests/eval/run_eval.py create mode 100644 services/search_docsite/tests/unit/test_run_eval.py diff --git a/services/search_docsite/tests/eval/__init__.py b/services/search_docsite/tests/eval/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/services/search_docsite/tests/eval/golden_queries.yaml b/services/search_docsite/tests/eval/golden_queries.yaml new file mode 100644 index 00000000..812849d6 --- /dev/null +++ b/services/search_docsite/tests/eval/golden_queries.yaml @@ -0,0 +1,27 @@ +# Starter set of representative job_chat/general-docs queries. Queries with an +# empty expected_doc_titles list are still run and reported (so this fixture +# is usable immediately), but excluded from the recall@k aggregate until a +# human curates the expected doc title(s) by inspecting real search results +# against the live docsite content (neither this script's author nor the +# implementer of this plan has that access). +queries: + - query: "how do I configure a webhook trigger" + expected_doc_titles: [] + - query: "how do I set up a cron trigger" + expected_doc_titles: [] + - query: "what is a run in OpenFn" + expected_doc_titles: [] + - query: "how do I use collections to store state between runs" + expected_doc_titles: [] + - query: "how do I configure a credential for an adaptor" + expected_doc_titles: [] + - query: "what is the difference between a job and a workflow" + expected_doc_titles: [] + - query: "how do I deploy a project using the CLI" + expected_doc_titles: [] + - query: "how do I debug a failed run" + expected_doc_titles: [] + - query: "what adaptors are available for HTTP requests" + expected_doc_titles: [] + - query: "how do I write a data transform function" + expected_doc_titles: [] diff --git a/services/search_docsite/tests/eval/run_eval.py b/services/search_docsite/tests/eval/run_eval.py new file mode 100644 index 00000000..d52eb1a5 --- /dev/null +++ b/services/search_docsite/tests/eval/run_eval.py @@ -0,0 +1,94 @@ +"""Offline recall@k + latency comparison between docsite search backends. + +Usage: poetry run python -m search_docsite.tests.eval.run_eval + +Compares the Postgres-backed DocsiteSearch (strategy='hybrid') against +LegacyPineconeDocsiteSearch (strategy='semantic') over the golden query set, +per the Phase-1 shadow-mode rollout plan: Postgres must match or beat +Pinecone's recall@5 (with <=1-query regression tolerance) and p95 latency +before DOCSITE_SEARCH_BACKEND is flipped to 'postgres' by default. +""" + +import time +from pathlib import Path + +import yaml +from search_docsite.pinecone_legacy_search import LegacyPineconeDocsiteSearch +from search_docsite.search_docsite import DocsiteSearch + +GOLDEN_QUERIES_PATH = Path(__file__).parent / "golden_queries.yaml" + + +def compute_recall_at_k(retrieved_titles: list, expected_titles: list) -> bool: + """True if any expected title appears among the retrieved titles.""" + return bool(set(retrieved_titles) & set(expected_titles)) + + +def run_eval(golden_queries: list, backend_cls, strategy: str, top_k: int = 5) -> dict: # noqa: ANN001 + """Run every golden query against one backend/strategy and return a report dict.""" + backend = backend_cls() + per_query = [] + latencies = [] + scored = 0 + skipped = 0 + hits = 0 + + for item in golden_queries: + query = item["query"] + expected_titles = item.get("expected_doc_titles", []) + + start = time.time() + results = backend.search(query, top_k=top_k, strategy=strategy, docs_type="general_docs") + elapsed = time.time() - start + latencies.append(elapsed) + + retrieved_titles = [r.metadata.get("doc_title") for r in results] + + if expected_titles: + hit = compute_recall_at_k(retrieved_titles, expected_titles) + scored += 1 + hits += int(hit) + else: + hit = None + skipped += 1 + + per_query.append({ + "query": query, + "retrieved_titles": retrieved_titles, + "expected_titles": expected_titles, + "hit": hit, + "latency_s": elapsed, + }) + + latencies_sorted = sorted(latencies) + p50 = latencies_sorted[len(latencies_sorted) // 2] if latencies_sorted else 0.0 + p95_index = min(len(latencies_sorted) - 1, int(len(latencies_sorted) * 0.95)) if latencies_sorted else 0 + p95 = latencies_sorted[p95_index] if latencies_sorted else 0.0 + + return { + "recall_at_k": (hits / scored) if scored else None, + "queries_scored": scored, + "queries_skipped": skipped, + "p50_latency_s": p50, + "p95_latency_s": p95, + "per_query": per_query, + } + + +def main() -> None: + with open(GOLDEN_QUERIES_PATH) as f: # noqa: PTH123 + golden_queries = yaml.safe_load(f)["queries"] + + postgres_report = run_eval(golden_queries, DocsiteSearch, strategy="hybrid") + pinecone_report = run_eval(golden_queries, LegacyPineconeDocsiteSearch, strategy="semantic") + + print(f"Postgres (hybrid): recall@5={postgres_report['recall_at_k']} " # noqa: T201 + f"(scored={postgres_report['queries_scored']}, skipped={postgres_report['queries_skipped']}) " + f"p50={postgres_report['p50_latency_s']:.3f}s p95={postgres_report['p95_latency_s']:.3f}s") + print(f"Pinecone (semantic): recall@5={pinecone_report['recall_at_k']} " # noqa: T201 + f"(scored={pinecone_report['queries_scored']}, skipped={pinecone_report['queries_skipped']}) " + f"p50={pinecone_report['p50_latency_s']:.3f}s p95={pinecone_report['p95_latency_s']:.3f}s") + + +if __name__ == "__main__": + main() diff --git a/services/search_docsite/tests/unit/test_run_eval.py b/services/search_docsite/tests/unit/test_run_eval.py new file mode 100644 index 00000000..fbd67c1d --- /dev/null +++ b/services/search_docsite/tests/unit/test_run_eval.py @@ -0,0 +1,66 @@ +"""Unit tests for the offline golden-query eval's scoring logic. +Backends are fully faked — no real search/DB/network involved.""" + +from unittest.mock import MagicMock + +from search_docsite.tests.eval.run_eval import compute_recall_at_k, run_eval + + +def test_compute_recall_at_k_true_when_any_expected_title_present(): # noqa: ANN201 + assert compute_recall_at_k(["Doc A", "Doc B"], ["Doc B", "Doc C"]) is True + + +def test_compute_recall_at_k_false_when_no_overlap(): # noqa: ANN201 + assert compute_recall_at_k(["Doc A"], ["Doc Z"]) is False + + +def make_fake_backend(titles_by_query): # noqa: ANN001, ANN201 + """Fake backend whose .search() returns SearchResults with the given doc_titles per query.""" + from embeddings.embeddings import SearchResult # noqa: PLC0415 + + def fake_search(query, top_k=None, strategy=None, docs_type=None): # noqa: ANN001, ANN202, ARG001 + return [SearchResult(f"text-{t}", {"doc_title": t, "docs_type": "general_docs"}, 0.9) for t in titles_by_query.get(query, [])] + + backend = MagicMock() + backend.search.side_effect = fake_search + backend_cls = MagicMock(return_value=backend) + return backend_cls + + +def test_run_eval_scores_labeled_queries_and_skips_unlabeled(): # noqa: ANN201 + golden_queries = [ + {"query": "how do I configure a webhook", "expected_doc_titles": ["Webhooks"]}, + {"query": "what is a run", "expected_doc_titles": []}, # unlabeled — skipped from recall aggregate + ] + backend_cls = make_fake_backend({"how do I configure a webhook": ["Webhooks", "Other Doc"], "what is a run": ["Runs"]}) + + report = run_eval(golden_queries, backend_cls, strategy="hybrid", top_k=5) + + assert report["queries_scored"] == 1 + assert report["queries_skipped"] == 1 + assert report["recall_at_k"] == 1.0 + assert len(report["per_query"]) == 2 # noqa: PLR2004 + + +def test_run_eval_computes_recall_across_multiple_labeled_queries(): # noqa: ANN201 + golden_queries = [ + {"query": "q1", "expected_doc_titles": ["A"]}, + {"query": "q2", "expected_doc_titles": ["Z"]}, # backend won't return Z -> miss + ] + backend_cls = make_fake_backend({"q1": ["A"], "q2": ["B"]}) + + report = run_eval(golden_queries, backend_cls, strategy="semantic", top_k=5) + + assert report["queries_scored"] == 2 # noqa: PLR2004 + assert report["recall_at_k"] == 0.5 # noqa: PLR2004 + + +def test_run_eval_reports_latency_percentiles(): # noqa: ANN201 + golden_queries = [{"query": "q1", "expected_doc_titles": ["A"]}] + backend_cls = make_fake_backend({"q1": ["A"]}) + + report = run_eval(golden_queries, backend_cls, strategy="hybrid", top_k=5) + + assert "p50_latency_s" in report + assert "p95_latency_s" in report + assert report["p50_latency_s"] >= 0 From 1d780db2dcfdfa40ffde340221f7d973d377b7f1 Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Sat, 25 Jul 2026 15:21:02 +0800 Subject: [PATCH 10/48] fix: remove orphaned threshold arg and revert unrequested noqa suppressions --- .../integration/test_adaptor_docs_pipeline.py | 2 +- .../search_docsite/tests/eval/run_eval.py | 8 +++---- .../tests/unit/test_run_eval.py | 22 +++++++++---------- 3 files changed, 16 insertions(+), 16 deletions(-) diff --git a/services/job_chat/tests/integration/test_adaptor_docs_pipeline.py b/services/job_chat/tests/integration/test_adaptor_docs_pipeline.py index 78ec9979..8b668098 100644 --- a/services/job_chat/tests/integration/test_adaptor_docs_pipeline.py +++ b/services/job_chat/tests/integration/test_adaptor_docs_pipeline.py @@ -192,7 +192,7 @@ def test_search_docs_returns_general_docs_only(): from job_chat.retrieve_docs import search_docs queries = [{"query": "http adaptor merge() function"}] - results = search_docs(queries, top_k=3, threshold=0.5) + results = search_docs(queries, top_k=3) print(results) assert isinstance(results, list), "Should return a list" diff --git a/services/search_docsite/tests/eval/run_eval.py b/services/search_docsite/tests/eval/run_eval.py index d52eb1a5..23b341ea 100644 --- a/services/search_docsite/tests/eval/run_eval.py +++ b/services/search_docsite/tests/eval/run_eval.py @@ -24,7 +24,7 @@ def compute_recall_at_k(retrieved_titles: list, expected_titles: list) -> bool: return bool(set(retrieved_titles) & set(expected_titles)) -def run_eval(golden_queries: list, backend_cls, strategy: str, top_k: int = 5) -> dict: # noqa: ANN001 +def run_eval(golden_queries: list, backend_cls, strategy: str, top_k: int = 5) -> dict: """Run every golden query against one backend/strategy and return a report dict.""" backend = backend_cls() per_query = [] @@ -76,16 +76,16 @@ def run_eval(golden_queries: list, backend_cls, strategy: str, top_k: int = 5) - def main() -> None: - with open(GOLDEN_QUERIES_PATH) as f: # noqa: PTH123 + with open(GOLDEN_QUERIES_PATH) as f: golden_queries = yaml.safe_load(f)["queries"] postgres_report = run_eval(golden_queries, DocsiteSearch, strategy="hybrid") pinecone_report = run_eval(golden_queries, LegacyPineconeDocsiteSearch, strategy="semantic") - print(f"Postgres (hybrid): recall@5={postgres_report['recall_at_k']} " # noqa: T201 + print(f"Postgres (hybrid): recall@5={postgres_report['recall_at_k']} " f"(scored={postgres_report['queries_scored']}, skipped={postgres_report['queries_skipped']}) " f"p50={postgres_report['p50_latency_s']:.3f}s p95={postgres_report['p95_latency_s']:.3f}s") - print(f"Pinecone (semantic): recall@5={pinecone_report['recall_at_k']} " # noqa: T201 + print(f"Pinecone (semantic): recall@5={pinecone_report['recall_at_k']} " f"(scored={pinecone_report['queries_scored']}, skipped={pinecone_report['queries_skipped']}) " f"p50={pinecone_report['p50_latency_s']:.3f}s p95={pinecone_report['p95_latency_s']:.3f}s") diff --git a/services/search_docsite/tests/unit/test_run_eval.py b/services/search_docsite/tests/unit/test_run_eval.py index fbd67c1d..e76800a8 100644 --- a/services/search_docsite/tests/unit/test_run_eval.py +++ b/services/search_docsite/tests/unit/test_run_eval.py @@ -6,19 +6,19 @@ from search_docsite.tests.eval.run_eval import compute_recall_at_k, run_eval -def test_compute_recall_at_k_true_when_any_expected_title_present(): # noqa: ANN201 +def test_compute_recall_at_k_true_when_any_expected_title_present(): assert compute_recall_at_k(["Doc A", "Doc B"], ["Doc B", "Doc C"]) is True -def test_compute_recall_at_k_false_when_no_overlap(): # noqa: ANN201 +def test_compute_recall_at_k_false_when_no_overlap(): assert compute_recall_at_k(["Doc A"], ["Doc Z"]) is False -def make_fake_backend(titles_by_query): # noqa: ANN001, ANN201 +def make_fake_backend(titles_by_query): """Fake backend whose .search() returns SearchResults with the given doc_titles per query.""" - from embeddings.embeddings import SearchResult # noqa: PLC0415 + from embeddings.embeddings import SearchResult - def fake_search(query, top_k=None, strategy=None, docs_type=None): # noqa: ANN001, ANN202, ARG001 + def fake_search(query, top_k=None, strategy=None, docs_type=None): return [SearchResult(f"text-{t}", {"doc_title": t, "docs_type": "general_docs"}, 0.9) for t in titles_by_query.get(query, [])] backend = MagicMock() @@ -27,7 +27,7 @@ def fake_search(query, top_k=None, strategy=None, docs_type=None): # noqa: ANN0 return backend_cls -def test_run_eval_scores_labeled_queries_and_skips_unlabeled(): # noqa: ANN201 +def test_run_eval_scores_labeled_queries_and_skips_unlabeled(): golden_queries = [ {"query": "how do I configure a webhook", "expected_doc_titles": ["Webhooks"]}, {"query": "what is a run", "expected_doc_titles": []}, # unlabeled — skipped from recall aggregate @@ -39,10 +39,10 @@ def test_run_eval_scores_labeled_queries_and_skips_unlabeled(): # noqa: ANN201 assert report["queries_scored"] == 1 assert report["queries_skipped"] == 1 assert report["recall_at_k"] == 1.0 - assert len(report["per_query"]) == 2 # noqa: PLR2004 + assert len(report["per_query"]) == 2 -def test_run_eval_computes_recall_across_multiple_labeled_queries(): # noqa: ANN201 +def test_run_eval_computes_recall_across_multiple_labeled_queries(): golden_queries = [ {"query": "q1", "expected_doc_titles": ["A"]}, {"query": "q2", "expected_doc_titles": ["Z"]}, # backend won't return Z -> miss @@ -51,11 +51,11 @@ def test_run_eval_computes_recall_across_multiple_labeled_queries(): # noqa: AN report = run_eval(golden_queries, backend_cls, strategy="semantic", top_k=5) - assert report["queries_scored"] == 2 # noqa: PLR2004 - assert report["recall_at_k"] == 0.5 # noqa: PLR2004 + assert report["queries_scored"] == 2 + assert report["recall_at_k"] == 0.5 -def test_run_eval_reports_latency_percentiles(): # noqa: ANN201 +def test_run_eval_reports_latency_percentiles(): golden_queries = [{"query": "q1", "expected_doc_titles": ["A"]}] backend_cls = make_fake_backend({"q1": ["A"]}) From 34d47cd001a51d80470570b7180154a09af034a5 Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Tue, 28 Jul 2026 12:54:56 +0800 Subject: [PATCH 11/48] fix: Restore back threshold to 0.8 --- services/job_chat/rag.yaml | 1 + services/job_chat/retrieve_docs.py | 20 ++++++++----- .../job_chat/tests/unit/test_retrieve_docs.py | 28 +++++++++++++++++-- 3 files changed, 39 insertions(+), 10 deletions(-) diff --git a/services/job_chat/rag.yaml b/services/job_chat/rag.yaml index dbfe8cf7..13caa0de 100644 --- a/services/job_chat/rag.yaml +++ b/services/job_chat/rag.yaml @@ -2,6 +2,7 @@ config_version: 1.0 model: "claude-opus" llm_search_decision: "claude-sonnet" llm_retrieval: "claude-sonnet" +threshold: 0.8 top_k: 5 temperature: 0 prompts_version: 1.0 diff --git a/services/job_chat/retrieve_docs.py b/services/job_chat/retrieve_docs.py index 19be0b2b..7490807f 100644 --- a/services/job_chat/retrieve_docs.py +++ b/services/job_chat/retrieve_docs.py @@ -77,7 +77,7 @@ def retrieve_knowledge(content, history, code="", adaptor="", api_key=None, stre search_queries, generate_queries_usage = generate_queries(content, client, user_context) with sentry_sdk.start_span(description="search_documentation"): try: - search_results = search_docs(search_queries, top_k=config["top_k"]) + search_results = search_docs(search_queries, top_k=config["top_k"], threshold=config["threshold"]) search_results = list(set(search_results)) search_results_sections = list(set(result.metadata["doc_title"] for result in search_results)) except Exception as e: @@ -169,20 +169,26 @@ def generate_queries(content, client, user_context=""): return (answer_parsed, usage) -def search_docs(search_queries, top_k): +def search_docs(search_queries, top_k, threshold=None): """Search the docsite store using search queries. Defaults to the legacy Pinecone backend; set DOCSITE_SEARCH_BACKEND=postgres to switch to the Postgres-backed hybrid search. Set DOCSITE_SHADOW_POSTGRES=true to also run the Postgres backend afterward (synchronously — this adds latency to the response for the duration of the shadow-mode verification window) and log a comparison against the authoritative backend. A future improvement could - make this call async/fire-and-forget to remove the added latency.""" + make this call async/fire-and-forget to remove the added latency. + + :param threshold: Score threshold below which results are dropped. Only + meaningful for the Pinecone/semantic strategy (a cosine-similarity + cutoff); the Postgres hybrid strategy ignores it (RRF-fused rank isn't + a comparable score), so it's forwarded unconditionally and has no + effect there.""" backend = os.environ.get("DOCSITE_SEARCH_BACKEND", "pinecone") if backend == "postgres": - primary_results = _run_backend_search(DocsiteSearch, "hybrid", search_queries, top_k) + primary_results = _run_backend_search(DocsiteSearch, "hybrid", search_queries, top_k, threshold) else: - primary_results = _run_backend_search(LegacyPineconeDocsiteSearch, "semantic", search_queries, top_k) + primary_results = _run_backend_search(LegacyPineconeDocsiteSearch, "semantic", search_queries, top_k, threshold) if os.environ.get("DOCSITE_SHADOW_POSTGRES", "").lower() == "true" and backend != "postgres": _log_shadow_comparison(search_queries, top_k, primary_results) @@ -190,11 +196,11 @@ def search_docs(search_queries, top_k): return primary_results -def _run_backend_search(backend_cls, strategy, search_queries, top_k): +def _run_backend_search(backend_cls, strategy, search_queries, top_k, threshold=None): searcher = backend_cls() results = [] for q in search_queries: - results.extend(searcher.search(q.get("query"), top_k=top_k, strategy=strategy, docs_type="general_docs")) + results.extend(searcher.search(q.get("query"), top_k=top_k, threshold=threshold, strategy=strategy, docs_type="general_docs")) return results diff --git a/services/job_chat/tests/unit/test_retrieve_docs.py b/services/job_chat/tests/unit/test_retrieve_docs.py index 1d271c0c..77a7b3aa 100644 --- a/services/job_chat/tests/unit/test_retrieve_docs.py +++ b/services/job_chat/tests/unit/test_retrieve_docs.py @@ -96,7 +96,25 @@ def test_search_docs_defaults_to_legacy_pinecone_backend(monkeypatch): results = search_docs([{"query": "q"}], top_k=5) assert [r.metadata["doc_title"] for r in results] == ["A"] - mock_legacy_cls.return_value.search.assert_called_once_with("q", top_k=5, strategy="semantic", docs_type="general_docs") + mock_legacy_cls.return_value.search.assert_called_once_with( + "q", top_k=5, threshold=None, strategy="semantic", docs_type="general_docs" + ) + + +def test_search_docs_passes_threshold_through_to_semantic_backend(monkeypatch): + """Threshold is a score cutoff that only makes sense for the Pinecone/semantic + path; it must still be forwarded there (this regressed once already when the + backend flag was introduced — rag.yaml's threshold silently stopped applying).""" + monkeypatch.delenv("DOCSITE_SEARCH_BACKEND", raising=False) + monkeypatch.delenv("DOCSITE_SHADOW_POSTGRES", raising=False) + + with patch.object(rd, "LegacyPineconeDocsiteSearch") as mock_legacy_cls: + mock_legacy_cls.return_value.search.return_value = [_fake_result("A")] + search_docs([{"query": "q"}], top_k=5, threshold=0.8) + + mock_legacy_cls.return_value.search.assert_called_once_with( + "q", top_k=5, threshold=0.8, strategy="semantic", docs_type="general_docs" + ) def test_search_docs_switches_to_postgres_backend_when_flagged(monkeypatch): @@ -108,7 +126,9 @@ def test_search_docs_switches_to_postgres_backend_when_flagged(monkeypatch): results = search_docs([{"query": "q"}], top_k=5) assert [r.metadata["doc_title"] for r in results] == ["B"] - mock_pg_cls.return_value.search.assert_called_once_with("q", top_k=5, strategy="hybrid", docs_type="general_docs") + mock_pg_cls.return_value.search.assert_called_once_with( + "q", top_k=5, threshold=None, strategy="hybrid", docs_type="general_docs" + ) def test_search_docs_shadow_mode_logs_comparison_without_changing_result(monkeypatch): @@ -123,7 +143,9 @@ def test_search_docs_shadow_mode_logs_comparison_without_changing_result(monkeyp # Primary (Pinecone) result is what's returned, unaffected by the shadow call assert [r.metadata["doc_title"] for r in results] == ["A"] - mock_pg_cls.return_value.search.assert_called_once_with("q", top_k=5, strategy="hybrid", docs_type="general_docs") + mock_pg_cls.return_value.search.assert_called_once_with( + "q", top_k=5, threshold=None, strategy="hybrid", docs_type="general_docs" + ) def test_search_docs_shadow_mode_swallows_postgres_errors(monkeypatch): From 10500064580760a08f0c5a380a8cc9772aa404bc Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Tue, 28 Jul 2026 14:06:14 +0800 Subject: [PATCH 12/48] fix: apply the 0.7 relevance gate to both search_documentation backends --- .../search_documentation.py | 19 +++++++++++++------ .../tests/unit/test_search_documentation.py | 16 +++++++++++++--- 2 files changed, 26 insertions(+), 9 deletions(-) diff --git a/services/tools/search_documentation/search_documentation.py b/services/tools/search_documentation/search_documentation.py index 676fe681..fd0f61cc 100644 --- a/services/tools/search_documentation/search_documentation.py +++ b/services/tools/search_documentation/search_documentation.py @@ -47,13 +47,20 @@ def _search_implementation(query: str, num_results: int) -> Dict: """ logger.info(f"Searching documentation for: {query[:100]}...") + # Both backends use semantic search with the same cosine-similarity cutoff, so + # results are directly comparable. Hybrid (RRF) is deliberately not used here: + # its score has no calibratable scale, so it can be neither thresholded nor + # rendered as a relevance figure. It stays available via the search_docsite + # service and run_eval for evaluation. backend = os.environ.get("DOCSITE_SEARCH_BACKEND", "pinecone") - if backend == "postgres": - docsite_search = DocsiteSearch() - search_results = docsite_search.search(query=query, top_k=num_results, strategy='hybrid') - else: - docsite_search = LegacyPineconeDocsiteSearch() - search_results = docsite_search.search(query=query, top_k=num_results, strategy='semantic') + search_cls = DocsiteSearch if backend == "postgres" else LegacyPineconeDocsiteSearch + docsite_search = search_cls() + search_results = docsite_search.search( + query=query, + top_k=num_results, + threshold=0.7, # Only return relevant results + strategy='semantic', + ) logger.info(f"Found {len(search_results)} documentation results") diff --git a/services/tools/search_documentation/tests/unit/test_search_documentation.py b/services/tools/search_documentation/tests/unit/test_search_documentation.py index 915afef5..20689e18 100644 --- a/services/tools/search_documentation/tests/unit/test_search_documentation.py +++ b/services/tools/search_documentation/tests/unit/test_search_documentation.py @@ -7,20 +7,30 @@ def test_search_implementation_uses_legacy_pinecone_by_default(monkeypatch): + """Characterization test: pins the exact call main made. The 0.7 threshold is + a quality gate on live traffic — it was silently dropped once already.""" monkeypatch.delenv("DOCSITE_SEARCH_BACKEND", raising=False) with patch.object(m, "LegacyPineconeDocsiteSearch") as mock_legacy_cls: mock_legacy_cls.return_value.search.return_value = [] m._search_implementation("how do I use webhooks", 5) - mock_legacy_cls.return_value.search.assert_called_once_with(query="how do I use webhooks", top_k=5, strategy="semantic") + mock_legacy_cls.return_value.search.assert_called_once_with( + query="how do I use webhooks", top_k=5, threshold=0.7, strategy="semantic" + ) -def test_search_implementation_switches_to_postgres_when_flagged(monkeypatch): +def test_search_implementation_uses_semantic_with_same_threshold_on_postgres(monkeypatch): + """Postgres must apply the identical quality gate. Hybrid's RRF score has a + ~0.033 ceiling — it cannot be thresholded, and renders as a constant 0.03 to + the LLM. Semantic returns true cosine similarity, directly comparable to + Pinecone, so the same 0.7 cutoff means the same thing on both backends.""" monkeypatch.setenv("DOCSITE_SEARCH_BACKEND", "postgres") with patch.object(m, "DocsiteSearch") as mock_pg_cls: mock_pg_cls.return_value.search.return_value = [] m._search_implementation("how do I use webhooks", 5) - mock_pg_cls.return_value.search.assert_called_once_with(query="how do I use webhooks", top_k=5, strategy="hybrid") + mock_pg_cls.return_value.search.assert_called_once_with( + query="how do I use webhooks", top_k=5, threshold=0.7, strategy="semantic" + ) From b28b5fcbe51ff7953901326d0001b479a08e0c4d Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Tue, 28 Jul 2026 14:16:58 +0800 Subject: [PATCH 13/48] refactor: remove shadow-mode Postgres comparison from job_chat --- services/job_chat/retrieve_docs.py | 47 +++---------------- .../job_chat/tests/unit/test_retrieve_docs.py | 28 ----------- 2 files changed, 7 insertions(+), 68 deletions(-) diff --git a/services/job_chat/retrieve_docs.py b/services/job_chat/retrieve_docs.py index 7490807f..93ee0450 100644 --- a/services/job_chat/retrieve_docs.py +++ b/services/job_chat/retrieve_docs.py @@ -1,6 +1,5 @@ import json import os -import time import anthropic import sentry_sdk @@ -172,58 +171,26 @@ def generate_queries(content, client, user_context=""): def search_docs(search_queries, top_k, threshold=None): """Search the docsite store using search queries. Defaults to the legacy Pinecone backend; set DOCSITE_SEARCH_BACKEND=postgres to switch to the - Postgres-backed hybrid search. Set DOCSITE_SHADOW_POSTGRES=true to also run - the Postgres backend afterward (synchronously — this adds latency to the - response for the duration of the shadow-mode verification window) and log a - comparison against the authoritative backend. A future improvement could - make this call async/fire-and-forget to remove the added latency. + Postgres-backed hybrid search. :param threshold: Score threshold below which results are dropped. Only meaningful for the Pinecone/semantic strategy (a cosine-similarity - cutoff); the Postgres hybrid strategy ignores it (RRF-fused rank isn't - a comparable score), so it's forwarded unconditionally and has no - effect there.""" + cutoff).""" backend = os.environ.get("DOCSITE_SEARCH_BACKEND", "pinecone") if backend == "postgres": - primary_results = _run_backend_search(DocsiteSearch, "hybrid", search_queries, top_k, threshold) - else: - primary_results = _run_backend_search(LegacyPineconeDocsiteSearch, "semantic", search_queries, top_k, threshold) - - if os.environ.get("DOCSITE_SHADOW_POSTGRES", "").lower() == "true" and backend != "postgres": - _log_shadow_comparison(search_queries, top_k, primary_results) - - return primary_results + return _run_backend_search(DocsiteSearch, "hybrid", search_queries, top_k, threshold) + return _run_backend_search(LegacyPineconeDocsiteSearch, "semantic", search_queries, top_k, threshold) def _run_backend_search(backend_cls, strategy, search_queries, top_k, threshold=None): searcher = backend_cls() results = [] for q in search_queries: - results.extend(searcher.search(q.get("query"), top_k=top_k, threshold=threshold, strategy=strategy, docs_type="general_docs")) - return results - - -def _log_shadow_comparison(search_queries, top_k, primary_results): - """Best-effort: run the Postgres hybrid backend synchronously (adds latency - to this request while shadow mode is on) and log a comparison against the - authoritative (Pinecone) results. Never raises.""" - try: - start = time.time() - shadow_results = _run_backend_search(DocsiteSearch, "hybrid", search_queries, top_k) - elapsed = time.time() - start - - primary_titles = {r.metadata.get("doc_title") for r in primary_results} - shadow_titles = {r.metadata.get("doc_title") for r in shadow_results} - overlap = len(primary_titles & shadow_titles) - - logger.info( - "docsite_shadow_comparison " - f"primary_count={len(primary_results)} shadow_count={len(shadow_results)} " - f"title_overlap={overlap} shadow_latency_s={elapsed:.3f}" + results.extend( + searcher.search(q.get("query"), top_k=top_k, threshold=threshold, strategy=strategy, docs_type="general_docs"), ) - except Exception as e: - logger.warning(f"Shadow Postgres docsite search failed: {e}") + return results def format_context(adaptor, code, history): """Optionally add more context about the user's job for the LLM.""" diff --git a/services/job_chat/tests/unit/test_retrieve_docs.py b/services/job_chat/tests/unit/test_retrieve_docs.py index 77a7b3aa..e508f984 100644 --- a/services/job_chat/tests/unit/test_retrieve_docs.py +++ b/services/job_chat/tests/unit/test_retrieve_docs.py @@ -131,31 +131,3 @@ def test_search_docs_switches_to_postgres_backend_when_flagged(monkeypatch): ) -def test_search_docs_shadow_mode_logs_comparison_without_changing_result(monkeypatch): - monkeypatch.delenv("DOCSITE_SEARCH_BACKEND", raising=False) - monkeypatch.setenv("DOCSITE_SHADOW_POSTGRES", "true") - - with patch.object(rd, "LegacyPineconeDocsiteSearch") as mock_legacy_cls, \ - patch.object(rd, "DocsiteSearch") as mock_pg_cls: - mock_legacy_cls.return_value.search.return_value = [_fake_result("A")] - mock_pg_cls.return_value.search.return_value = [_fake_result("A")] - results = search_docs([{"query": "q"}], top_k=5) - - # Primary (Pinecone) result is what's returned, unaffected by the shadow call - assert [r.metadata["doc_title"] for r in results] == ["A"] - mock_pg_cls.return_value.search.assert_called_once_with( - "q", top_k=5, threshold=None, strategy="hybrid", docs_type="general_docs" - ) - - -def test_search_docs_shadow_mode_swallows_postgres_errors(monkeypatch): - monkeypatch.delenv("DOCSITE_SEARCH_BACKEND", raising=False) - monkeypatch.setenv("DOCSITE_SHADOW_POSTGRES", "true") - - with patch.object(rd, "LegacyPineconeDocsiteSearch") as mock_legacy_cls, \ - patch.object(rd, "DocsiteSearch") as mock_pg_cls: - mock_legacy_cls.return_value.search.return_value = [_fake_result("A")] - mock_pg_cls.return_value.search.side_effect = RuntimeError("boom") - results = search_docs([{"query": "q"}], top_k=5) # must not raise - - assert [r.metadata["doc_title"] for r in results] == ["A"] From ba0769b4f6881719d834c08e6b7a00a8bf6cb5fa Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Tue, 28 Jul 2026 14:42:44 +0800 Subject: [PATCH 14/48] fix: cast hybrid RRF score to float8 so results stay JSON-serializable --- services/search_docsite/search_docsite.py | 7 +++-- .../tests/unit/test_docsite_search.py | 28 +++++++++++++++++++ 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/services/search_docsite/search_docsite.py b/services/search_docsite/search_docsite.py index dced1ccf..c3010f52 100644 --- a/services/search_docsite/search_docsite.py +++ b/services/search_docsite/search_docsite.py @@ -159,7 +159,7 @@ def _hybrid_search(self, conn, batch_id, query, top_k, doc_title, docs_type): SELECT COALESCE(s.text, k.text) AS text, COALESCE(s.doc_title, k.doc_title) AS doc_title, COALESCE(s.docs_type, k.docs_type) AS docs_type, - COALESCE(1.0 / (60 + s.rnk), 0) + COALESCE(1.0 / (60 + k.rnk), 0) AS score + COALESCE(1.0::float8 / (60 + s.rnk), 0) + COALESCE(1.0::float8 / (60 + k.rnk), 0) AS score FROM semantic s FULL OUTER JOIN keyword k ON s.id = k.id ORDER BY score DESC LIMIT %(max_k)s @@ -172,7 +172,10 @@ def _hybrid_search(self, conn, batch_id, query, top_k, doc_title, docs_type): cur.execute(sql, params) rows = cur.fetchall() - results = [SearchResult(text, {"doc_title": title, "docs_type": dtype}, score) for text, title, dtype, score in rows] + results = [ + SearchResult(text, {"doc_title": title, "docs_type": dtype}, float(score)) + for text, title, dtype, score in rows + ] logger.info(f"Hybrid search returned {len(results)} results") return results diff --git a/services/search_docsite/tests/unit/test_docsite_search.py b/services/search_docsite/tests/unit/test_docsite_search.py index 5e5424a9..a48a8e31 100644 --- a/services/search_docsite/tests/unit/test_docsite_search.py +++ b/services/search_docsite/tests/unit/test_docsite_search.py @@ -5,6 +5,8 @@ lazy `_embeddings` attribute, matching the DocsiteIndexer test pattern. """ +import json +from decimal import Decimal from unittest.mock import MagicMock, patch import pytest @@ -134,3 +136,29 @@ def test_hybrid_search_runs_rrf_query_and_returns_results(): params = cur.execute.call_args[0][1] assert params["candidate_k"] == 50 assert params["max_k"] == 5 + + +def test_hybrid_search_score_is_json_serializable_float(): + """Postgres returns RRF as `numeric`, which psycopg2 hands back as Decimal. + Decimal is not JSON-serializable, and entry.py's json.dump sits outside its + try/except — so this would kill the process, not return a 500.""" + conn, cur = make_conn() + cur.fetchall.return_value = [("a", "Doc A", "general_docs", Decimal("0.032"))] + ds = make_search() + + results = ds._hybrid_search(conn, batch_id=1, query="webhook", top_k=5, doc_title=None, docs_type="general_docs") + + assert isinstance(results[0].score, float) + json.dumps(results[0].to_json()) # must not raise + + +def test_hybrid_search_casts_rrf_to_float8_in_sql(): + """Belt and braces: the SQL itself must not produce numeric in the first place.""" + conn, cur = make_conn() + cur.fetchall.return_value = [] + ds = make_search() + + ds._hybrid_search(conn, batch_id=1, query="webhook", top_k=5, doc_title=None, docs_type=None) + + sql = cur.execute.call_args[0][0] + assert "float8" in sql From 73e84420e53fa3b7e1822ccb54b695e80cdc9c07 Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Tue, 28 Jul 2026 14:51:20 +0800 Subject: [PATCH 15/48] fix: use semantic strategy on Postgres so the relevance gate survives cutover --- services/job_chat/retrieve_docs.py | 19 ++++++++------- .../job_chat/tests/unit/test_retrieve_docs.py | 10 ++++---- services/search_docsite/search_docsite.py | 11 ++++++++- .../tests/unit/test_docsite_search.py | 23 +++++++++++++++++++ 4 files changed, 50 insertions(+), 13 deletions(-) diff --git a/services/job_chat/retrieve_docs.py b/services/job_chat/retrieve_docs.py index 93ee0450..2d7c9cd7 100644 --- a/services/job_chat/retrieve_docs.py +++ b/services/job_chat/retrieve_docs.py @@ -171,16 +171,19 @@ def generate_queries(content, client, user_context=""): def search_docs(search_queries, top_k, threshold=None): """Search the docsite store using search queries. Defaults to the legacy Pinecone backend; set DOCSITE_SEARCH_BACKEND=postgres to switch to the - Postgres-backed hybrid search. + Postgres-backed search. - :param threshold: Score threshold below which results are dropped. Only - meaningful for the Pinecone/semantic strategy (a cosine-similarity - cutoff).""" - backend = os.environ.get("DOCSITE_SEARCH_BACKEND", "pinecone") + Both backends use semantic search with the same cosine-similarity cutoff, so + results are directly comparable and the quality gate survives the cutover. + Hybrid (RRF) is deliberately not used here: its score has no calibratable + scale, so a threshold cannot be applied to it. It stays available via the + search_docsite service and run_eval for evaluation. - if backend == "postgres": - return _run_backend_search(DocsiteSearch, "hybrid", search_queries, top_k, threshold) - return _run_backend_search(LegacyPineconeDocsiteSearch, "semantic", search_queries, top_k, threshold) + :param threshold: Cosine-similarity cutoff, applied identically by both + backends.""" + backend = os.environ.get("DOCSITE_SEARCH_BACKEND", "pinecone") + backend_cls = DocsiteSearch if backend == "postgres" else LegacyPineconeDocsiteSearch + return _run_backend_search(backend_cls, "semantic", search_queries, top_k, threshold) def _run_backend_search(backend_cls, strategy, search_queries, top_k, threshold=None): diff --git a/services/job_chat/tests/unit/test_retrieve_docs.py b/services/job_chat/tests/unit/test_retrieve_docs.py index e508f984..5dcd4ff9 100644 --- a/services/job_chat/tests/unit/test_retrieve_docs.py +++ b/services/job_chat/tests/unit/test_retrieve_docs.py @@ -117,17 +117,19 @@ def test_search_docs_passes_threshold_through_to_semantic_backend(monkeypatch): ) -def test_search_docs_switches_to_postgres_backend_when_flagged(monkeypatch): +def test_search_docs_uses_semantic_with_threshold_on_postgres_backend(monkeypatch): + """Postgres must apply the identical 0.8 cosine gate as Pinecone. Hybrid's RRF + score cannot be thresholded, so using it here would silently drop the quality + filter at cutover. Semantic returns comparable cosine scores.""" monkeypatch.setenv("DOCSITE_SEARCH_BACKEND", "postgres") - monkeypatch.delenv("DOCSITE_SHADOW_POSTGRES", raising=False) with patch.object(rd, "DocsiteSearch") as mock_pg_cls: mock_pg_cls.return_value.search.return_value = [_fake_result("B")] - results = search_docs([{"query": "q"}], top_k=5) + results = search_docs([{"query": "q"}], top_k=5, threshold=0.8) assert [r.metadata["doc_title"] for r in results] == ["B"] mock_pg_cls.return_value.search.assert_called_once_with( - "q", top_k=5, threshold=None, strategy="hybrid", docs_type="general_docs" + "q", top_k=5, threshold=0.8, strategy="semantic", docs_type="general_docs" ) diff --git a/services/search_docsite/search_docsite.py b/services/search_docsite/search_docsite.py index c3010f52..f2b10efc 100644 --- a/services/search_docsite/search_docsite.py +++ b/services/search_docsite/search_docsite.py @@ -39,12 +39,21 @@ def search(self, query, top_k=None, threshold=None, strategy='semantic', doc_tit :param query: Search query string :param top_k: Number of results to return - :param threshold: Score threshold (only meaningful for strategy='semantic') + :param threshold: Score threshold. Only valid for strategy='semantic' — the + keyword (FTS rank) and hybrid (RRF rank) scores are not on a comparable + scale, so passing a threshold with them raises rather than being ignored. :param strategy: 'semantic' | 'keyword' | 'hybrid' (default: 'semantic') :param doc_title: Filter by document title :param docs_type: Filter by document type :return: List of SearchResult objects """ + if threshold is not None and strategy != 'semantic': + raise ApolloError( + 400, + f"threshold is only supported for strategy='semantic', got '{strategy}'", + type="BAD_REQUEST", + ) + conn = get_db_connection() register_vector_type(conn) try: diff --git a/services/search_docsite/tests/unit/test_docsite_search.py b/services/search_docsite/tests/unit/test_docsite_search.py index a48a8e31..8308a74c 100644 --- a/services/search_docsite/tests/unit/test_docsite_search.py +++ b/services/search_docsite/tests/unit/test_docsite_search.py @@ -53,6 +53,29 @@ def test_search_raises_on_unknown_strategy(): assert exc.value.code == 400 +@pytest.mark.parametrize("strategy", ["hybrid", "keyword"]) +def test_search_rejects_threshold_for_non_semantic_strategies(strategy): + """RRF/FTS scores are not comparable to a cosine cutoff. Silently ignoring a + threshold here is a landmine: 0.8 against a 0.033-max score would drop + every result with no error.""" + conn, _ = make_conn() + ds = make_search(batch_id=1) + with patched(conn)[0], patched(conn)[1]: + with pytest.raises(ApolloError) as exc: + ds.search("query", threshold=0.8, strategy=strategy) + assert exc.value.code == 400 + + +@pytest.mark.parametrize("strategy", ["hybrid", "keyword"]) +def test_search_allows_none_threshold_for_non_semantic_strategies(strategy): + conn, _ = make_conn() + ds = make_search(batch_id=1) + with patched(conn)[0], patched(conn)[1], \ + patch.object(ds, "_keyword_search", return_value=["r"]), \ + patch.object(ds, "_hybrid_search", return_value=["r"]): + assert ds.search("query", threshold=None, strategy=strategy) == ["r"] + + # --- _resolve_current_batch -------------------------------------------------- def test_resolve_current_batch_returns_newest_complete_batch_id(): From 5fd8ed8416bbada81a027df5c16c174ab8e57ce5 Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Tue, 28 Jul 2026 15:05:10 +0800 Subject: [PATCH 16/48] feat: add per-request backend selection to search_docsite --- services/search_docsite/README.md | 18 +++-- .../search_docsite/pinecone_legacy_search.py | 10 ++- services/search_docsite/search_docsite.py | 30 +++++--- .../tests/unit/test_pinecone_legacy_search.py | 22 ++++++ .../tests/unit/test_search_docsite_main.py | 77 +++++++++++++++++++ 5 files changed, 141 insertions(+), 16 deletions(-) create mode 100644 services/search_docsite/tests/unit/test_search_docsite_main.py diff --git a/services/search_docsite/README.md b/services/search_docsite/README.md index ecbe0146..c7380344 100644 --- a/services/search_docsite/README.md +++ b/services/search_docsite/README.md @@ -26,16 +26,24 @@ bun py search_docsite tmp/payload.json -O ## Implementation The service uses the DocsiteSearch class to query the database (Pinecone). It embeds semantic search queries using OpenAI. +To compare backends on the same query, run the service twice with different +`backend` values and diff the results. This replaces the shadow-mode comparison +that was considered for the Postgres migration. + ## Payload Reference The input payload is a JSON object with the following structure: ```js { - "query": "What is Asana", // Input query - "collection_name": "Docsite-20250225", // Name of the collection in the vector database - "docs_type": "adaptor_docs", // Filter for document type adaptor_docs, adaptor_functions, general_docs (optional) - "doc_title": "Asana", // Filter for document title (optional) - "top_k": 5 // Adjust the number of search results (optional) + "query": "What is Asana", // Input query (required) + "backend": "pinecone", // 'pinecone' | 'postgres'. Defaults to DOCSITE_SEARCH_BACKEND, itself defaulting to pinecone. + "docs_type": "adaptor_docs", // Filter for adaptor_docs | adaptor_functions | general_docs (optional) + "doc_title": "Asana", // Filter for document title (optional) + "top_k": 5, // Number of search results (optional) + "threshold": 0.8, // Cosine cutoff. Only valid with strategy 'semantic'. (optional) + "strategy": "semantic", // Postgres backend only: 'semantic' | 'keyword' | 'hybrid' + "batch_id": 12, // Postgres backend only: pin a specific batch (optional) + "collection_name": "docsite-..." // Pinecone backend only: pin a namespace (optional) } ``` diff --git a/services/search_docsite/pinecone_legacy_search.py b/services/search_docsite/pinecone_legacy_search.py index de887e99..5718e2aa 100644 --- a/services/search_docsite/pinecone_legacy_search.py +++ b/services/search_docsite/pinecone_legacy_search.py @@ -35,8 +35,14 @@ def search(self, query, top_k=None, threshold=None, strategy='semantic', doc_tit filters = self._build_filter(doc_title=doc_title, docs_type=docs_type) logger.info("Metadata filters built") - if strategy == 'semantic': - return self._semantic_search(query=query, top_k=top_k, threshold=threshold, filters=filters) + if strategy != 'semantic': + raise ApolloError( + 400, + f"The Pinecone backend only supports strategy='semantic', got '{strategy}'", + type="BAD_REQUEST", + ) + + return self._semantic_search(query=query, top_k=top_k, threshold=threshold, filters=filters) def _semantic_search(self, query, top_k=None, threshold=None, filters=None): if top_k is None and threshold is None: diff --git a/services/search_docsite/search_docsite.py b/services/search_docsite/search_docsite.py index f2b10efc..f757a5f7 100644 --- a/services/search_docsite/search_docsite.py +++ b/services/search_docsite/search_docsite.py @@ -4,6 +4,7 @@ from pgvector.psycopg2 import register_vector from util import create_logger, ApolloError, get_db_connection from embeddings.embeddings import SearchResult +from search_docsite.pinecone_legacy_search import LegacyPineconeDocsiteSearch logger = create_logger("DocsiteSearch") @@ -189,6 +190,12 @@ def _hybrid_search(self, conn, batch_id, query, top_k, doc_title, docs_type): return results +BACKEND_INDEX_PARAMS = { + "postgres": ["batch_id", "default_top_k"], + "pinecone": ["collection_name", "index_name", "default_top_k", "embeddings"], +} + + def main(data): logger.info("Starting...") @@ -196,21 +203,23 @@ def main(data): missing = [field for field in required_fields if field not in data] if missing: logger.error(f"Missing required fields in data: {', '.join(missing)}") - return + return None + + backend = data.get("backend") or os.environ.get("DOCSITE_SEARCH_BACKEND", "pinecone") + if backend not in BACKEND_INDEX_PARAMS: + raise ApolloError( + 400, + f"Unknown backend '{backend}'. Expected 'pinecone' or 'postgres'", + type="BAD_REQUEST", + ) - index_params = {} search_params = {"query": data["query"]} - optional_search_params = ["docs_type", "doc_title", "top_k", "threshold", "strategy"] - optional_index_params = ["batch_id", "default_top_k"] - for key in optional_search_params: if key in data: search_params[key] = data[key] - for key in optional_index_params: - if key in data: - index_params[key] = data[key] + index_params = {key: data[key] for key in BACKEND_INDEX_PARAMS[backend] if key in data} load_dotenv(override=True) openai_api_key = os.environ.get('OPENAI_API_KEY') @@ -219,7 +228,10 @@ def main(data): logger.error(msg) raise ApolloError(500, msg, type="BAD_REQUEST") - docsite_search = DocsiteSearch(**index_params) + search_cls = DocsiteSearch if backend == "postgres" else LegacyPineconeDocsiteSearch + logger.info(f"Searching docsite via the {backend} backend") + + docsite_search = search_cls(**index_params) results = docsite_search.search(**search_params) return [result.to_json() for result in results] diff --git a/services/search_docsite/tests/unit/test_pinecone_legacy_search.py b/services/search_docsite/tests/unit/test_pinecone_legacy_search.py index 70978750..5d6b2bf2 100644 --- a/services/search_docsite/tests/unit/test_pinecone_legacy_search.py +++ b/services/search_docsite/tests/unit/test_pinecone_legacy_search.py @@ -138,3 +138,25 @@ def test_get_most_recent_namespace_raises_when_none_valid(): with pytest.raises(ApolloError) as exc: ds._get_most_recent_namespace() assert exc.value.code == 404 + + +# --- strategy guard ------------------------------------------------------------- + +@pytest.mark.parametrize("strategy", ["hybrid", "keyword", "nonsense"]) +def test_legacy_search_raises_on_unsupported_strategy(strategy): + """Previously fell off the end of the method and returned None implicitly, + which surfaces downstream as a confusing TypeError. Reachable now that the + backend is selectable per request.""" + ds = make_search() + with pytest.raises(ApolloError) as exc: + ds.search("query", strategy=strategy) + assert exc.value.code == 400 + + +def test_legacy_search_still_dispatches_semantic(): + ds = make_search() + ds.vectorstore.similarity_search_with_score.return_value = [(FakeDoc("a"), 0.9)] + + results = ds.search("query", strategy="semantic") + + assert [r.text for r in results] == ["a"] diff --git a/services/search_docsite/tests/unit/test_search_docsite_main.py b/services/search_docsite/tests/unit/test_search_docsite_main.py new file mode 100644 index 00000000..2db6d5af --- /dev/null +++ b/services/search_docsite/tests/unit/test_search_docsite_main.py @@ -0,0 +1,77 @@ +"""Unit tests for search_docsite.main's backend selection. + +The `backend` payload field is the shadow-mode replacement: it lets the same +query be run against both backends on demand for manual comparison. +""" + +from unittest.mock import patch + +import pytest + +import search_docsite.search_docsite as m +from util import ApolloError + + +def test_main_defaults_to_pinecone_backend(monkeypatch): + monkeypatch.delenv("DOCSITE_SEARCH_BACKEND", raising=False) + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + + with patch.object(m, "LegacyPineconeDocsiteSearch") as mock_legacy_cls: + mock_legacy_cls.return_value.search.return_value = [] + m.main({"query": "webhooks"}) + + mock_legacy_cls.assert_called_once_with() + mock_legacy_cls.return_value.search.assert_called_once_with(query="webhooks") + + +def test_main_payload_backend_overrides_env(monkeypatch): + monkeypatch.setenv("DOCSITE_SEARCH_BACKEND", "pinecone") + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + + with patch.object(m, "DocsiteSearch") as mock_pg_cls: + mock_pg_cls.return_value.search.return_value = [] + m.main({"query": "webhooks", "backend": "postgres"}) + + mock_pg_cls.return_value.search.assert_called_once_with(query="webhooks") + + +def test_main_env_selects_postgres_when_no_payload_override(monkeypatch): + monkeypatch.setenv("DOCSITE_SEARCH_BACKEND", "postgres") + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + + with patch.object(m, "DocsiteSearch") as mock_pg_cls: + mock_pg_cls.return_value.search.return_value = [] + m.main({"query": "webhooks"}) + + mock_pg_cls.return_value.search.assert_called_once_with(query="webhooks") + + +def test_main_rejects_unknown_backend(monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + + with pytest.raises(ApolloError) as exc: + m.main({"query": "webhooks", "backend": "sqlite"}) + + assert exc.value.code == 400 + + +def test_main_routes_index_params_per_backend(monkeypatch): + """The two classes take different constructor params. batch_id is meaningless + to Pinecone; collection_name is meaningless to Postgres.""" + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + + with patch.object(m, "DocsiteSearch") as mock_pg_cls: + mock_pg_cls.return_value.search.return_value = [] + m.main({"query": "q", "backend": "postgres", "batch_id": 3, "collection_name": "ignored"}) + + mock_pg_cls.assert_called_once_with(batch_id=3) + + +def test_main_routes_collection_name_to_pinecone_only(monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + + with patch.object(m, "LegacyPineconeDocsiteSearch") as mock_legacy_cls: + mock_legacy_cls.return_value.search.return_value = [] + m.main({"query": "q", "backend": "pinecone", "collection_name": "docsite-202501010000", "batch_id": 9}) + + mock_legacy_cls.assert_called_once_with(collection_name="docsite-202501010000") From cbf3ae952498a29ddc3462cfbb3dbb64dfb3f681 Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Tue, 28 Jul 2026 15:36:18 +0800 Subject: [PATCH 17/48] feat: restore Pinecone write path behind embed_docsite target param --- poetry.lock | 122 +++++++++++- pyproject.toml | 1 + services/embed_docsite/README.md | 22 ++- services/embed_docsite/embed_docsite.py | 81 ++++++-- .../embed_docsite/pinecone_legacy_indexer.py | 182 ++++++++++++++++++ services/embed_docsite/tests/unit/conftest.py | 15 ++ .../tests/unit/test_embed_docsite.py | 47 ++++- 7 files changed, 449 insertions(+), 21 deletions(-) create mode 100644 services/embed_docsite/pinecone_legacy_indexer.py create mode 100644 services/embed_docsite/tests/unit/conftest.py diff --git a/poetry.lock b/poetry.lock index 0be6229a..d6f43638 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1724,6 +1724,102 @@ files = [ {file = "packaging-24.2.tar.gz", hash = "sha256:c228a6dc5e932d346bc5739379109d49e8853dd8223571c7c5b55260edc0b97f"}, ] +[[package]] +name = "pandas" +version = "2.3.3" +description = "Powerful data structures for data analysis, time series, and statistics" +optional = false +python-versions = ">=3.9" +groups = ["main"] +files = [ + {file = "pandas-2.3.3-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:376c6446ae31770764215a6c937f72d917f214b43560603cd60da6408f183b6c"}, + {file = "pandas-2.3.3-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:e19d192383eab2f4ceb30b412b22ea30690c9e618f78870357ae1d682912015a"}, + {file = "pandas-2.3.3-cp310-cp310-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5caf26f64126b6c7aec964f74266f435afef1c1b13da3b0636c7518a1fa3e2b1"}, + {file = "pandas-2.3.3-cp310-cp310-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:dd7478f1463441ae4ca7308a70e90b33470fa593429f9d4c578dd00d1fa78838"}, + {file = "pandas-2.3.3-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:4793891684806ae50d1288c9bae9330293ab4e083ccd1c5e383c34549c6e4250"}, + {file = "pandas-2.3.3-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:28083c648d9a99a5dd035ec125d42439c6c1c525098c58af0fc38dd1a7a1b3d4"}, + {file = "pandas-2.3.3-cp310-cp310-win_amd64.whl", hash = "sha256:503cf027cf9940d2ceaa1a93cfb5f8c8c7e6e90720a2850378f0b3f3b1e06826"}, + {file = "pandas-2.3.3-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:602b8615ebcc4a0c1751e71840428ddebeb142ec02c786e8ad6b1ce3c8dec523"}, + {file = "pandas-2.3.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:8fe25fc7b623b0ef6b5009149627e34d2a4657e880948ec3c840e9402e5c1b45"}, + {file = "pandas-2.3.3-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b468d3dad6ff947df92dcb32ede5b7bd41a9b3cceef0a30ed925f6d01fb8fa66"}, + {file = "pandas-2.3.3-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b98560e98cb334799c0b07ca7967ac361a47326e9b4e5a7dfb5ab2b1c9d35a1b"}, + {file = "pandas-2.3.3-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:1d37b5848ba49824e5c30bedb9c830ab9b7751fd049bc7914533e01c65f79791"}, + {file = "pandas-2.3.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:db4301b2d1f926ae677a751eb2bd0e8c5f5319c9cb3f88b0becbbb0b07b34151"}, + {file = "pandas-2.3.3-cp311-cp311-win_amd64.whl", hash = "sha256:f086f6fe114e19d92014a1966f43a3e62285109afe874f067f5abbdcbb10e59c"}, + {file = "pandas-2.3.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:6d21f6d74eb1725c2efaa71a2bfc661a0689579b58e9c0ca58a739ff0b002b53"}, + {file = "pandas-2.3.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:3fd2f887589c7aa868e02632612ba39acb0b8948faf5cc58f0850e165bd46f35"}, + {file = "pandas-2.3.3-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ecaf1e12bdc03c86ad4a7ea848d66c685cb6851d807a26aa245ca3d2017a1908"}, + {file = "pandas-2.3.3-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b3d11d2fda7eb164ef27ffc14b4fcab16a80e1ce67e9f57e19ec0afaf715ba89"}, + {file = "pandas-2.3.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:a68e15f780eddf2b07d242e17a04aa187a7ee12b40b930bfdd78070556550e98"}, + {file = "pandas-2.3.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:371a4ab48e950033bcf52b6527eccb564f52dc826c02afd9a1bc0ab731bba084"}, + {file = "pandas-2.3.3-cp312-cp312-win_amd64.whl", hash = "sha256:a16dcec078a01eeef8ee61bf64074b4e524a2a3f4b3be9326420cabe59c4778b"}, + {file = "pandas-2.3.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:56851a737e3470de7fa88e6131f41281ed440d29a9268dcbf0002da5ac366713"}, + {file = "pandas-2.3.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:bdcd9d1167f4885211e401b3036c0c8d9e274eee67ea8d0758a256d60704cfe8"}, + {file = "pandas-2.3.3-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e32e7cc9af0f1cc15548288a51a3b681cc2a219faa838e995f7dc53dbab1062d"}, + {file = "pandas-2.3.3-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:318d77e0e42a628c04dc56bcef4b40de67918f7041c2b061af1da41dcff670ac"}, + {file = "pandas-2.3.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:4e0a175408804d566144e170d0476b15d78458795bb18f1304fb94160cabf40c"}, + {file = "pandas-2.3.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:93c2d9ab0fc11822b5eece72ec9587e172f63cff87c00b062f6e37448ced4493"}, + {file = "pandas-2.3.3-cp313-cp313-win_amd64.whl", hash = "sha256:f8bfc0e12dc78f777f323f55c58649591b2cd0c43534e8355c51d3fede5f4dee"}, + {file = "pandas-2.3.3-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:75ea25f9529fdec2d2e93a42c523962261e567d250b0013b16210e1d40d7c2e5"}, + {file = "pandas-2.3.3-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:74ecdf1d301e812db96a465a525952f4dde225fdb6d8e5a521d47e1f42041e21"}, + {file = "pandas-2.3.3-cp313-cp313t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6435cb949cb34ec11cc9860246ccb2fdc9ecd742c12d3304989017d53f039a78"}, + {file = "pandas-2.3.3-cp313-cp313t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:900f47d8f20860de523a1ac881c4c36d65efcb2eb850e6948140fa781736e110"}, + {file = "pandas-2.3.3-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:a45c765238e2ed7d7c608fc5bc4a6f88b642f2f01e70c0c23d2224dd21829d86"}, + {file = "pandas-2.3.3-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:c4fc4c21971a1a9f4bdb4c73978c7f7256caa3e62b323f70d6cb80db583350bc"}, + {file = "pandas-2.3.3-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:ee15f284898e7b246df8087fc82b87b01686f98ee67d85a17b7ab44143a3a9a0"}, + {file = "pandas-2.3.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:1611aedd912e1ff81ff41c745822980c49ce4a7907537be8692c8dbc31924593"}, + {file = "pandas-2.3.3-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6d2cefc361461662ac48810cb14365a365ce864afe85ef1f447ff5a1e99ea81c"}, + {file = "pandas-2.3.3-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ee67acbbf05014ea6c763beb097e03cd629961c8a632075eeb34247120abcb4b"}, + {file = "pandas-2.3.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:c46467899aaa4da076d5abc11084634e2d197e9460643dd455ac3db5856b24d6"}, + {file = "pandas-2.3.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:6253c72c6a1d990a410bc7de641d34053364ef8bcd3126f7e7450125887dffe3"}, + {file = "pandas-2.3.3-cp314-cp314-win_amd64.whl", hash = "sha256:1b07204a219b3b7350abaae088f451860223a52cfb8a6c53358e7948735158e5"}, + {file = "pandas-2.3.3-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:2462b1a365b6109d275250baaae7b760fd25c726aaca0054649286bcfbb3e8ec"}, + {file = "pandas-2.3.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:0242fe9a49aa8b4d78a4fa03acb397a58833ef6199e9aa40a95f027bb3a1b6e7"}, + {file = "pandas-2.3.3-cp314-cp314t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a21d830e78df0a515db2b3d2f5570610f5e6bd2e27749770e8bb7b524b89b450"}, + {file = "pandas-2.3.3-cp314-cp314t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2e3ebdb170b5ef78f19bfb71b0dc5dc58775032361fa188e814959b74d726dd5"}, + {file = "pandas-2.3.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:d051c0e065b94b7a3cea50eb1ec32e912cd96dba41647eb24104b6c6c14c5788"}, + {file = "pandas-2.3.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:3869faf4bd07b3b66a9f462417d0ca3a9df29a9f6abd5d0d0dbab15dac7abe87"}, + {file = "pandas-2.3.3-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:c503ba5216814e295f40711470446bc3fd00f0faea8a086cbc688808e26f92a2"}, + {file = "pandas-2.3.3-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:a637c5cdfa04b6d6e2ecedcb81fc52ffb0fd78ce2ebccc9ea964df9f658de8c8"}, + {file = "pandas-2.3.3-cp39-cp39-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:854d00d556406bffe66a4c0802f334c9ad5a96b4f1f868adf036a21b11ef13ff"}, + {file = "pandas-2.3.3-cp39-cp39-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bf1f8a81d04ca90e32a0aceb819d34dbd378a98bf923b6398b9a3ec0bf44de29"}, + {file = "pandas-2.3.3-cp39-cp39-musllinux_1_2_aarch64.whl", hash = "sha256:23ebd657a4d38268c7dfbdf089fbc31ea709d82e4923c5ffd4fbd5747133ce73"}, + {file = "pandas-2.3.3-cp39-cp39-musllinux_1_2_x86_64.whl", hash = "sha256:5554c929ccc317d41a5e3d1234f3be588248e61f08a74dd17c9eabb535777dc9"}, + {file = "pandas-2.3.3-cp39-cp39-win_amd64.whl", hash = "sha256:d3e28b3e83862ccf4d85ff19cf8c20b2ae7e503881711ff2d534dc8f761131aa"}, + {file = "pandas-2.3.3.tar.gz", hash = "sha256:e05e1af93b977f7eafa636d043f9f94c7ee3ac81af99c13508215942e64c993b"}, +] + +[package.dependencies] +numpy = {version = ">=1.23.2", markers = "python_version == \"3.11\""} +python-dateutil = ">=2.8.2" +pytz = ">=2020.1" +tzdata = ">=2022.7" + +[package.extras] +all = ["PyQt5 (>=5.15.9)", "SQLAlchemy (>=2.0.0)", "adbc-driver-postgresql (>=0.8.0)", "adbc-driver-sqlite (>=0.8.0)", "beautifulsoup4 (>=4.11.2)", "bottleneck (>=1.3.6)", "dataframe-api-compat (>=0.1.7)", "fastparquet (>=2022.12.0)", "fsspec (>=2022.11.0)", "gcsfs (>=2022.11.0)", "html5lib (>=1.1)", "hypothesis (>=6.46.1)", "jinja2 (>=3.1.2)", "lxml (>=4.9.2)", "matplotlib (>=3.6.3)", "numba (>=0.56.4)", "numexpr (>=2.8.4)", "odfpy (>=1.4.1)", "openpyxl (>=3.1.0)", "pandas-gbq (>=0.19.0)", "psycopg2 (>=2.9.6)", "pyarrow (>=10.0.1)", "pymysql (>=1.0.2)", "pyreadstat (>=1.2.0)", "pytest (>=7.3.2)", "pytest-xdist (>=2.2.0)", "python-calamine (>=0.1.7)", "pyxlsb (>=1.0.10)", "qtpy (>=2.3.0)", "s3fs (>=2022.11.0)", "scipy (>=1.10.0)", "tables (>=3.8.0)", "tabulate (>=0.9.0)", "xarray (>=2022.12.0)", "xlrd (>=2.0.1)", "xlsxwriter (>=3.0.5)", "zstandard (>=0.19.0)"] +aws = ["s3fs (>=2022.11.0)"] +clipboard = ["PyQt5 (>=5.15.9)", "qtpy (>=2.3.0)"] +compression = ["zstandard (>=0.19.0)"] +computation = ["scipy (>=1.10.0)", "xarray (>=2022.12.0)"] +consortium-standard = ["dataframe-api-compat (>=0.1.7)"] +excel = ["odfpy (>=1.4.1)", "openpyxl (>=3.1.0)", "python-calamine (>=0.1.7)", "pyxlsb (>=1.0.10)", "xlrd (>=2.0.1)", "xlsxwriter (>=3.0.5)"] +feather = ["pyarrow (>=10.0.1)"] +fss = ["fsspec (>=2022.11.0)"] +gcp = ["gcsfs (>=2022.11.0)", "pandas-gbq (>=0.19.0)"] +hdf5 = ["tables (>=3.8.0)"] +html = ["beautifulsoup4 (>=4.11.2)", "html5lib (>=1.1)", "lxml (>=4.9.2)"] +mysql = ["SQLAlchemy (>=2.0.0)", "pymysql (>=1.0.2)"] +output-formatting = ["jinja2 (>=3.1.2)", "tabulate (>=0.9.0)"] +parquet = ["pyarrow (>=10.0.1)"] +performance = ["bottleneck (>=1.3.6)", "numba (>=0.56.4)", "numexpr (>=2.8.4)"] +plot = ["matplotlib (>=3.6.3)"] +postgresql = ["SQLAlchemy (>=2.0.0)", "adbc-driver-postgresql (>=0.8.0)", "psycopg2 (>=2.9.6)"] +pyarrow = ["pyarrow (>=10.0.1)"] +spss = ["pyreadstat (>=1.2.0)"] +sql-other = ["SQLAlchemy (>=2.0.0)", "adbc-driver-postgresql (>=0.8.0)", "adbc-driver-sqlite (>=0.8.0)"] +test = ["hypothesis (>=6.46.1)", "pytest (>=7.3.2)", "pytest-xdist (>=2.2.0)"] +xml = ["lxml (>=4.9.2)"] + [[package]] name = "pgvector" version = "0.3.6" @@ -2284,6 +2380,18 @@ files = [ [package.extras] cli = ["click (>=5.0)"] +[[package]] +name = "pytz" +version = "2026.3.post1" +description = "World timezone definitions, modern and historical" +optional = false +python-versions = "*" +groups = ["main"] +files = [ + {file = "pytz-2026.3.post1-py2.py3-none-any.whl", hash = "sha256:dd95840dd199baea12d9cc096a1d452caa6596a1c1e4b5f3dbd1541855d5e815"}, + {file = "pytz-2026.3.post1.tar.gz", hash = "sha256:2211d3fcf9a797d3405cac96ac7f61d80e6a644f72a3309607282fe8a2010c5d"}, +] + [[package]] name = "pyyaml" version = "6.0.3" @@ -2973,6 +3081,18 @@ files = [ [package.dependencies] typing-extensions = ">=4.12.0" +[[package]] +name = "tzdata" +version = "2026.3" +description = "Provider of IANA time zone data" +optional = false +python-versions = ">=2" +groups = ["main"] +files = [ + {file = "tzdata-2026.3-py2.py3-none-any.whl", hash = "sha256:dc096730c87af6cab1b171c9d532be840741ff5d459015e7f6947bd7d7e54931"}, + {file = "tzdata-2026.3.tar.gz", hash = "sha256:4a1518b8993086a7982523e071643f3c0e5f213e75b21318e78bcabfff9d1415"}, +] + [[package]] name = "urllib3" version = "2.7.0" @@ -3640,4 +3760,4 @@ cffi = ["cffi (>=1.17,<2.0) ; platform_python_implementation != \"PyPy\" and pyt [metadata] lock-version = "2.1" python-versions = "3.11.*" -content-hash = "d7763120d68c74066844df66ac246b4ab1d2f596437b158c5d20770baa98f7a1" +content-hash = "8e4463c369e05f92ff1c77750943b311f4af4699d23c487b032e7836166242ca" diff --git a/pyproject.toml b/pyproject.toml index 7f76b4ae..e3c82a6e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,6 +27,7 @@ langfuse = "^4.14.1" opentelemetry-instrumentation-anthropic = "^0.62.1" opentelemetry-instrumentation-threading = "0.65b0" pgvector = "^0.3.6" +pandas = "^2.2" [tool.poetry.group.dev] optional = false diff --git a/services/embed_docsite/README.md b/services/embed_docsite/README.md index 6129f303..b3c2530b 100644 --- a/services/embed_docsite/README.md +++ b/services/embed_docsite/README.md @@ -27,14 +27,26 @@ The service uses the DocsiteProcessor to download the documentation and chunk it 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 { - "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. + "target": "pinecone", // 'pinecone' | 'postgres'. Defaults to pinecone. Chooses the write destination. + "docs_to_upload": ["adaptor_docs", "general_docs", "adaptor_functions"], + "docs_to_ignore": ["job-examples.md", "release-notes.md"], + "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", + "max_total_collections": 3, + + // Postgres target only: + "keep_batches": 2 // Number of recent complete batches to retain when pruning } ``` diff --git a/services/embed_docsite/embed_docsite.py b/services/embed_docsite/embed_docsite.py index e8c6f972..95936ffe 100644 --- a/services/embed_docsite/embed_docsite.py +++ b/services/embed_docsite/embed_docsite.py @@ -8,20 +8,48 @@ register_vector_type, ) from embed_docsite.docsite_processor import DocsiteProcessor +from embed_docsite.pinecone_legacy_indexer import LegacyPineconeDocsiteIndexer from util import ApolloError, create_logger, get_db_connection logger = create_logger("embed_docsite") +VALID_TARGETS = ("pinecone", "postgres") + + +def _collect_documents(docs_to_upload, docs_to_ignore, chunk_target_length, chunk_min_length): + """Download and chunk every requested docs_type. Shared by both targets.""" + documents = [] + metadata_dict = {} + for docs_type in docs_to_upload: + processor = DocsiteProcessor( + docs_type=docs_type, + docs_to_ignore=docs_to_ignore, + target_length=chunk_target_length, + min_length=chunk_min_length, + ) + type_documents, type_metadata = processor.get_preprocessed_docs() + documents.extend(type_documents) + metadata_dict.update(type_metadata) + return documents, metadata_dict + def main(data: dict) -> dict: logger.info("Starting...") + target = data.get("target", "pinecone") docs_to_upload = data.get("docs_to_upload", ALL_DOCS_TYPES) docs_to_ignore = data.get("docs_to_ignore", ["job-examples.md", "release-notes.md"]) chunk_target_length = data.get("chunk_target_length", 1000) chunk_min_length = data.get("chunk_min_length", 700) keep_batches = data.get("keep_batches", 2) + if target not in VALID_TARGETS: + raise ApolloError( + 400, + f"Unknown target '{target}'. Expected 'pinecone' or 'postgres'", + type="BAD_REQUEST", + ) + load_dotenv(override=True) openai_api_key = data.get("OPENAI_API_KEY") or os.environ.get("OPENAI_API_KEY") @@ -30,6 +58,45 @@ def main(data: dict) -> dict: logger.error(msg) raise ApolloError(500, f"{msg}. Add to payload or environment", type="BAD_REQUEST") + documents, metadata_dict = _collect_documents( + docs_to_upload, docs_to_ignore, chunk_target_length, chunk_min_length, + ) + + if target == "pinecone": + return _upload_to_pinecone(data, documents, metadata_dict, docs_to_upload) + + return _upload_to_postgres( + documents, metadata_dict, docs_to_upload, chunk_target_length, chunk_min_length, keep_batches, + ) + + +def _upload_to_pinecone(data, documents, metadata_dict, docs_to_upload): + """Legacy write path. Deliberately opens no Postgres connection.""" + pinecone_api_key = data.get("PINECONE_API_KEY") or os.environ.get("PINECONE_API_KEY") + if not pinecone_api_key: + msg = "Missing API key: PINECONE_API_KEY" + logger.error(msg) + raise ApolloError(500, f"{msg}. Add to payload or environment", type="BAD_REQUEST") + + index_params = { + key: data[key] + for key in ("collection_name", "index_name", "max_total_collections") + if key in data + } + indexer = LegacyPineconeDocsiteIndexer(**index_params) + indexer.insert_documents(documents, metadata_dict) + + return { + "target": "pinecone", + "collection_name": indexer.collection_name, + "docs_types": docs_to_upload, + "chunk_count": len(documents), + } + + +def _upload_to_postgres( + documents, metadata_dict, docs_to_upload, chunk_target_length, chunk_min_length, keep_batches, +): indexer = DocsiteIndexer( chunk_target_length=chunk_target_length, chunk_min_length=chunk_min_length, @@ -42,19 +109,6 @@ def main(data: dict) -> dict: try: create_table_if_not_exists(conn) - documents = [] - metadata_dict = {} - for docs_type in docs_to_upload: - processor = DocsiteProcessor( - docs_type=docs_type, - docs_to_ignore=docs_to_ignore, - target_length=chunk_target_length, - min_length=chunk_min_length, - ) - type_documents, type_metadata = processor.get_preprocessed_docs() - documents.extend(type_documents) - metadata_dict.update(type_metadata) - batch_id = indexer.start_batch(conn, docs_to_upload) chunk_count = indexer.insert_documents(conn, batch_id, documents, metadata_dict) copied = indexer.copy_forward_missing_docs_types(conn, batch_id, docs_to_upload) @@ -63,6 +117,7 @@ def main(data: dict) -> dict: pruned = indexer.prune_old_batches(conn) return { + "target": "postgres", "batch_id": batch_id, "docs_types": docs_to_upload, "chunk_count": chunk_count, diff --git a/services/embed_docsite/pinecone_legacy_indexer.py b/services/embed_docsite/pinecone_legacy_indexer.py new file mode 100644 index 00000000..4094c79f --- /dev/null +++ b/services/embed_docsite/pinecone_legacy_indexer.py @@ -0,0 +1,182 @@ +import os +import time +from datetime import datetime +import pandas as pd +from pinecone import Pinecone, ServerlessSpec +from langchain_pinecone import PineconeVectorStore +from langchain_openai import OpenAIEmbeddings +from langchain_community.document_loaders import DataFrameLoader +from util import create_logger, ApolloError + +logger = create_logger("LegacyPineconeDocsiteIndexer") + +class LegacyPineconeDocsiteIndexer: + """ + Legacy Pinecone-backed docsite indexer, preserved as the write-side rollback + path for the Postgres migration. Selected via embed_docsite's `target` + payload param. Deleted alongside pinecone_legacy_search.py in the cleanup PR. + + Initialize vectorstore and insert new documents. Create a new index if needed. + + :param collection_name: Vectorstore collection name (namespace) to store documents + :param index_name: Vectorstore index name (default: docsite) + :param embeddings: LangChain embedding type (default: OpenAIEmbeddings()) + :param dimension: Embedding dimension (default: 1536 for OpenAI Embeddings) + :param max_total_collections: Max total collections in index. Delete old collections by date if exceeded after a new upload (default: 50) + """ + def __init__(self, collection_name=None, index_name="docsite", embeddings=OpenAIEmbeddings(), dimension=1536, max_total_collections=50): + self.collection_name = collection_name if collection_name is not None else f"docsite-{datetime.now().strftime('%Y%m%d%H%M')}" + self.index_name = index_name + self.embeddings = embeddings + self.dimension = dimension + self.max_total_collections = max_total_collections + self.pc = Pinecone(api_key=os.environ.get("PINECONE_API_KEY")) + + if not self.index_exists(): + self.create_index() + + self.index = self.pc.Index(self.index_name) + self.vectorstore = PineconeVectorStore(index_name=index_name, namespace=self.collection_name, embedding=embeddings) + + def insert_documents(self, inputs, metadata_dict): + """ + Create the index if it does not exist and insert the input documents. + + :param inputs: Dictionary containing name, docs_type, and doc_chunk + :param metadata_dict: Metadata dict with document titles as keys (from DocsiteProcessor) + :return: Initialized indices + """ + + # Get vector count before insertion for verification + try: + stats = self.index.describe_index_stats() + vectors_before = stats.namespaces.get(self.collection_name, {}).get("vector_count", 0) + logger.info(f"Current vector count in namespace '{self.collection_name}': {vectors_before}") + except Exception as e: + logger.warning(f"Could not get vector count before insertion: {str(e)}") + vectors_before = 0 + + df = self.preprocess_metadata(inputs=inputs, metadata_dict=metadata_dict) + logger.info(f"Input metadata preprocessed") + loader = DataFrameLoader(df, page_content_column="text") + docs = loader.load() + logger.info(f"Inputs processed into LangChain docs") + logger.info(f"Uploading {len(docs)} documents to index...") + + idx = self.vectorstore.add_documents( + documents=docs + ) + sleep_time = 10 + max_wait_time = 150 + elapsed_time = 0 + logger.info(f"Waiting up to {max_wait_time}s to verify upload count") + + while elapsed_time < max_wait_time: + time.sleep(sleep_time) + elapsed_time += sleep_time + + # Verify the upload by checking the vector count + try: + stats = self.index.describe_index_stats() + vectors_after = stats.namespaces.get(self.collection_name, {}).get("vector_count", 0) + logger.info(f"Vector count after {elapsed_time}s: {vectors_after}") + + if vectors_after >= vectors_before + len(docs): + logger.info(f"Successfully added {vectors_after - vectors_before} vectors to namespace '{self.collection_name}'") + break + else: + logger.warning(f"No new vectors were added to namespace '{self.collection_name}' after {sleep_time}s") + except Exception as e: + logger.warning(f"Could not verify vector insertion: {str(e)}") + + if vectors_after <= vectors_before: + logger.warning(f"Could not verify full dataset upload to namespace '{self.collection_name}' after {max_wait_time}s") + + self.delete_old_collections(self.max_total_collections) + + return idx + + def delete_collection(self): + """ + Deletes the entire collection (namespace) and all its contents. + This operation cannot be undone and removes both the collection structure and all vectors/documents within it. + """ + self.index.delete(delete_all=True, namespace=self.collection_name) + + def delete_old_collections(self, max_total_collections): + """Retrieve docsite uploads by collection name from Pinecone and delete them if there are more than max_total_collections.""" + + logger.info(f"Fetching outdated docsite collections") + pc = Pinecone(api_key=os.environ.get("PINECONE_API_KEY")) + index = pc.Index("docsite") + index_stats = index.describe_index_stats() + namespaces = index_stats.get('namespaces', {}).keys() + valid_namespaces = sorted( + (ns for ns in namespaces if ns.startswith("docsite-") and ns[8:].isdigit() and len(ns) == 16), + reverse=False + ) + if len(valid_namespaces) > max_total_collections: + logger.info(f"Deleting outdated docsite collections") + for old_collection in valid_namespaces[:max_total_collections]: + self.index.delete(delete_all=True, namespace=old_collection) + logger.info(f"Deleted collection {old_collection}") + + if not valid_namespaces: + logger.info(f"No valid namespaces found in the index when deleting old collections.") + + def create_index(self): + """Creates a new Pinecone index if it does not exist.""" + + if not self.index_exists(): + self.pc.create_index( + name=self.index_name, + dimension=self.dimension, + metric="cosine", + spec=ServerlessSpec(cloud="aws", region="us-east-1") + ) + while not self.pc.describe_index(self.index_name).status["ready"]: + time.sleep(1) + + def index_exists(self): + """Check if the index exists in Pinecone.""" + existing_indexes = [index_info["name"] for index_info in self.pc.list_indexes()] + + return self.index_name in existing_indexes + + def preprocess_metadata(self, inputs, page_content_column="text", add_chunk_as_metadata=False, metadata_cols=None, metadata_dict=None): + """ + Create a DataFrame for indexing from input documents and metadata. + + :param inputs: Dictionary containing name, docs_type, and doc_chunk + :param page_content_column: Name of the field which will be embedded (default: text) + :param add_chunk_as_metadata: Copy the text to embed as a separate metadata field (default: False) + :param metadata_cols: Optional list of metadata columns to include (default: None) + :param metadata_dict: Dictionary mapping names to metadata dictionaries (default: None) + :return: pandas.DataFrame with text and metadata columns + """ + + # Create DataFrame from the inputs (doc_chunk, name, docs_type) + df = pd.DataFrame(inputs) + + # Rename some columns for metadata upload + df = df.rename(columns={"doc_chunk": page_content_column, "name": "doc_title"}) + + df["doc_title"] = df["doc_title"].str.replace(".md$", "", regex=True) + + # Optionally add chunk to metadata for keyword searching + if add_chunk_as_metadata: + df["embedding_text"] = df[page_content_column] + + # Add further metadata columns if specified + if metadata_cols: + for col in metadata_cols: + df[col] = metadata_dict.get(inputs["name"], {}).get(col) + + return df + + + + + + + \ No newline at end of file diff --git a/services/embed_docsite/tests/unit/conftest.py b/services/embed_docsite/tests/unit/conftest.py new file mode 100644 index 00000000..bb48d19b --- /dev/null +++ b/services/embed_docsite/tests/unit/conftest.py @@ -0,0 +1,15 @@ +"""Test config for embed_docsite unit tests. + +Importing `embed_docsite.pinecone_legacy_indexer` pulls in +`LegacyPineconeDocsiteIndexer`, whose module-level `OpenAIEmbeddings()` +default arg validates credentials at construction (openai 2.x / langchain-openai 1.x). +A key must therefore exist at import time. + +Dummy placeholders only: unit tests mock every real network call, so no real +key is ever used. `setdefault` means a real key (from services/.env) wins. +""" + +import os + +os.environ.setdefault("OPENAI_API_KEY", "sk-test-dummy") +os.environ.setdefault("PINECONE_API_KEY", "pc-test-dummy") diff --git a/services/embed_docsite/tests/unit/test_embed_docsite.py b/services/embed_docsite/tests/unit/test_embed_docsite.py index 2f2a9d28..05c16824 100644 --- a/services/embed_docsite/tests/unit/test_embed_docsite.py +++ b/services/embed_docsite/tests/unit/test_embed_docsite.py @@ -23,7 +23,7 @@ def test_main_orchestrates_full_batch_lifecycle_and_returns_summary(): patch.object(m, "DocsiteProcessor", return_value=fake_processor), \ patch.object(m, "DocsiteIndexer", return_value=fake_indexer), \ patch.dict("os.environ", {"OPENAI_API_KEY": "sk-test"}): - result = m.main({"docs_to_upload": ["general_docs"]}) + result = m.main({"docs_to_upload": ["general_docs"], "target": "postgres"}) fake_indexer.start_batch.assert_called_once_with(fake_conn, ["general_docs"]) fake_indexer.insert_documents.assert_called_once() @@ -34,6 +34,7 @@ def test_main_orchestrates_full_batch_lifecycle_and_returns_summary(): fake_conn.close.assert_called_once() assert result == { + "target": "postgres", "batch_id": 7, "docs_types": ["general_docs"], "chunk_count": 10, @@ -71,7 +72,49 @@ def test_main_defaults_docs_to_upload_to_all_types(): patch.object(m, "DocsiteProcessor", return_value=fake_processor) as mock_processor_cls, \ patch.object(m, "DocsiteIndexer", return_value=fake_indexer), \ patch.dict("os.environ", {"OPENAI_API_KEY": "sk-test"}): - m.main({}) + m.main({"target": "postgres"}) called_docs_types = [call.kwargs["docs_type"] for call in mock_processor_cls.call_args_list] assert called_docs_types == m.ALL_DOCS_TYPES + + +def test_main_defaults_to_pinecone_target(): + """Default must match main's behavior: write to Pinecone, never open Postgres.""" + fake_indexer = MagicMock() + fake_processor = MagicMock() + fake_processor.get_preprocessed_docs.return_value = ([], {}) + + with patch.object(m, "get_db_connection") as mock_get_conn, \ + patch.object(m, "DocsiteProcessor", return_value=fake_processor), \ + patch.object(m, "LegacyPineconeDocsiteIndexer", return_value=fake_indexer) as mock_legacy_cls, \ + patch.dict("os.environ", {"OPENAI_API_KEY": "sk-test", "PINECONE_API_KEY": "pc-test"}): + m.main({"docs_to_upload": ["general_docs"]}) + + mock_legacy_cls.assert_called_once() + mock_get_conn.assert_not_called() + + +def test_main_pinecone_target_does_not_require_postgres_url(): + """With target=pinecone the service must not touch Postgres at all, so + POSTGRES_URL need not be set — matching main's dependency surface.""" + fake_processor = MagicMock() + fake_processor.get_preprocessed_docs.return_value = ([], {}) + + with patch.object(m, "get_db_connection") as mock_get_conn, \ + patch.object(m, "DocsiteProcessor", return_value=fake_processor), \ + patch.object(m, "LegacyPineconeDocsiteIndexer"), \ + patch.dict("os.environ", {"OPENAI_API_KEY": "sk-test", "PINECONE_API_KEY": "pc-test"}, clear=True): + m.main({}) + + mock_get_conn.assert_not_called() + + +def test_main_rejects_unknown_target(): + from util import ApolloError + import pytest + + with patch.dict("os.environ", {"OPENAI_API_KEY": "sk-test"}): + with pytest.raises(ApolloError) as exc: + m.main({"target": "elasticsearch"}) + + assert exc.value.code == 400 From ed07783d12d4473cc72ae1ef445d174ba53ce731 Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Tue, 28 Jul 2026 15:49:16 +0800 Subject: [PATCH 18/48] fix: add type annotations to embed_docsite's new helper functions --- services/embed_docsite/embed_docsite.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/services/embed_docsite/embed_docsite.py b/services/embed_docsite/embed_docsite.py index 95936ffe..f6f6a204 100644 --- a/services/embed_docsite/embed_docsite.py +++ b/services/embed_docsite/embed_docsite.py @@ -16,7 +16,9 @@ VALID_TARGETS = ("pinecone", "postgres") -def _collect_documents(docs_to_upload, docs_to_ignore, chunk_target_length, chunk_min_length): +def _collect_documents( + docs_to_upload: list, docs_to_ignore: list, chunk_target_length: int, chunk_min_length: int, +) -> tuple[list, dict]: """Download and chunk every requested docs_type. Shared by both targets.""" documents = [] metadata_dict = {} @@ -70,7 +72,7 @@ def main(data: dict) -> dict: ) -def _upload_to_pinecone(data, documents, metadata_dict, docs_to_upload): +def _upload_to_pinecone(data: dict, documents: list, metadata_dict: dict, docs_to_upload: list) -> dict: """Legacy write path. Deliberately opens no Postgres connection.""" pinecone_api_key = data.get("PINECONE_API_KEY") or os.environ.get("PINECONE_API_KEY") if not pinecone_api_key: @@ -95,8 +97,13 @@ def _upload_to_pinecone(data, documents, metadata_dict, docs_to_upload): def _upload_to_postgres( - documents, metadata_dict, docs_to_upload, chunk_target_length, chunk_min_length, keep_batches, -): + documents: list, + metadata_dict: dict, + docs_to_upload: list, + chunk_target_length: int, + chunk_min_length: int, + keep_batches: int, +) -> dict: indexer = DocsiteIndexer( chunk_target_length=chunk_target_length, chunk_min_length=chunk_min_length, From 131830e6d445a37830a5fde70dc850e22b57ef40 Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Tue, 28 Jul 2026 16:03:43 +0800 Subject: [PATCH 19/48] feat: add versioned migration runner for the docsite schema --- services/db_migrations.py | 66 +++++++++++ services/embed_docsite/docsite_indexer.py | 41 ------- services/embed_docsite/embed_docsite.py | 3 - .../tests/unit/test_db_migrations.py | 109 ++++++++++++++++++ .../tests/unit/test_docsite_indexer.py | 11 -- .../tests/unit/test_embed_docsite.py | 2 - .../0001_docsite_batches_and_chunks.sql} | 6 +- services/util.py | 13 ++- 8 files changed, 189 insertions(+), 62 deletions(-) create mode 100644 services/db_migrations.py create mode 100644 services/embed_docsite/tests/unit/test_db_migrations.py rename services/{embed_docsite/schema.sql => migrations/0001_docsite_batches_and_chunks.sql} (87%) diff --git a/services/db_migrations.py b/services/db_migrations.py new file mode 100644 index 00000000..26f573ac --- /dev/null +++ b/services/db_migrations.py @@ -0,0 +1,66 @@ +"""Versioned schema migrations for the Python-owned docs database (POSTGRES_URL). + +Mirrors platform/src/db/migrate.ts, which owns the TypeScript-side auth database. +The two runners are deliberately kept separate but symmetrical: .sql files applied +in lexical order, applied filenames recorded so re-runs are a no-op, and an +advisory lock so concurrent starters queue rather than collide. + +The tracking table (_migrations_docs) and lock key (8314_2026) are both distinct +from the TypeScript runner's, because APOLLO_CLIENTS_DB_URL falls back to +POSTGRES_URL in local development 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() -> list: + """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) -> int: + """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) diff --git a/services/embed_docsite/docsite_indexer.py b/services/embed_docsite/docsite_indexer.py index e4de7c44..d47a02b2 100644 --- a/services/embed_docsite/docsite_indexer.py +++ b/services/embed_docsite/docsite_indexer.py @@ -8,47 +8,6 @@ ALL_DOCS_TYPES = ["adaptor_docs", "general_docs", "adaptor_functions"] -CREATE_TABLES_SQL = """ -CREATE EXTENSION IF NOT EXISTS vector; - -CREATE TABLE IF NOT EXISTS docsite_batches ( - id BIGSERIAL PRIMARY KEY, - status VARCHAR(20) NOT NULL DEFAULT 'building' - CHECK (status IN ('building', 'complete', 'failed')), - docs_types TEXT[] NOT NULL, - chunk_target_length INT NOT NULL, - chunk_min_length INT NOT NULL, - embedding_model VARCHAR(100) NOT NULL, - chunk_count INT, - started_at TIMESTAMPTZ NOT NULL DEFAULT now(), - completed_at TIMESTAMPTZ -); - -CREATE TABLE IF NOT EXISTS docsite_chunks ( - id BIGSERIAL PRIMARY KEY, - batch_id BIGINT NOT NULL REFERENCES docsite_batches(id) ON DELETE CASCADE, - doc_title VARCHAR(500) NOT NULL, - docs_type VARCHAR(50) NOT NULL, - chunk_index INT NOT NULL, - text TEXT NOT NULL, - embedding vector(1536) NOT NULL, - text_search tsvector GENERATED ALWAYS AS (to_tsvector('english', text)) STORED, - created_at TIMESTAMPTZ NOT NULL DEFAULT now() -); - -CREATE INDEX IF NOT EXISTS idx_docsite_chunks_batch ON docsite_chunks(batch_id); -CREATE INDEX IF NOT EXISTS idx_docsite_chunks_doc_title ON docsite_chunks(batch_id, doc_title); -CREATE INDEX IF NOT EXISTS idx_docsite_chunks_docs_type ON docsite_chunks(batch_id, docs_type); -CREATE INDEX IF NOT EXISTS idx_docsite_chunks_fts ON docsite_chunks USING gin(text_search); -""" - - -def create_table_if_not_exists(conn): - """Create the docsite_batches/docsite_chunks tables and pgvector extension if missing.""" - with conn.cursor() as cur: - cur.execute(CREATE_TABLES_SQL) - conn.commit() - def register_vector_type(conn): """Register the pgvector adapter on this connection so Python lists convert to the `vector` type.""" diff --git a/services/embed_docsite/embed_docsite.py b/services/embed_docsite/embed_docsite.py index f6f6a204..c4da2194 100644 --- a/services/embed_docsite/embed_docsite.py +++ b/services/embed_docsite/embed_docsite.py @@ -4,7 +4,6 @@ from embed_docsite.docsite_indexer import ( ALL_DOCS_TYPES, DocsiteIndexer, - create_table_if_not_exists, register_vector_type, ) from embed_docsite.docsite_processor import DocsiteProcessor @@ -114,8 +113,6 @@ def _upload_to_postgres( register_vector_type(conn) try: - create_table_if_not_exists(conn) - batch_id = indexer.start_batch(conn, docs_to_upload) chunk_count = indexer.insert_documents(conn, batch_id, documents, metadata_dict) copied = indexer.copy_forward_missing_docs_types(conn, batch_id, docs_to_upload) diff --git a/services/embed_docsite/tests/unit/test_db_migrations.py b/services/embed_docsite/tests/unit/test_db_migrations.py new file mode 100644 index 00000000..06c043cf --- /dev/null +++ b/services/embed_docsite/tests/unit/test_db_migrations.py @@ -0,0 +1,109 @@ +"""Unit tests for the Python migration runner. + +Mirrors the TypeScript runner in platform/src/db/migrate.ts: lexical ordering, +already-applied files skipped, advisory lock taken. The connection and cursor +are MagicMocks — the repo-root conftest blocks real psycopg2.connect in unit +tests anyway. +""" + +from unittest.mock import MagicMock, patch + +import db_migrations as m + + +def make_conn(): + conn = MagicMock() + cur = MagicMock() + conn.cursor.return_value.__enter__.return_value = cur + return conn, cur + + +def test_run_migrations_takes_advisory_lock_before_applying(): + conn, cur = make_conn() + cur.fetchall.return_value = [] + + with patch.object(m, "_migration_files", return_value=[]): + m.run_migrations(conn) + + first_sql = cur.execute.call_args_list[0][0][0] + assert "pg_advisory_xact_lock" in first_sql + assert "8314" in str(cur.execute.call_args_list[0]) + + +def test_run_migrations_creates_tracking_table(): + conn, cur = make_conn() + cur.fetchall.return_value = [] + + with patch.object(m, "_migration_files", return_value=[]): + m.run_migrations(conn) + + all_sql = " ".join(str(call) for call in cur.execute.call_args_list) + assert "_migrations_docs" in all_sql + + +def test_migration_files_returns_sql_files_in_lexical_order(tmp_path): + """The sort lives in _migration_files, so it must be tested against the real + filesystem — patching that function out (as the apply tests below do) would + bypass the very ordering being asserted. Files are created out of order and + a non-.sql file is included to prove it is filtered.""" + (tmp_path / "0002_second.sql").write_text("SELECT 2;", encoding="utf-8") + (tmp_path / "0010_tenth.sql").write_text("SELECT 10;", encoding="utf-8") + (tmp_path / "0001_first.sql").write_text("SELECT 1;", encoding="utf-8") + (tmp_path / "notes.md").write_text("not a migration", encoding="utf-8") + + with patch.object(m, "MIGRATIONS_DIR", tmp_path): + names = [p.name for p in m._migration_files()] + + assert names == ["0001_first.sql", "0002_second.sql", "0010_tenth.sql"] + + +def test_migration_files_returns_empty_when_dir_missing(tmp_path): + with patch.object(m, "MIGRATIONS_DIR", tmp_path / "does_not_exist"): + assert m._migration_files() == [] + + +def test_run_migrations_applies_pending_files_in_the_order_given(tmp_path): + conn, cur = make_conn() + cur.fetchall.return_value = [] + + first = tmp_path / "0001_first.sql" + first.write_text("SELECT 1;", encoding="utf-8") + second = tmp_path / "0002_second.sql" + second.write_text("SELECT 2;", encoding="utf-8") + + with patch.object(m, "_migration_files", return_value=[first, second]): + applied = m.run_migrations(conn) + + assert applied == 2 + executed = [str(call) for call in cur.execute.call_args_list] + first_idx = next(i for i, e in enumerate(executed) if "SELECT 1;" in e) + second_idx = next(i for i, e in enumerate(executed) if "SELECT 2;" in e) + assert first_idx < second_idx + + +def test_run_migrations_skips_already_applied_files(tmp_path): + conn, cur = make_conn() + cur.fetchall.return_value = [("0001_first.sql",)] + + first = tmp_path / "0001_first.sql" + first.write_text("SELECT 1;", encoding="utf-8") + + with patch.object(m, "_migration_files", return_value=[first]): + applied = m.run_migrations(conn) + + assert applied == 0 + executed = " ".join(str(call) for call in cur.execute.call_args_list) + assert "SELECT 1;" not in executed + + +def test_run_migrations_commits_once_at_the_end(tmp_path): + conn, cur = make_conn() + cur.fetchall.return_value = [] + + first = tmp_path / "0001_first.sql" + first.write_text("SELECT 1;", encoding="utf-8") + + with patch.object(m, "_migration_files", return_value=[first]): + m.run_migrations(conn) + + conn.commit.assert_called_once() diff --git a/services/embed_docsite/tests/unit/test_docsite_indexer.py b/services/embed_docsite/tests/unit/test_docsite_indexer.py index 95519816..066ded2f 100644 --- a/services/embed_docsite/tests/unit/test_docsite_indexer.py +++ b/services/embed_docsite/tests/unit/test_docsite_indexer.py @@ -19,17 +19,6 @@ def make_conn(): return conn, cur -def test_create_table_if_not_exists_executes_schema_sql(): - conn, cur = make_conn() - m.create_table_if_not_exists(conn) - - executed_sql = cur.execute.call_args[0][0] - assert "CREATE EXTENSION IF NOT EXISTS vector" in executed_sql - assert "CREATE TABLE IF NOT EXISTS docsite_batches" in executed_sql - assert "CREATE TABLE IF NOT EXISTS docsite_chunks" in executed_sql - conn.commit.assert_called_once() - - def test_register_vector_type_calls_pgvector_register(): conn = MagicMock() with patch.object(m, "register_vector") as mock_register: diff --git a/services/embed_docsite/tests/unit/test_embed_docsite.py b/services/embed_docsite/tests/unit/test_embed_docsite.py index 05c16824..3dc7ffe0 100644 --- a/services/embed_docsite/tests/unit/test_embed_docsite.py +++ b/services/embed_docsite/tests/unit/test_embed_docsite.py @@ -19,7 +19,6 @@ def test_main_orchestrates_full_batch_lifecycle_and_returns_summary(): with patch.object(m, "get_db_connection", return_value=fake_conn), \ patch.object(m, "register_vector_type"), \ - patch.object(m, "create_table_if_not_exists"), \ patch.object(m, "DocsiteProcessor", return_value=fake_processor), \ patch.object(m, "DocsiteIndexer", return_value=fake_indexer), \ patch.dict("os.environ", {"OPENAI_API_KEY": "sk-test"}): @@ -68,7 +67,6 @@ def test_main_defaults_docs_to_upload_to_all_types(): with patch.object(m, "get_db_connection", return_value=fake_conn), \ patch.object(m, "register_vector_type"), \ - patch.object(m, "create_table_if_not_exists"), \ patch.object(m, "DocsiteProcessor", return_value=fake_processor) as mock_processor_cls, \ patch.object(m, "DocsiteIndexer", return_value=fake_indexer), \ patch.dict("os.environ", {"OPENAI_API_KEY": "sk-test"}): diff --git a/services/embed_docsite/schema.sql b/services/migrations/0001_docsite_batches_and_chunks.sql similarity index 87% rename from services/embed_docsite/schema.sql rename to services/migrations/0001_docsite_batches_and_chunks.sql index 11928bf9..0b8dfa53 100644 --- a/services/embed_docsite/schema.sql +++ b/services/migrations/0001_docsite_batches_and_chunks.sql @@ -1,7 +1,5 @@ --- services/embed_docsite/schema.sql --- Schema for docsite chunk storage (Postgres + pgvector). --- Note: This table is automatically created by the embed_docsite service. --- See README.md for example queries. +-- Docsite chunk storage (Postgres + pgvector). +-- Applied by services/db_migrations.py; recorded in _migrations_docs. CREATE EXTENSION IF NOT EXISTS vector; diff --git a/services/util.py b/services/util.py index 45d6d0f7..3ba1d0d5 100644 --- a/services/util.py +++ b/services/util.py @@ -112,6 +112,11 @@ def apollo(name: str, payload: dict) -> dict: def get_db_connection() -> "psycopg2.extensions.connection": """Get database connection from POSTGRES_URL environment variable. + Applies any pending schema migrations before handing the connection back. + bridge.ts spawns a fresh Python process per request, so there is no + long-lived process to cache "already migrated" in; the runner's tracking + table makes repeat calls a cheap no-op. + Returns: psycopg2.connection: Database connection @@ -121,7 +126,13 @@ def get_db_connection() -> "psycopg2.extensions.connection": db_url = os.environ.get("POSTGRES_URL") if not db_url: raise ApolloError(500, "Missing POSTGRES_URL environment variable", type="DATABASE_ERROR") - return psycopg2.connect(db_url) + + conn = psycopg2.connect(db_url) + + from db_migrations import run_migrations + run_migrations(conn) + + return conn def sum_usage(*usage_objects): From ab6b2233c4e193fa88f2e956ba3fb6112ae36d83 Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Tue, 28 Jul 2026 16:23:44 +0800 Subject: [PATCH 20/48] feat: report backend agreement in the offline docsite eval --- .../search_docsite/tests/eval/run_eval.py | 30 +++++++++++ .../tests/unit/test_run_eval.py | 50 ++++++++++++++++++- 2 files changed, 79 insertions(+), 1 deletion(-) diff --git a/services/search_docsite/tests/eval/run_eval.py b/services/search_docsite/tests/eval/run_eval.py index 23b341ea..a0a4e184 100644 --- a/services/search_docsite/tests/eval/run_eval.py +++ b/services/search_docsite/tests/eval/run_eval.py @@ -24,6 +24,30 @@ def compute_recall_at_k(retrieved_titles: list, expected_titles: list) -> bool: return bool(set(retrieved_titles) & set(expected_titles)) +def compute_agreement(report_a: dict, report_b: dict) -> dict: + """Doc-title overlap between two backends' result sets, per query and averaged. + + Needs no ground truth, so it gives a usable comparison signal before + golden_queries.yaml is curated. This is the measure the deleted shadow mode + computed on live traffic; it belongs here, offline, instead. + """ + per_query = [] + for a, b in zip(report_a["per_query"], report_b["per_query"]): + titles_a = {t for t in a["retrieved_titles"] if t is not None} + titles_b = {t for t in b["retrieved_titles"] if t is not None} + union = titles_a | titles_b + overlap = titles_a & titles_b + per_query.append({ + "query": a["query"], + "overlap": len(overlap), + "union": len(union), + "jaccard": (len(overlap) / len(union)) if union else 0.0, + }) + + mean_jaccard = (sum(q["jaccard"] for q in per_query) / len(per_query)) if per_query else 0.0 + return {"per_query": per_query, "mean_jaccard": mean_jaccard} + + def run_eval(golden_queries: list, backend_cls, strategy: str, top_k: int = 5) -> dict: """Run every golden query against one backend/strategy and return a report dict.""" backend = backend_cls() @@ -89,6 +113,12 @@ def main() -> None: f"(scored={pinecone_report['queries_scored']}, skipped={pinecone_report['queries_skipped']}) " f"p50={pinecone_report['p50_latency_s']:.3f}s p95={pinecone_report['p95_latency_s']:.3f}s") + agreement = compute_agreement(postgres_report, pinecone_report) + print(f"Backend agreement: mean doc-title Jaccard={agreement['mean_jaccard']:.3f} " + f"across {len(agreement['per_query'])} queries") + for q in agreement["per_query"]: + print(f" {q['jaccard']:.2f} overlap={q['overlap']}/{q['union']} {q['query']}") + if __name__ == "__main__": main() diff --git a/services/search_docsite/tests/unit/test_run_eval.py b/services/search_docsite/tests/unit/test_run_eval.py index e76800a8..83599cf6 100644 --- a/services/search_docsite/tests/unit/test_run_eval.py +++ b/services/search_docsite/tests/unit/test_run_eval.py @@ -3,7 +3,8 @@ from unittest.mock import MagicMock -from search_docsite.tests.eval.run_eval import compute_recall_at_k, run_eval +import pytest +from search_docsite.tests.eval.run_eval import compute_agreement, compute_recall_at_k, run_eval def test_compute_recall_at_k_true_when_any_expected_title_present(): @@ -64,3 +65,50 @@ def test_run_eval_reports_latency_percentiles(): assert "p50_latency_s" in report assert "p95_latency_s" in report assert report["p50_latency_s"] >= 0 + + +def test_compute_agreement_reports_perfect_overlap(): + report_a = {"per_query": [{"query": "q", "retrieved_titles": ["A", "B"]}]} + report_b = {"per_query": [{"query": "q", "retrieved_titles": ["B", "A"]}]} + + agreement = compute_agreement(report_a, report_b) + + assert agreement["per_query"][0]["overlap"] == 2 + assert agreement["per_query"][0]["jaccard"] == 1.0 + assert agreement["mean_jaccard"] == 1.0 + + +def test_compute_agreement_reports_partial_overlap(): + report_a = {"per_query": [{"query": "q", "retrieved_titles": ["A", "B"]}]} + report_b = {"per_query": [{"query": "q", "retrieved_titles": ["B", "C"]}]} + + agreement = compute_agreement(report_a, report_b) + + assert agreement["per_query"][0]["overlap"] == 1 + assert agreement["per_query"][0]["union"] == 3 + assert agreement["per_query"][0]["jaccard"] == pytest.approx(1 / 3) + + +def test_compute_agreement_handles_both_backends_returning_nothing(): + report_a = {"per_query": [{"query": "q", "retrieved_titles": []}]} + report_b = {"per_query": [{"query": "q", "retrieved_titles": []}]} + + agreement = compute_agreement(report_a, report_b) + + assert agreement["per_query"][0]["jaccard"] == 0.0 + assert agreement["mean_jaccard"] == 0.0 + + +def test_compute_agreement_means_across_queries(): + report_a = {"per_query": [ + {"query": "q1", "retrieved_titles": ["A"]}, + {"query": "q2", "retrieved_titles": ["X"]}, + ]} + report_b = {"per_query": [ + {"query": "q1", "retrieved_titles": ["A"]}, + {"query": "q2", "retrieved_titles": ["Y"]}, + ]} + + agreement = compute_agreement(report_a, report_b) + + assert agreement["mean_jaccard"] == pytest.approx(0.5) From 378aeabf52209092f7fae3ecb86ca90f251c289e Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Tue, 28 Jul 2026 16:44:40 +0800 Subject: [PATCH 21/48] docs: correct stale LegacyPineconeDocsiteSearch docstring --- services/search_docsite/pinecone_legacy_search.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/services/search_docsite/pinecone_legacy_search.py b/services/search_docsite/pinecone_legacy_search.py index 5718e2aa..c131e72b 100644 --- a/services/search_docsite/pinecone_legacy_search.py +++ b/services/search_docsite/pinecone_legacy_search.py @@ -10,10 +10,15 @@ class LegacyPineconeDocsiteSearch: """ - Legacy Pinecone-backed docsite search, preserved for the Postgres-migration - shadow-mode comparison window and as a rollback path. Not used by any - mounted service directly — see services/job_chat/retrieve_docs.py's - DOCSITE_SEARCH_BACKEND/DOCSITE_SHADOW_POSTGRES flags. + Legacy Pinecone-backed docsite search, preserved from before the Postgres + migration as the rollback path. Since DOCSITE_SEARCH_BACKEND defaults to + "pinecone", this implementation is what actually serves search traffic + today whenever that flag is unset or explicitly set to "pinecone". It is + constructed directly by services/search_docsite/search_docsite.py's + main(), services/job_chat/retrieve_docs.py's search_docs(), and + services/tools/search_documentation/search_documentation.py's + _search_implementation(), each of which select it via the + DOCSITE_SEARCH_BACKEND flag. :param collection_name: Vectorstore collection name (namespace) to store documents :param index_name: Vectorstore index name (default: docsite) From 16437303b26a59234f2949372f9bd0f7399ca754 Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Tue, 28 Jul 2026 22:47:31 +0800 Subject: [PATCH 22/48] refactor: revert unrelated churn in docsite_processor --- services/embed_docsite/docsite_processor.py | 61 ++++++++++--------- .../tests/unit/test_docsite_processor.py | 2 +- 2 files changed, 32 insertions(+), 31 deletions(-) diff --git a/services/embed_docsite/docsite_processor.py b/services/embed_docsite/docsite_processor.py index 3660930e..c49f74a1 100644 --- a/services/embed_docsite/docsite_processor.py +++ b/services/embed_docsite/docsite_processor.py @@ -1,17 +1,15 @@ import json import os import re - import nltk +from embed_docsite.github_utils import get_docs +from util import create_logger try: nltk.data.find('tokenizers/punkt_tab') except LookupError: nltk.download('punkt_tab', quiet=True) -from embed_docsite.github_utils import get_docs -from util import create_logger - logger = create_logger("DocsiteProcessor") class DocsiteProcessor: @@ -48,38 +46,41 @@ def get_preprocessed_docs(self): return chunks, metadata_dict def _chunk_adaptor_docs(self, json_data): - """Extract and clean docs from adaptor data, and chunk according to self.target_length/self.min_length.""" + """Extract and clean docs from adaptor data, and chunk according to a target and minimum chunk sizes.""" output = [] metadata_dict = dict() - + for item in json_data: if isinstance(item, dict) and "docs" in item and "name" in item: if item["name"] in self.docs_to_ignore: continue - + docs = item["docs"] name = item["name"] + # Decode JSON string try: docs = json.loads(docs) except json.JSONDecodeError: pass - + docs = self._clean_html(docs) - item["docs"] = docs + # Save all fields for adding to metadata later + item["docs"] = docs # replace docs with cleaned text metadata_dict[name] = item + # Split by headers, and where needed, sentences splits = self._split_by_headers(docs) splits = self._split_oversized_chunks(chunks=splits, target_length=self.target_length) chunks = self._accumulate_chunks(splits=splits, target_length=self.target_length, overlap=self.overlap, min_length=self.min_length) for chunk in chunks: output.append({"name": name, "docs_type": self.docs_type, "doc_chunk": chunk}) - + self._write_chunks_to_file(chunks=output, file_name=f"{self.docs_type}_chunks.json") - return output, metadata_dict + return output, metadata_dict def _clean_html(self, text): """Remove HTML tags while preserving essential formatting.""" @@ -88,7 +89,7 @@ def _clean_html(self, text): text = re.sub(r'<\/?strong>', '**', text) # Convert to bold text = re.sub(r'<[^>]+>', '', text) # Remove other HTML tags - return text.rstrip() + return text.strip() def _split_by_headers(self, text): """Split text into chunks based on Markdown headers (# and ##) and code blocks.""" @@ -99,7 +100,7 @@ def _split_by_headers(self, text): def _split_oversized_chunks(self, chunks, target_length): """Check if chunks are over the target lengths, and split them further if needed.""" result = [] - + for chunk in chunks: if len(chunk) <= target_length: result.append(chunk) @@ -107,7 +108,7 @@ def _split_oversized_chunks(self, chunks, target_length): # Chunk is too big, split by newlines lines = chunk.split('\n') current_chunk = "" - + for line in lines: # If adding this line would exceed target size and we already have content if len(current_chunk) + len(line) + 1 > target_length and current_chunk: @@ -118,11 +119,11 @@ def _split_oversized_chunks(self, chunks, target_length): if current_chunk: current_chunk += '\n' current_chunk += line - + # Add the last chunk if current_chunk: result.append(current_chunk) - + return result def _accumulate_chunks(self, splits, target_length, overlap, min_length): @@ -130,25 +131,25 @@ def _accumulate_chunks(self, splits, target_length, overlap, min_length): accumulated = [] current_chunk = "" last_overlap_length = 0 - + for split in splits: if len(current_chunk) + len(split) <= target_length: current_chunk += split - elif len(current_chunk) >= min_length: - accumulated.append(current_chunk) # Store the completed chunk - - # add overlap - if self.docs_type == "adaptor_functions": - overlap_sections = " ".join(current_chunk.split("\n")[-overlap:]) - else: - # Split by sentences (doesn't split code) - overlap_sections = " ".join(nltk.sent_tokenize(current_chunk)[-overlap:]) - current_chunk = overlap_sections + split # Start a new chunk - last_overlap_length = len(overlap_sections) else: - # Current chunk is too small, add the next split even though it exceeds target_length - current_chunk += split + if len(current_chunk) >= min_length: + accumulated.append(current_chunk) # Store the completed chunk + # add overlap + if self.docs_type == "adaptor_functions": + overlap_sections = " ".join(current_chunk.split("\n")[-overlap:]) + else: + overlap_sections = " ".join(nltk.sent_tokenize(current_chunk)[-overlap:]) # Split by sentences (doesn't split code) + current_chunk = overlap_sections + split # Start a new chunk + last_overlap_length = len(overlap_sections) + else: + # Current chunk is too small, add the next split even though it exceeds target_length + current_chunk += split + if current_chunk: if len(current_chunk) >= min_length or len(accumulated)==0: accumulated.append(current_chunk) diff --git a/services/embed_docsite/tests/unit/test_docsite_processor.py b/services/embed_docsite/tests/unit/test_docsite_processor.py index 52020f1d..d390df57 100644 --- a/services/embed_docsite/tests/unit/test_docsite_processor.py +++ b/services/embed_docsite/tests/unit/test_docsite_processor.py @@ -15,7 +15,7 @@ def make_processor(**kwargs): def test_clean_html_converts_tags(): p = make_processor() result = p._clean_html("

Hello

x bold drop") - assert result == "\nHello\n `x` **bold** drop" + assert result == "Hello\n `x` **bold** drop" def test_split_by_headers_splits_on_markdown_headers(): From 8a99866f3dbabe8dbf3c0a3b2e5d5472ec2e4566 Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Tue, 28 Jul 2026 22:55:01 +0800 Subject: [PATCH 23/48] refactor: revert import and formatting churn in docsite consumers Co-Authored-By: Claude Opus 5 --- services/job_chat/retrieve_docs.py | 24 ++++++++++--------- .../search_documentation.py | 11 +++++---- 2 files changed, 20 insertions(+), 15 deletions(-) diff --git a/services/job_chat/retrieve_docs.py b/services/job_chat/retrieve_docs.py index 2d7c9cd7..22b247f5 100644 --- a/services/job_chat/retrieve_docs.py +++ b/services/job_chat/retrieve_docs.py @@ -1,24 +1,22 @@ -import json import os - +import json import anthropic -import sentry_sdk from anthropic import ( APIConnectionError, - AuthenticationError, BadRequestError, - InternalServerError, - NotFoundError, + AuthenticationError, PermissionDeniedError, - RateLimitError, + NotFoundError, UnprocessableEntityError, + RateLimitError, + InternalServerError, ) +import sentry_sdk from langfuse import observe +from util import ApolloError, create_logger from models import resolve_model -from search_docsite.pinecone_legacy_search import LegacyPineconeDocsiteSearch from search_docsite.search_docsite import DocsiteSearch -from util import ApolloError, create_logger - +from search_docsite.pinecone_legacy_search import LegacyPineconeDocsiteSearch from .rag_config_loader import ConfigLoader logger = create_logger("job_chat.retrieve_docs") @@ -76,7 +74,11 @@ def retrieve_knowledge(content, history, code="", adaptor="", api_key=None, stre search_queries, generate_queries_usage = generate_queries(content, client, user_context) with sentry_sdk.start_span(description="search_documentation"): try: - search_results = search_docs(search_queries, top_k=config["top_k"], threshold=config["threshold"]) + search_results = search_docs( + search_queries, + top_k=config["top_k"], + threshold=config["threshold"] + ) search_results = list(set(search_results)) search_results_sections = list(set(result.metadata["doc_title"] for result in search_results)) except Exception as e: diff --git a/services/tools/search_documentation/search_documentation.py b/services/tools/search_documentation/search_documentation.py index fd0f61cc..5470e2de 100644 --- a/services/tools/search_documentation/search_documentation.py +++ b/services/tools/search_documentation/search_documentation.py @@ -7,16 +7,16 @@ """ import os import sys -from dataclasses import dataclass from pathlib import Path from typing import Dict +from dataclasses import dataclass # Import utilities from services directory sys.path.append(str(Path(__file__).parent.parent.parent)) -from search_docsite.pinecone_legacy_search import LegacyPineconeDocsiteSearch +from util import create_logger, ApolloError from search_docsite.search_docsite import DocsiteSearch -from util import ApolloError, create_logger +from search_docsite.pinecone_legacy_search import LegacyPineconeDocsiteSearch logger = create_logger(__name__) @@ -54,12 +54,15 @@ def _search_implementation(query: str, num_results: int) -> Dict: # service and run_eval for evaluation. backend = os.environ.get("DOCSITE_SEARCH_BACKEND", "pinecone") search_cls = DocsiteSearch if backend == "postgres" else LegacyPineconeDocsiteSearch + # Initialize docsite search docsite_search = search_cls() + + # Search with threshold for quality results search_results = docsite_search.search( query=query, top_k=num_results, threshold=0.7, # Only return relevant results - strategy='semantic', + strategy='semantic' ) logger.info(f"Found {len(search_results)} documentation results") From 0a021a22f5afa3e8724737fcaa29dc75239dea56 Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Tue, 28 Jul 2026 23:06:57 +0800 Subject: [PATCH 24/48] refactor: drop type annotations from internal helpers --- services/db_migrations.py | 4 ++-- services/embed_docsite/docsite_indexer.py | 15 +++++++-------- services/embed_docsite/embed_docsite.py | 15 +++------------ services/search_docsite/tests/eval/run_eval.py | 8 ++++---- 4 files changed, 16 insertions(+), 26 deletions(-) diff --git a/services/db_migrations.py b/services/db_migrations.py index 26f573ac..a96b5ebd 100644 --- a/services/db_migrations.py +++ b/services/db_migrations.py @@ -29,14 +29,14 @@ """ -def _migration_files() -> list: +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) -> int: +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 diff --git a/services/embed_docsite/docsite_indexer.py b/services/embed_docsite/docsite_indexer.py index d47a02b2..55333c84 100644 --- a/services/embed_docsite/docsite_indexer.py +++ b/services/embed_docsite/docsite_indexer.py @@ -1,4 +1,3 @@ -from typing import Optional from psycopg2.extras import execute_values from langchain_openai import OpenAIEmbeddings from pgvector.psycopg2 import register_vector @@ -41,7 +40,7 @@ def embeddings(self): self._embeddings = OpenAIEmbeddings() return self._embeddings - def start_batch(self, conn, docs_types: list) -> int: + def start_batch(self, conn, docs_types): """Insert a new 'building' batch row and return its id.""" sql = """ INSERT INTO docsite_batches (status, docs_types, chunk_target_length, chunk_min_length, embedding_model) @@ -55,7 +54,7 @@ def start_batch(self, conn, docs_types: list) -> int: logger.info(f"Started batch {batch_id} for docs_types={docs_types}") return batch_id - def insert_documents(self, conn, batch_id: int, documents: list, metadata_dict: dict) -> int: + def insert_documents(self, conn, batch_id, documents, metadata_dict): """Embed and bulk-insert chunks for this batch. Returns the number of chunks inserted.""" if not documents: return 0 @@ -82,7 +81,7 @@ def insert_documents(self, conn, batch_id: int, documents: list, metadata_dict: logger.info(f"Inserted {len(rows)} chunks into batch {batch_id}") return len(rows) - def _embed_in_batches(self, texts: list, batch_size: int = 100) -> list: + def _embed_in_batches(self, texts, batch_size=100): """Call the OpenAI embeddings API in batches of batch_size texts.""" embeddings = [] for i in range(0, len(texts), batch_size): @@ -90,7 +89,7 @@ def _embed_in_batches(self, texts: list, batch_size: int = 100) -> list: embeddings.extend(self.embeddings.embed_documents(batch)) return embeddings - def copy_forward_missing_docs_types(self, conn, batch_id: int, docs_types_present: list) -> int: + def copy_forward_missing_docs_types(self, conn, batch_id, docs_types_present): """Copy chunks for docs_types NOT in this run from the previous complete batch, so every complete batch is a full snapshot across all docs_types. Returns rows copied.""" missing_types = [t for t in ALL_DOCS_TYPES if t not in docs_types_present] @@ -120,7 +119,7 @@ def copy_forward_missing_docs_types(self, conn, batch_id: int, docs_types_presen logger.info(f"Copied {copied} chunks forward for docs_types={missing_types} from batch {previous_batch_id}") return copied - def build_index(self, conn, batch_id: int) -> None: + def build_index(self, conn, batch_id): """Build a per-batch partial HNSW index. Runs outside a transaction (autocommit).""" conn.autocommit = True try: @@ -137,7 +136,7 @@ def build_index(self, conn, batch_id: int) -> None: conn.autocommit = False logger.info(f"Built HNSW index for batch {batch_id}") - def promote_batch(self, conn, batch_id: int, chunk_count: int) -> None: + def promote_batch(self, conn, batch_id, chunk_count): """Flip a batch to 'complete' — the moment it becomes visible to readers.""" with conn.cursor() as cur: cur.execute( @@ -147,7 +146,7 @@ def promote_batch(self, conn, batch_id: int, chunk_count: int) -> None: conn.commit() logger.info(f"Promoted batch {batch_id} ({chunk_count} chunks)") - def prune_old_batches(self, conn, keep_batches: Optional[int] = None) -> list: + def prune_old_batches(self, conn, keep_batches=None): """Delete complete batches older than the newest `keep_batches`, dropping their partial indexes first. Returns the list of pruned batch ids.""" keep = keep_batches if keep_batches is not None else self.keep_batches diff --git a/services/embed_docsite/embed_docsite.py b/services/embed_docsite/embed_docsite.py index c4da2194..4a2433c2 100644 --- a/services/embed_docsite/embed_docsite.py +++ b/services/embed_docsite/embed_docsite.py @@ -15,9 +15,7 @@ VALID_TARGETS = ("pinecone", "postgres") -def _collect_documents( - docs_to_upload: list, docs_to_ignore: list, chunk_target_length: int, chunk_min_length: int, -) -> tuple[list, dict]: +def _collect_documents(docs_to_upload, docs_to_ignore, chunk_target_length, chunk_min_length): """Download and chunk every requested docs_type. Shared by both targets.""" documents = [] metadata_dict = {} @@ -71,7 +69,7 @@ def main(data: dict) -> dict: ) -def _upload_to_pinecone(data: dict, documents: list, metadata_dict: dict, docs_to_upload: list) -> dict: +def _upload_to_pinecone(data, documents, metadata_dict, docs_to_upload): """Legacy write path. Deliberately opens no Postgres connection.""" pinecone_api_key = data.get("PINECONE_API_KEY") or os.environ.get("PINECONE_API_KEY") if not pinecone_api_key: @@ -95,14 +93,7 @@ def _upload_to_pinecone(data: dict, documents: list, metadata_dict: dict, docs_t } -def _upload_to_postgres( - documents: list, - metadata_dict: dict, - docs_to_upload: list, - chunk_target_length: int, - chunk_min_length: int, - keep_batches: int, -) -> dict: +def _upload_to_postgres(documents, metadata_dict, docs_to_upload, chunk_target_length, chunk_min_length, keep_batches): indexer = DocsiteIndexer( chunk_target_length=chunk_target_length, chunk_min_length=chunk_min_length, diff --git a/services/search_docsite/tests/eval/run_eval.py b/services/search_docsite/tests/eval/run_eval.py index a0a4e184..02ed54a3 100644 --- a/services/search_docsite/tests/eval/run_eval.py +++ b/services/search_docsite/tests/eval/run_eval.py @@ -19,12 +19,12 @@ GOLDEN_QUERIES_PATH = Path(__file__).parent / "golden_queries.yaml" -def compute_recall_at_k(retrieved_titles: list, expected_titles: list) -> bool: +def compute_recall_at_k(retrieved_titles, expected_titles): """True if any expected title appears among the retrieved titles.""" return bool(set(retrieved_titles) & set(expected_titles)) -def compute_agreement(report_a: dict, report_b: dict) -> dict: +def compute_agreement(report_a, report_b): """Doc-title overlap between two backends' result sets, per query and averaged. Needs no ground truth, so it gives a usable comparison signal before @@ -48,7 +48,7 @@ def compute_agreement(report_a: dict, report_b: dict) -> dict: return {"per_query": per_query, "mean_jaccard": mean_jaccard} -def run_eval(golden_queries: list, backend_cls, strategy: str, top_k: int = 5) -> dict: +def run_eval(golden_queries, backend_cls, strategy, top_k=5): """Run every golden query against one backend/strategy and return a report dict.""" backend = backend_cls() per_query = [] @@ -99,7 +99,7 @@ def run_eval(golden_queries: list, backend_cls, strategy: str, top_k: int = 5) - } -def main() -> None: +def main(): with open(GOLDEN_QUERIES_PATH) as f: golden_queries = yaml.safe_load(f)["queries"] From f2f540fc2dfbc15d87b43961298880f44f683a5c Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Tue, 28 Jul 2026 23:26:17 +0800 Subject: [PATCH 25/48] refactor: extract shared docsite backend resolver --- services/job_chat/retrieve_docs.py | 38 +++++---------- .../job_chat/tests/unit/test_retrieve_docs.py | 48 +++++++++---------- services/search_docsite/search_docsite.py | 19 +++++--- .../tests/unit/test_search_docsite_main.py | 22 +++++++++ .../search_documentation.py | 8 +--- .../tests/unit/test_search_documentation.py | 32 ++++--------- 6 files changed, 79 insertions(+), 88 deletions(-) diff --git a/services/job_chat/retrieve_docs.py b/services/job_chat/retrieve_docs.py index 22b247f5..59db8d83 100644 --- a/services/job_chat/retrieve_docs.py +++ b/services/job_chat/retrieve_docs.py @@ -15,8 +15,7 @@ from langfuse import observe from util import ApolloError, create_logger from models import resolve_model -from search_docsite.search_docsite import DocsiteSearch -from search_docsite.pinecone_legacy_search import LegacyPineconeDocsiteSearch +from search_docsite.search_docsite import resolve_backend from .rag_config_loader import ConfigLoader logger = create_logger("job_chat.retrieve_docs") @@ -171,31 +170,20 @@ def generate_queries(content, client, user_context=""): return (answer_parsed, usage) def search_docs(search_queries, top_k, threshold=None): - """Search the docsite store using search queries. Defaults to the legacy - Pinecone backend; set DOCSITE_SEARCH_BACKEND=postgres to switch to the - Postgres-backed search. - - Both backends use semantic search with the same cosine-similarity cutoff, so - results are directly comparable and the quality gate survives the cutover. - Hybrid (RRF) is deliberately not used here: its score has no calibratable - scale, so a threshold cannot be applied to it. It stays available via the - search_docsite service and run_eval for evaluation. - - :param threshold: Cosine-similarity cutoff, applied identically by both - backends.""" - backend = os.environ.get("DOCSITE_SEARCH_BACKEND", "pinecone") - backend_cls = DocsiteSearch if backend == "postgres" else LegacyPineconeDocsiteSearch - return _run_backend_search(backend_cls, "semantic", search_queries, top_k, threshold) - - -def _run_backend_search(backend_cls, strategy, search_queries, top_k, threshold=None): - searcher = backend_cls() - results = [] + """Search the docsite store using search queries.""" + searcher = resolve_backend()() + search_results = [] for q in search_queries: - results.extend( - searcher.search(q.get("query"), top_k=top_k, threshold=threshold, strategy=strategy, docs_type="general_docs"), + query_search_result = searcher.search( + q.get("query"), + top_k=top_k, + threshold=threshold, + strategy="semantic", + docs_type="general_docs" ) - return results + search_results.extend(query_search_result) + + return search_results def format_context(adaptor, code, history): """Optionally add more context about the user's job for the LLM.""" diff --git a/services/job_chat/tests/unit/test_retrieve_docs.py b/services/job_chat/tests/unit/test_retrieve_docs.py index 5dcd4ff9..d53387e1 100644 --- a/services/job_chat/tests/unit/test_retrieve_docs.py +++ b/services/job_chat/tests/unit/test_retrieve_docs.py @@ -80,56 +80,52 @@ def test_call_llm_wraps_unexpected_error_as_apollo_error(): assert exc.value.type == "UNKNOWN_ERROR" -# --- search_docs (backend flag + shadow mode) ----------------------------------- +# --- search_docs ---------------------------------------------------------------- def _fake_result(title): return SearchResult(f"text for {title}", {"doc_title": title, "docs_type": "general_docs"}, 0.9) -def test_search_docs_defaults_to_legacy_pinecone_backend(monkeypatch): +def test_search_docs_forwards_query_args_to_resolved_backend(monkeypatch): monkeypatch.delenv("DOCSITE_SEARCH_BACKEND", raising=False) - monkeypatch.delenv("DOCSITE_SHADOW_POSTGRES", raising=False) - with patch.object(rd, "LegacyPineconeDocsiteSearch") as mock_legacy_cls: - mock_legacy_cls.return_value.search.return_value = [_fake_result("A")] + with patch.object(rd, "resolve_backend") as mock_resolve: + backend = mock_resolve.return_value.return_value + backend.search.return_value = [_fake_result("A")] results = search_docs([{"query": "q"}], top_k=5) assert [r.metadata["doc_title"] for r in results] == ["A"] - mock_legacy_cls.return_value.search.assert_called_once_with( + backend.search.assert_called_once_with( "q", top_k=5, threshold=None, strategy="semantic", docs_type="general_docs" ) -def test_search_docs_passes_threshold_through_to_semantic_backend(monkeypatch): - """Threshold is a score cutoff that only makes sense for the Pinecone/semantic - path; it must still be forwarded there (this regressed once already when the - backend flag was introduced — rag.yaml's threshold silently stopped applying).""" +def test_search_docs_passes_threshold_through(monkeypatch): + """Threshold is a cosine-similarity cutoff that must reach the backend — this + regressed once already when the backend flag was introduced, and rag.yaml's + threshold silently stopped applying.""" monkeypatch.delenv("DOCSITE_SEARCH_BACKEND", raising=False) - monkeypatch.delenv("DOCSITE_SHADOW_POSTGRES", raising=False) - with patch.object(rd, "LegacyPineconeDocsiteSearch") as mock_legacy_cls: - mock_legacy_cls.return_value.search.return_value = [_fake_result("A")] + with patch.object(rd, "resolve_backend") as mock_resolve: + backend = mock_resolve.return_value.return_value + backend.search.return_value = [_fake_result("A")] search_docs([{"query": "q"}], top_k=5, threshold=0.8) - mock_legacy_cls.return_value.search.assert_called_once_with( + backend.search.assert_called_once_with( "q", top_k=5, threshold=0.8, strategy="semantic", docs_type="general_docs" ) -def test_search_docs_uses_semantic_with_threshold_on_postgres_backend(monkeypatch): - """Postgres must apply the identical 0.8 cosine gate as Pinecone. Hybrid's RRF - score cannot be thresholded, so using it here would silently drop the quality - filter at cutover. Semantic returns comparable cosine scores.""" - monkeypatch.setenv("DOCSITE_SEARCH_BACKEND", "postgres") +def test_search_docs_accumulates_results_across_queries(monkeypatch): + monkeypatch.delenv("DOCSITE_SEARCH_BACKEND", raising=False) - with patch.object(rd, "DocsiteSearch") as mock_pg_cls: - mock_pg_cls.return_value.search.return_value = [_fake_result("B")] - results = search_docs([{"query": "q"}], top_k=5, threshold=0.8) + with patch.object(rd, "resolve_backend") as mock_resolve: + backend = mock_resolve.return_value.return_value + backend.search.side_effect = [[_fake_result("A")], [_fake_result("B")]] + results = search_docs([{"query": "q1"}, {"query": "q2"}], top_k=5) - assert [r.metadata["doc_title"] for r in results] == ["B"] - mock_pg_cls.return_value.search.assert_called_once_with( - "q", top_k=5, threshold=0.8, strategy="semantic", docs_type="general_docs" - ) + assert [r.metadata["doc_title"] for r in results] == ["A", "B"] + assert backend.search.call_count == 2 diff --git a/services/search_docsite/search_docsite.py b/services/search_docsite/search_docsite.py index f757a5f7..6ccadc37 100644 --- a/services/search_docsite/search_docsite.py +++ b/services/search_docsite/search_docsite.py @@ -196,6 +196,17 @@ def _hybrid_search(self, conn, batch_id, query, top_k, doc_title, docs_type): } +def resolve_backend(override=None): + """Return the search class for the configured backend. + + :param override: Backend name that is prioritised over DOCSITE_SEARCH_BACKEND + """ + name = override or os.environ.get("DOCSITE_SEARCH_BACKEND", "pinecone") + if name not in BACKEND_INDEX_PARAMS: + raise ApolloError(400, f"Unknown backend '{name}'. Expected 'pinecone' or 'postgres'", type="BAD_REQUEST") + return DocsiteSearch if name == "postgres" else LegacyPineconeDocsiteSearch + + def main(data): logger.info("Starting...") @@ -206,12 +217,7 @@ def main(data): return None backend = data.get("backend") or os.environ.get("DOCSITE_SEARCH_BACKEND", "pinecone") - if backend not in BACKEND_INDEX_PARAMS: - raise ApolloError( - 400, - f"Unknown backend '{backend}'. Expected 'pinecone' or 'postgres'", - type="BAD_REQUEST", - ) + search_cls = resolve_backend(backend) search_params = {"query": data["query"]} optional_search_params = ["docs_type", "doc_title", "top_k", "threshold", "strategy"] @@ -228,7 +234,6 @@ def main(data): logger.error(msg) raise ApolloError(500, msg, type="BAD_REQUEST") - search_cls = DocsiteSearch if backend == "postgres" else LegacyPineconeDocsiteSearch logger.info(f"Searching docsite via the {backend} backend") docsite_search = search_cls(**index_params) diff --git a/services/search_docsite/tests/unit/test_search_docsite_main.py b/services/search_docsite/tests/unit/test_search_docsite_main.py index 2db6d5af..98bb0e31 100644 --- a/services/search_docsite/tests/unit/test_search_docsite_main.py +++ b/services/search_docsite/tests/unit/test_search_docsite_main.py @@ -75,3 +75,25 @@ def test_main_routes_collection_name_to_pinecone_only(monkeypatch): m.main({"query": "q", "backend": "pinecone", "collection_name": "docsite-202501010000", "batch_id": 9}) mock_legacy_cls.assert_called_once_with(collection_name="docsite-202501010000") + + +def test_resolve_backend_defaults_to_pinecone(monkeypatch): + monkeypatch.delenv("DOCSITE_SEARCH_BACKEND", raising=False) + assert m.resolve_backend() is m.LegacyPineconeDocsiteSearch + + +def test_resolve_backend_reads_env(monkeypatch): + monkeypatch.setenv("DOCSITE_SEARCH_BACKEND", "postgres") + assert m.resolve_backend() is m.DocsiteSearch + + +def test_resolve_backend_override_wins_over_env(monkeypatch): + monkeypatch.setenv("DOCSITE_SEARCH_BACKEND", "pinecone") + assert m.resolve_backend("postgres") is m.DocsiteSearch + + +def test_resolve_backend_rejects_unknown_name(monkeypatch): + monkeypatch.delenv("DOCSITE_SEARCH_BACKEND", raising=False) + with pytest.raises(ApolloError) as exc: + m.resolve_backend("sqlite") + assert exc.value.code == 400 diff --git a/services/tools/search_documentation/search_documentation.py b/services/tools/search_documentation/search_documentation.py index 5470e2de..65697a20 100644 --- a/services/tools/search_documentation/search_documentation.py +++ b/services/tools/search_documentation/search_documentation.py @@ -5,7 +5,6 @@ 1. As a standalone service via entry.py: bun py tools/search_documentation 2. As a tool by supervisor via search_documentation_tool() """ -import os import sys from pathlib import Path from typing import Dict @@ -15,8 +14,7 @@ sys.path.append(str(Path(__file__).parent.parent.parent)) from util import create_logger, ApolloError -from search_docsite.search_docsite import DocsiteSearch -from search_docsite.pinecone_legacy_search import LegacyPineconeDocsiteSearch +from search_docsite.search_docsite import resolve_backend logger = create_logger(__name__) @@ -52,10 +50,8 @@ def _search_implementation(query: str, num_results: int) -> Dict: # its score has no calibratable scale, so it can be neither thresholded nor # rendered as a relevance figure. It stays available via the search_docsite # service and run_eval for evaluation. - backend = os.environ.get("DOCSITE_SEARCH_BACKEND", "pinecone") - search_cls = DocsiteSearch if backend == "postgres" else LegacyPineconeDocsiteSearch # Initialize docsite search - docsite_search = search_cls() + docsite_search = resolve_backend()() # Search with threshold for quality results search_results = docsite_search.search( diff --git a/services/tools/search_documentation/tests/unit/test_search_documentation.py b/services/tools/search_documentation/tests/unit/test_search_documentation.py index 20689e18..569f8af9 100644 --- a/services/tools/search_documentation/tests/unit/test_search_documentation.py +++ b/services/tools/search_documentation/tests/unit/test_search_documentation.py @@ -1,36 +1,20 @@ -"""Unit tests for search_documentation's backend-flag selection. -No prior tests existed for this module.""" +"""Unit tests for search_documentation's delegation to the resolved backend.""" from unittest.mock import patch import tools.search_documentation.search_documentation as m -def test_search_implementation_uses_legacy_pinecone_by_default(monkeypatch): - """Characterization test: pins the exact call main made. The 0.7 threshold is - a quality gate on live traffic — it was silently dropped once already.""" +def test_search_implementation_applies_quality_gate_to_resolved_backend(monkeypatch): + """The 0.7 threshold is a quality gate on live traffic — it was silently + dropped once already.""" monkeypatch.delenv("DOCSITE_SEARCH_BACKEND", raising=False) - with patch.object(m, "LegacyPineconeDocsiteSearch") as mock_legacy_cls: - mock_legacy_cls.return_value.search.return_value = [] + with patch.object(m, "resolve_backend") as mock_resolve: + backend = mock_resolve.return_value.return_value + backend.search.return_value = [] m._search_implementation("how do I use webhooks", 5) - mock_legacy_cls.return_value.search.assert_called_once_with( - query="how do I use webhooks", top_k=5, threshold=0.7, strategy="semantic" - ) - - -def test_search_implementation_uses_semantic_with_same_threshold_on_postgres(monkeypatch): - """Postgres must apply the identical quality gate. Hybrid's RRF score has a - ~0.033 ceiling — it cannot be thresholded, and renders as a constant 0.03 to - the LLM. Semantic returns true cosine similarity, directly comparable to - Pinecone, so the same 0.7 cutoff means the same thing on both backends.""" - monkeypatch.setenv("DOCSITE_SEARCH_BACKEND", "postgres") - - with patch.object(m, "DocsiteSearch") as mock_pg_cls: - mock_pg_cls.return_value.search.return_value = [] - m._search_implementation("how do I use webhooks", 5) - - mock_pg_cls.return_value.search.assert_called_once_with( + backend.search.assert_called_once_with( query="how do I use webhooks", top_k=5, threshold=0.7, strategy="semantic" ) From 62b77db9ac3b35217412be13363e88b027811c92 Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Tue, 28 Jul 2026 23:35:16 +0800 Subject: [PATCH 26/48] docs: focus docsite comments on invariants --- services/db_migrations.py | 14 ++++++-------- services/embed_docsite/docsite_indexer.py | 6 +++--- services/job_chat/retrieve_docs.py | 8 +++++++- services/search_docsite/pinecone_legacy_search.py | 12 +++--------- services/search_docsite/search_docsite.py | 5 ++--- services/search_docsite/tests/eval/run_eval.py | 10 ++++------ .../tests/unit/test_search_docsite_main.py | 6 +++--- services/util.py | 4 +--- 8 files changed, 29 insertions(+), 36 deletions(-) diff --git a/services/db_migrations.py b/services/db_migrations.py index a96b5ebd..47e09f74 100644 --- a/services/db_migrations.py +++ b/services/db_migrations.py @@ -1,13 +1,11 @@ """Versioned schema migrations for the Python-owned docs database (POSTGRES_URL). -Mirrors platform/src/db/migrate.ts, which owns the TypeScript-side auth database. -The two runners are deliberately kept separate but symmetrical: .sql files applied -in lexical order, applied filenames recorded so re-runs are a no-op, and an -advisory lock so concurrent starters queue rather than collide. - -The tracking table (_migrations_docs) and lock key (8314_2026) are both distinct -from the TypeScript runner's, because APOLLO_CLIENTS_DB_URL falls back to -POSTGRES_URL in local development and both runners can target one database. +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 diff --git a/services/embed_docsite/docsite_indexer.py b/services/embed_docsite/docsite_indexer.py index 55333c84..0df3c4a3 100644 --- a/services/embed_docsite/docsite_indexer.py +++ b/services/embed_docsite/docsite_indexer.py @@ -19,7 +19,7 @@ class DocsiteIndexer: A batch is a full, self-consistent snapshot across all docs_types. Batches are built invisibly (status='building'), then promoted to 'complete' - atomically — replacing Pinecone's timestamped-namespace-per-run pattern. + atomically, so readers only ever see a finished batch. :param chunk_target_length: Target chunk size in characters (default: 1000) :param chunk_min_length: Minimum chunk size before merging with the next split (default: 700) @@ -34,8 +34,8 @@ def __init__(self, chunk_target_length=1000, chunk_min_length=700, keep_batches= @property def embeddings(self): - """Lazily construct the OpenAI embeddings client (avoids eager credential - validation at import/instantiation time).""" + """Lazily construct the OpenAI embeddings client, so importing this module + needs no credentials.""" if self._embeddings is None: self._embeddings = OpenAIEmbeddings() return self._embeddings diff --git a/services/job_chat/retrieve_docs.py b/services/job_chat/retrieve_docs.py index 59db8d83..523310b8 100644 --- a/services/job_chat/retrieve_docs.py +++ b/services/job_chat/retrieve_docs.py @@ -170,7 +170,13 @@ def generate_queries(content, client, user_context=""): return (answer_parsed, usage) def search_docs(search_queries, top_k, threshold=None): - """Search the docsite store using search queries.""" + """Search the docsite store. Both backends run semantic search, so the + threshold applies identically. + + Set DOCSITE_SEARCH_BACKEND=postgres to use Postgres. + + :param threshold: Cosine-similarity cutoff + """ searcher = resolve_backend()() search_results = [] for q in search_queries: diff --git a/services/search_docsite/pinecone_legacy_search.py b/services/search_docsite/pinecone_legacy_search.py index c131e72b..b1b0016b 100644 --- a/services/search_docsite/pinecone_legacy_search.py +++ b/services/search_docsite/pinecone_legacy_search.py @@ -10,15 +10,9 @@ class LegacyPineconeDocsiteSearch: """ - Legacy Pinecone-backed docsite search, preserved from before the Postgres - migration as the rollback path. Since DOCSITE_SEARCH_BACKEND defaults to - "pinecone", this implementation is what actually serves search traffic - today whenever that flag is unset or explicitly set to "pinecone". It is - constructed directly by services/search_docsite/search_docsite.py's - main(), services/job_chat/retrieve_docs.py's search_docs(), and - services/tools/search_documentation/search_documentation.py's - _search_implementation(), each of which select it via the - DOCSITE_SEARCH_BACKEND flag. + Legacy Pinecone-backed docsite search, still the default backend and the + rollback path for the Postgres migration. Selected via resolve_backend() + in services/search_docsite/search_docsite.py. :param collection_name: Vectorstore collection name (namespace) to store documents :param index_name: Vectorstore index name (default: docsite) diff --git a/services/search_docsite/search_docsite.py b/services/search_docsite/search_docsite.py index 6ccadc37..492be93c 100644 --- a/services/search_docsite/search_docsite.py +++ b/services/search_docsite/search_docsite.py @@ -40,9 +40,8 @@ def search(self, query, top_k=None, threshold=None, strategy='semantic', doc_tit :param query: Search query string :param top_k: Number of results to return - :param threshold: Score threshold. Only valid for strategy='semantic' — the - keyword (FTS rank) and hybrid (RRF rank) scores are not on a comparable - scale, so passing a threshold with them raises rather than being ignored. + :param threshold: Cosine-similarity cutoff. Valid only for + strategy='semantic'; raises for other strategies. :param strategy: 'semantic' | 'keyword' | 'hybrid' (default: 'semantic') :param doc_title: Filter by document title :param docs_type: Filter by document type diff --git a/services/search_docsite/tests/eval/run_eval.py b/services/search_docsite/tests/eval/run_eval.py index 02ed54a3..9044a19f 100644 --- a/services/search_docsite/tests/eval/run_eval.py +++ b/services/search_docsite/tests/eval/run_eval.py @@ -3,10 +3,9 @@ Usage: poetry run python -m search_docsite.tests.eval.run_eval Compares the Postgres-backed DocsiteSearch (strategy='hybrid') against -LegacyPineconeDocsiteSearch (strategy='semantic') over the golden query set, -per the Phase-1 shadow-mode rollout plan: Postgres must match or beat -Pinecone's recall@5 (with <=1-query regression tolerance) and p95 latency -before DOCSITE_SEARCH_BACKEND is flipped to 'postgres' by default. +LegacyPineconeDocsiteSearch (strategy='semantic') over the golden query set. +Postgres should match or beat Pinecone's recall@5 and p95 latency before +DOCSITE_SEARCH_BACKEND is flipped to 'postgres' by default. """ import time @@ -28,8 +27,7 @@ def compute_agreement(report_a, report_b): """Doc-title overlap between two backends' result sets, per query and averaged. Needs no ground truth, so it gives a usable comparison signal before - golden_queries.yaml is curated. This is the measure the deleted shadow mode - computed on live traffic; it belongs here, offline, instead. + golden_queries.yaml is curated. """ per_query = [] for a, b in zip(report_a["per_query"], report_b["per_query"]): diff --git a/services/search_docsite/tests/unit/test_search_docsite_main.py b/services/search_docsite/tests/unit/test_search_docsite_main.py index 98bb0e31..19af97d9 100644 --- a/services/search_docsite/tests/unit/test_search_docsite_main.py +++ b/services/search_docsite/tests/unit/test_search_docsite_main.py @@ -1,7 +1,7 @@ -"""Unit tests for search_docsite.main's backend selection. +"""Unit tests for search_docsite.main's backend selection and resolve_backend. -The `backend` payload field is the shadow-mode replacement: it lets the same -query be run against both backends on demand for manual comparison. +The `backend` payload field lets the same query be run against either backend +on demand for manual comparison. """ from unittest.mock import patch diff --git a/services/util.py b/services/util.py index 3ba1d0d5..abbc3bf8 100644 --- a/services/util.py +++ b/services/util.py @@ -113,9 +113,7 @@ def get_db_connection() -> "psycopg2.extensions.connection": """Get database connection from POSTGRES_URL environment variable. Applies any pending schema migrations before handing the connection back. - bridge.ts spawns a fresh Python process per request, so there is no - long-lived process to cache "already migrated" in; the runner's tracking - table makes repeat calls a cheap no-op. + The runner's tracking table makes repeat calls a cheap no-op. Returns: psycopg2.connection: Database connection From b8d753be72e2f611bba4ca0effb7426f5cdcb61d Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Wed, 29 Jul 2026 13:22:26 +0800 Subject: [PATCH 27/48] fix: run docsite migrations from the indexer, not every connection --- services/embed_docsite/embed_docsite.py | 4 ++++ .../tests/unit/test_db_migrations.py | 15 ++++++++++++ .../tests/unit/test_embed_docsite.py | 24 +++++++++++++++++++ services/util.py | 13 ++++------ 4 files changed, 48 insertions(+), 8 deletions(-) diff --git a/services/embed_docsite/embed_docsite.py b/services/embed_docsite/embed_docsite.py index 4a2433c2..8c5ba01b 100644 --- a/services/embed_docsite/embed_docsite.py +++ b/services/embed_docsite/embed_docsite.py @@ -1,5 +1,6 @@ import os +from db_migrations import run_migrations from dotenv import load_dotenv from embed_docsite.docsite_indexer import ( ALL_DOCS_TYPES, @@ -101,6 +102,9 @@ def _upload_to_postgres(documents, metadata_dict, docs_to_upload, chunk_target_l ) conn = get_db_connection() + # Order matters: register_vector looks up pgvector's type oid, so the + # extension must exist before we register. + run_migrations(conn) register_vector_type(conn) try: diff --git a/services/embed_docsite/tests/unit/test_db_migrations.py b/services/embed_docsite/tests/unit/test_db_migrations.py index 06c043cf..6468434a 100644 --- a/services/embed_docsite/tests/unit/test_db_migrations.py +++ b/services/embed_docsite/tests/unit/test_db_migrations.py @@ -107,3 +107,18 @@ def test_run_migrations_commits_once_at_the_end(tmp_path): m.run_migrations(conn) conn.commit.assert_called_once() + + +def test_get_db_connection_does_not_run_migrations(): + """Migrations belong to the indexer, not to every reader. CREATE EXTENSION + needs privileges managed Postgres withholds, so a reader that triggers it + 500s on a deployment that never enabled the Postgres docsite backend.""" + import util + + with patch.object(util, "psycopg2") as mock_psycopg2, \ + patch.object(m, "run_migrations") as mock_run, \ + patch.dict("os.environ", {"POSTGRES_URL": "postgresql://user@host/db"}): + conn = util.get_db_connection() + + assert conn is mock_psycopg2.connect.return_value + mock_run.assert_not_called() diff --git a/services/embed_docsite/tests/unit/test_embed_docsite.py b/services/embed_docsite/tests/unit/test_embed_docsite.py index 3dc7ffe0..647ccd0d 100644 --- a/services/embed_docsite/tests/unit/test_embed_docsite.py +++ b/services/embed_docsite/tests/unit/test_embed_docsite.py @@ -116,3 +116,27 @@ def test_main_rejects_unknown_target(): m.main({"target": "elasticsearch"}) assert exc.value.code == 400 + + +def test_postgres_upload_runs_migrations_before_registering_vector_type(): + """Order is load-bearing, not incidental: register_vector runs + to_regtype('vector') and raises unless CREATE EXTENSION has already run.""" + calls = [] + fake_conn = MagicMock() + fake_indexer = MagicMock() + fake_indexer.start_batch.return_value = 7 + fake_indexer.insert_documents.return_value = 1 + fake_indexer.copy_forward_missing_docs_types.return_value = 0 + fake_indexer.prune_old_batches.return_value = [] + fake_processor = MagicMock() + fake_processor.get_preprocessed_docs.return_value = ([], {}) + + with patch.object(m, "get_db_connection", return_value=fake_conn), \ + patch.object(m, "run_migrations", side_effect=lambda _conn: calls.append("migrate")), \ + patch.object(m, "register_vector_type", side_effect=lambda _conn: calls.append("register")), \ + patch.object(m, "DocsiteProcessor", return_value=fake_processor), \ + patch.object(m, "DocsiteIndexer", return_value=fake_indexer), \ + patch.dict("os.environ", {"OPENAI_API_KEY": "sk-test"}): + m.main({"docs_to_upload": ["general_docs"], "target": "postgres"}) + + assert calls == ["migrate", "register"] diff --git a/services/util.py b/services/util.py index abbc3bf8..b4164fe6 100644 --- a/services/util.py +++ b/services/util.py @@ -112,8 +112,10 @@ def apollo(name: str, payload: dict) -> dict: def get_db_connection() -> "psycopg2.extensions.connection": """Get database connection from POSTGRES_URL environment variable. - Applies any pending schema migrations before handing the connection back. - The runner's tracking table makes repeat calls a cheap no-op. + Returns a plain connection. Schema migrations belong to the indexer and are + run explicitly by embed_docsite: applying them here would put CREATE + EXTENSION in the path of every reader, including deployments that never use + the Postgres docsite backend and roles without the privilege to run it. Returns: psycopg2.connection: Database connection @@ -125,12 +127,7 @@ def get_db_connection() -> "psycopg2.extensions.connection": if not db_url: raise ApolloError(500, "Missing POSTGRES_URL environment variable", type="DATABASE_ERROR") - conn = psycopg2.connect(db_url) - - from db_migrations import run_migrations - run_migrations(conn) - - return conn + return psycopg2.connect(db_url) def sum_usage(*usage_objects): From 27ab6ac9986d6a146db86c61382ef4346af632eb Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Wed, 29 Jul 2026 13:25:05 +0800 Subject: [PATCH 28/48] Remove openai key test --- .../embed_docsite/tests/unit/test_embed_docsite.py | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/services/embed_docsite/tests/unit/test_embed_docsite.py b/services/embed_docsite/tests/unit/test_embed_docsite.py index 647ccd0d..9c72aeea 100644 --- a/services/embed_docsite/tests/unit/test_embed_docsite.py +++ b/services/embed_docsite/tests/unit/test_embed_docsite.py @@ -42,19 +42,6 @@ def test_main_orchestrates_full_batch_lifecycle_and_returns_summary(): "promoted": True, } - -def test_main_raises_when_openai_key_missing(): - from util import ApolloError - import pytest - - with patch.dict("os.environ", {}, clear=True): - with pytest.raises(ApolloError) as exc: - m.main({}) - - assert exc.value.code == 500 - assert "OPENAI_API_KEY" in exc.value.message - - def test_main_defaults_docs_to_upload_to_all_types(): fake_conn = MagicMock() fake_indexer = MagicMock() From c2476c03b8c4b17a2918e473b1c43f2d31d9c6b3 Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Wed, 29 Jul 2026 13:42:54 +0800 Subject: [PATCH 29/48] fix: return a clear 503 when the docsite schema is not initialised --- services/search_docsite/search_docsite.py | 33 +++++++++++++--- .../tests/unit/test_docsite_search.py | 39 +++++++++++++++++++ 2 files changed, 67 insertions(+), 5 deletions(-) diff --git a/services/search_docsite/search_docsite.py b/services/search_docsite/search_docsite.py index 492be93c..18fa8564 100644 --- a/services/search_docsite/search_docsite.py +++ b/services/search_docsite/search_docsite.py @@ -1,4 +1,5 @@ import os +import psycopg2 from dotenv import load_dotenv from langchain_openai import OpenAIEmbeddings from pgvector.psycopg2 import register_vector @@ -8,6 +9,10 @@ logger = create_logger("DocsiteSearch") +SCHEMA_MISSING_MESSAGE = ( + "Docsite schema not initialised — run embed_docsite with target=postgres" +) + def register_vector_type(conn): """Register the pgvector adapter on this connection.""" @@ -34,6 +39,22 @@ def embeddings(self): self._embeddings = OpenAIEmbeddings() return self._embeddings + def _connect(self): + """Open a connection with pgvector registered. + + A database the indexer has never touched fails here rather than at the + query: register_vector looks up pgvector's type oid and raises if the + extension is absent. 503, not 500 — it resolves by running the indexer, + with no redeploy. + """ + conn = get_db_connection() + try: + register_vector_type(conn) + except psycopg2.ProgrammingError as exc: + conn.close() + raise ApolloError(503, SCHEMA_MISSING_MESSAGE, type="DATABASE_ERROR") from exc + return conn + def search(self, query, top_k=None, threshold=None, strategy='semantic', doc_title=None, docs_type=None): """ Search docsite_chunks with optional filters. @@ -54,8 +75,7 @@ def search(self, query, top_k=None, threshold=None, strategy='semantic', doc_tit type="BAD_REQUEST", ) - conn = get_db_connection() - register_vector_type(conn) + conn = self._connect() try: batch_id = self._explicit_batch_id or self._resolve_current_batch(conn) @@ -72,9 +92,12 @@ def search(self, query, top_k=None, threshold=None, strategy='semantic', doc_tit def _resolve_current_batch(self, conn): """Find the newest complete batch id.""" - with conn.cursor() as cur: - cur.execute("SELECT id FROM docsite_batches WHERE status = 'complete' ORDER BY id DESC LIMIT 1") - row = cur.fetchone() + try: + with conn.cursor() as cur: + cur.execute("SELECT id FROM docsite_batches WHERE status = 'complete' ORDER BY id DESC LIMIT 1") + row = cur.fetchone() + except psycopg2.errors.UndefinedTable as exc: + raise ApolloError(503, SCHEMA_MISSING_MESSAGE, type="DATABASE_ERROR") from exc if row is None: raise ApolloError(404, "No complete docsite batch found", type="NOT_FOUND") return row[0] diff --git a/services/search_docsite/tests/unit/test_docsite_search.py b/services/search_docsite/tests/unit/test_docsite_search.py index 8308a74c..a9d96e52 100644 --- a/services/search_docsite/tests/unit/test_docsite_search.py +++ b/services/search_docsite/tests/unit/test_docsite_search.py @@ -9,6 +9,7 @@ from decimal import Decimal from unittest.mock import MagicMock, patch +import psycopg2 import pytest import search_docsite.search_docsite as m @@ -185,3 +186,41 @@ def test_hybrid_search_casts_rrf_to_float8_in_sql(): sql = cur.execute.call_args[0][0] assert "float8" in sql + + +# --- missing schema ---------------------------------------------------------- + +def test_search_maps_missing_pgvector_extension_to_503(): + """register_vector raises ProgrammingError('vector type not found in the + database') when CREATE EXTENSION has not run. Since migrations moved to the + indexer, that is the first thing a reader hits on an un-indexed database.""" + conn, _ = make_conn() + ds = make_search(batch_id=1) + + with patch.object(m, "get_db_connection", return_value=conn), \ + patch.object(m, "register_vector_type", + side_effect=psycopg2.ProgrammingError("vector type not found in the database")): + with pytest.raises(ApolloError) as exc: + ds.search("query") + + assert exc.value.code == 503 + assert "embed_docsite" in exc.value.message + conn.close.assert_called_once() + + +def test_search_maps_missing_docsite_tables_to_503(): + """Extension present, tables absent — the second way a database can be + un-indexed.""" + conn, cur = make_conn() + cur.execute.side_effect = psycopg2.errors.UndefinedTable( + 'relation "docsite_batches" does not exist', + ) + ds = make_search() + + with patch.object(m, "get_db_connection", return_value=conn), \ + patch.object(m, "register_vector_type"): + with pytest.raises(ApolloError) as exc: + ds.search("query") + + assert exc.value.code == 503 + assert "embed_docsite" in exc.value.message From 5d56511a53f85c9498b972d34e8b9926068cff81 Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Wed, 29 Jul 2026 13:55:16 +0800 Subject: [PATCH 30/48] fix: bind docsite embeddings as pgvector Vector --- poetry.lock | 13 ++++----- pyproject.toml | 2 +- services/embed_docsite/docsite_indexer.py | 10 +++++-- .../tests/unit/test_docsite_indexer.py | 5 ++-- services/search_docsite/search_docsite.py | 8 ++++-- .../tests/unit/test_docsite_search.py | 28 +++++++++++++++++++ 6 files changed, 51 insertions(+), 15 deletions(-) diff --git a/poetry.lock b/poetry.lock index d6f43638..3d85b154 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1822,19 +1822,16 @@ xml = ["lxml (>=4.9.2)"] [[package]] name = "pgvector" -version = "0.3.6" +version = "0.5.0" description = "pgvector support for Python" optional = false -python-versions = ">=3.8" +python-versions = ">=3.10" groups = ["main"] files = [ - {file = "pgvector-0.3.6-py3-none-any.whl", hash = "sha256:f6c269b3c110ccb7496bac87202148ed18f34b390a0189c783e351062400a75a"}, - {file = "pgvector-0.3.6.tar.gz", hash = "sha256:31d01690e6ea26cea8a633cde5f0f55f5b246d9c8292d68efdef8c22ec994ade"}, + {file = "pgvector-0.5.0-py3-none-any.whl", hash = "sha256:fedc9800894e6da2be51358d7b7c574bf34f247ca741a5a09513622135f5964f"}, + {file = "pgvector-0.5.0.tar.gz", hash = "sha256:07a9dcf735696879406983afc6eba9a787cef7c0cf6c367ca1a5779f036dee74"}, ] -[package.dependencies] -numpy = "*" - [[package]] name = "pinecone" version = "7.3.0" @@ -3760,4 +3757,4 @@ cffi = ["cffi (>=1.17,<2.0) ; platform_python_implementation != \"PyPy\" and pyt [metadata] lock-version = "2.1" python-versions = "3.11.*" -content-hash = "8e4463c369e05f92ff1c77750943b311f4af4699d23c487b032e7836166242ca" +content-hash = "e0f7dc02320564ac4dc7929bf5ff747716b4ab90810f4acc421db052f49014cc" diff --git a/pyproject.toml b/pyproject.toml index e3c82a6e..04470372 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,7 @@ psycopg2-binary = "^2.9.10" langfuse = "^4.14.1" opentelemetry-instrumentation-anthropic = "^0.62.1" opentelemetry-instrumentation-threading = "0.65b0" -pgvector = "^0.3.6" +pgvector = "^0.5.0" pandas = "^2.2" [tool.poetry.group.dev] diff --git a/services/embed_docsite/docsite_indexer.py b/services/embed_docsite/docsite_indexer.py index 0df3c4a3..7562e403 100644 --- a/services/embed_docsite/docsite_indexer.py +++ b/services/embed_docsite/docsite_indexer.py @@ -1,5 +1,6 @@ from psycopg2.extras import execute_values from langchain_openai import OpenAIEmbeddings +from pgvector import Vector from pgvector.psycopg2 import register_vector from util import create_logger @@ -9,7 +10,12 @@ def register_vector_type(conn): - """Register the pgvector adapter on this connection so Python lists convert to the `vector` type.""" + """Register pgvector's psycopg2 adapters on this connection. + + Registers `Vector` and `numpy.ndarray` — *not* `list`. A bare list is sent + as `numeric[]`, which Postgres will cast on assignment but has no `<=>` + operator for, so writes succeed and searches fail. + """ register_vector(conn) @@ -68,7 +74,7 @@ def insert_documents(self, conn, batch_id, documents, metadata_dict): doc_title = doc["name"].removesuffix(".md") chunk_index = doc_title_indices.get(doc_title, 0) doc_title_indices[doc_title] = chunk_index + 1 - rows.append((batch_id, doc_title, doc["docs_type"], chunk_index, doc["doc_chunk"], embedding)) + rows.append((batch_id, doc_title, doc["docs_type"], chunk_index, doc["doc_chunk"], Vector(embedding))) insert_sql = """ INSERT INTO docsite_chunks (batch_id, doc_title, docs_type, chunk_index, text, embedding) diff --git a/services/embed_docsite/tests/unit/test_docsite_indexer.py b/services/embed_docsite/tests/unit/test_docsite_indexer.py index 066ded2f..db23b04a 100644 --- a/services/embed_docsite/tests/unit/test_docsite_indexer.py +++ b/services/embed_docsite/tests/unit/test_docsite_indexer.py @@ -10,6 +10,7 @@ from unittest.mock import MagicMock, patch import embed_docsite.docsite_indexer as m +from pgvector import Vector def make_conn(): @@ -60,8 +61,8 @@ def test_insert_documents_embeds_and_bulk_inserts(): assert count == 2 indexer._embeddings.embed_documents.assert_called_once_with(["chunk one", "chunk two"]) rows = mock_execute_values.call_args[0][2] - assert rows[0] == (7, "doc-a", "general_docs", 0, "chunk one", [0.1, 0.2]) - assert rows[1] == (7, "doc-a", "general_docs", 1, "chunk two", [0.3, 0.4]) + assert rows[0] == (7, "doc-a", "general_docs", 0, "chunk one", Vector([0.1, 0.2])) + assert rows[1] == (7, "doc-a", "general_docs", 1, "chunk two", Vector([0.3, 0.4])) conn.commit.assert_called_once() diff --git a/services/search_docsite/search_docsite.py b/services/search_docsite/search_docsite.py index 18fa8564..a066b1af 100644 --- a/services/search_docsite/search_docsite.py +++ b/services/search_docsite/search_docsite.py @@ -2,6 +2,7 @@ import psycopg2 from dotenv import load_dotenv from langchain_openai import OpenAIEmbeddings +from pgvector import Vector from pgvector.psycopg2 import register_vector from util import create_logger, ApolloError, get_db_connection from embeddings.embeddings import SearchResult @@ -107,7 +108,10 @@ def _semantic_search(self, conn, batch_id, query, top_k, threshold, doc_title, d top_k = self.default_top_k max_k = top_k or 50 - query_embedding = self.embeddings.embed_query(query) + # Vector, not list: psycopg2 registers adapters for Vector and ndarray + # only, and a bare list arrives as numeric[], for which `<=>` has no + # operator. + query_embedding = Vector(self.embeddings.embed_query(query)) sql = """ SELECT text, doc_title, docs_type, 1 - (embedding <=> %(query_embedding)s) AS score @@ -164,7 +168,7 @@ def _hybrid_search(self, conn, batch_id, query, top_k, doc_title, docs_type): max_k = top_k or self.default_top_k candidate_k = 50 - query_embedding = self.embeddings.embed_query(query) + query_embedding = Vector(self.embeddings.embed_query(query)) sql = """ WITH semantic AS ( diff --git a/services/search_docsite/tests/unit/test_docsite_search.py b/services/search_docsite/tests/unit/test_docsite_search.py index a9d96e52..cc68093e 100644 --- a/services/search_docsite/tests/unit/test_docsite_search.py +++ b/services/search_docsite/tests/unit/test_docsite_search.py @@ -13,6 +13,7 @@ import pytest import search_docsite.search_docsite as m +from pgvector import Vector from util import ApolloError @@ -224,3 +225,30 @@ def test_search_maps_missing_docsite_tables_to_503(): assert exc.value.code == 503 assert "embed_docsite" in exc.value.message + + +# --- vector binding ---------------------------------------------------------- + +def test_semantic_search_binds_the_embedding_as_a_vector(): + """psycopg2 has no adapter for list, so a list is sent as numeric[] and + `vector <=> numeric[]` resolves to no operator. Only Vector and ndarray are + registered by pgvector.""" + conn, cur = make_conn() + cur.fetchall.return_value = [] + ds = make_search() + + ds._semantic_search(conn, batch_id=1, query="q", top_k=5, threshold=None, doc_title=None, docs_type=None) + + params = cur.execute.call_args[0][1] + assert params["query_embedding"] == Vector([0.1, 0.2, 0.3]) + + +def test_hybrid_search_binds_the_embedding_as_a_vector(): + conn, cur = make_conn() + cur.fetchall.return_value = [] + ds = make_search() + + ds._hybrid_search(conn, batch_id=1, query="q", top_k=5, doc_title=None, docs_type=None) + + params = cur.execute.call_args[0][1] + assert params["query_embedding"] == Vector([0.1, 0.2, 0.3]) From 0725262a4ef4b4c0d419e0f57a317690d3372cbc Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Wed, 29 Jul 2026 14:08:43 +0800 Subject: [PATCH 31/48] fix: commit before toggling autocommit, and mark failed batches --- services/embed_docsite/docsite_indexer.py | 44 +++++-- services/embed_docsite/embed_docsite.py | 13 +++ .../embed_docsite/tests/unit/fake_conn.py | 107 +++++++++++++++++ .../tests/unit/test_docsite_indexer.py | 110 ++++++++++++------ .../tests/unit/test_embed_docsite.py | 48 ++++++++ 5 files changed, 277 insertions(+), 45 deletions(-) create mode 100644 services/embed_docsite/tests/unit/fake_conn.py diff --git a/services/embed_docsite/docsite_indexer.py b/services/embed_docsite/docsite_indexer.py index 7562e403..bca24ba9 100644 --- a/services/embed_docsite/docsite_indexer.py +++ b/services/embed_docsite/docsite_indexer.py @@ -1,3 +1,5 @@ +from contextlib import contextmanager + from psycopg2.extras import execute_values from langchain_openai import OpenAIEmbeddings from pgvector import Vector @@ -9,6 +11,23 @@ ALL_DOCS_TYPES = ["adaptor_docs", "general_docs", "adaptor_functions"] +@contextmanager +def autocommit(conn): + """Run DDL that cannot execute inside a transaction (CREATE/DROP INDEX + CONCURRENTLY). + + Commits first: psycopg2 refuses to change autocommit while a transaction is + open, and a preceding SELECT is enough to open one. + """ + previous = conn.autocommit + conn.commit() + conn.autocommit = True + try: + yield + finally: + conn.autocommit = previous + + def register_vector_type(conn): """Register pgvector's psycopg2 adapters on this connection. @@ -126,9 +145,8 @@ def copy_forward_missing_docs_types(self, conn, batch_id, docs_types_present): return copied def build_index(self, conn, batch_id): - """Build a per-batch partial HNSW index. Runs outside a transaction (autocommit).""" - conn.autocommit = True - try: + """Build a per-batch partial HNSW index.""" + with autocommit(conn): with conn.cursor() as cur: cur.execute( f""" @@ -138,8 +156,6 @@ def build_index(self, conn, batch_id): WHERE batch_id = {batch_id} """ ) - finally: - conn.autocommit = False logger.info(f"Built HNSW index for batch {batch_id}") def promote_batch(self, conn, batch_id, chunk_count): @@ -152,6 +168,19 @@ def promote_batch(self, conn, batch_id, chunk_count): conn.commit() logger.info(f"Promoted batch {batch_id} ({chunk_count} chunks)") + def fail_batch(self, conn, batch_id): + """Mark a batch 'failed' after an aborted build. + + Rolls back first: the connection is in an aborted transaction from + whatever error got us here, so any statement would raise + InFailedSqlTransaction. + """ + conn.rollback() + with conn.cursor() as cur: + cur.execute("UPDATE docsite_batches SET status = 'failed' WHERE id = %s", (batch_id,)) + conn.commit() + logger.info(f"Marked batch {batch_id} failed") + def prune_old_batches(self, conn, keep_batches=None): """Delete complete batches older than the newest `keep_batches`, dropping their partial indexes first. Returns the list of pruned batch ids.""" @@ -165,12 +194,9 @@ def prune_old_batches(self, conn, keep_batches=None): old_batch_ids = [row[0] for row in cur.fetchall()] for batch_id in old_batch_ids: - conn.autocommit = True - try: + with autocommit(conn): with conn.cursor() as cur: cur.execute(f"DROP INDEX CONCURRENTLY IF EXISTS idx_docsite_chunks_hnsw_{batch_id}") - finally: - conn.autocommit = False with conn.cursor() as cur: cur.execute("DELETE FROM docsite_batches WHERE id = %s", (batch_id,)) conn.commit() diff --git a/services/embed_docsite/embed_docsite.py b/services/embed_docsite/embed_docsite.py index 8c5ba01b..0ee35c82 100644 --- a/services/embed_docsite/embed_docsite.py +++ b/services/embed_docsite/embed_docsite.py @@ -94,6 +94,14 @@ def _upload_to_pinecone(data, documents, metadata_dict, docs_to_upload): } +def _mark_failed(indexer, conn, batch_id): + """Best-effort bookkeeping: never let it replace the error that caused it.""" + try: + indexer.fail_batch(conn, batch_id) + except Exception as exc: + logger.error(f"Could not mark batch {batch_id} failed: {exc}") + + def _upload_to_postgres(documents, metadata_dict, docs_to_upload, chunk_target_length, chunk_min_length, keep_batches): indexer = DocsiteIndexer( chunk_target_length=chunk_target_length, @@ -107,6 +115,7 @@ def _upload_to_postgres(documents, metadata_dict, docs_to_upload, chunk_target_l run_migrations(conn) register_vector_type(conn) + batch_id = None try: batch_id = indexer.start_batch(conn, docs_to_upload) chunk_count = indexer.insert_documents(conn, batch_id, documents, metadata_dict) @@ -124,6 +133,10 @@ def _upload_to_postgres(documents, metadata_dict, docs_to_upload, chunk_target_l "pruned_batches": pruned, "promoted": True, } + except Exception: + if batch_id is not None: + _mark_failed(indexer, conn, batch_id) + raise finally: conn.close() diff --git a/services/embed_docsite/tests/unit/fake_conn.py b/services/embed_docsite/tests/unit/fake_conn.py new file mode 100644 index 00000000..5d921d98 --- /dev/null +++ b/services/embed_docsite/tests/unit/fake_conn.py @@ -0,0 +1,107 @@ +"""A psycopg2-shaped connection that models transaction state. + +MagicMock let two defects through: `conn.autocommit = True` on a mock is an +inert attribute write, but psycopg2 raises `set_session cannot be used inside a +transaction` when a transaction is open — and a bare SELECT is enough to open +one. This models the state machine those defects depend on. +""" + +import psycopg2 +from psycopg2.extensions import STATUS_IN_TRANSACTION, STATUS_READY + + +class FakeCursor: + """Cursor serving rows from its connection's queued results.""" + + def __init__(self, conn): + self._conn = conn + self._rows = [] + self.rowcount = -1 + + def __enter__(self): + return self + + def __exit__(self, *_exc): + return False + + def execute(self, sql, params=None): + self._conn._record(sql, params) + self._rows, self.rowcount = self._conn._pop_result() + + def fetchone(self): + return self._rows[0] if self._rows else None + + def fetchall(self): + return list(self._rows) + + def close(self): + pass + + +class FakeConn: + """psycopg2-shaped connection modelling the transaction state machine. + + :param results: FIFO consumed one entry per execute(). An entry may be a + list of rows, an int meaning "no rows, but report this rowcount" (as + INSERT and DELETE do), or None. An exhausted queue yields no rows. + :param fail_on: substring; an execute() whose SQL contains it raises and + leaves the connection aborted, as psycopg2 would. + """ + + def __init__(self, results=None, fail_on=None): + # _guard off while __init__ populates state, so the autocommit check + # below does not fire before `status` exists. + object.__setattr__(self, "_guard", False) + self._results = list(results or []) + self._fail_on = fail_on + self._failed = False + self.status = STATUS_READY + self.executed = [] + self.commits = 0 + self.rollbacks = 0 + self.closed = False + self.autocommit = False + object.__setattr__(self, "_guard", True) + + def __setattr__(self, name, value): + """Only `autocommit` is guarded — it is the assignment psycopg2 rejects.""" + if name == "autocommit" and self.__dict__.get("_guard") and self.status != STATUS_READY: + raise psycopg2.ProgrammingError("set_session cannot be used inside a transaction") + object.__setattr__(self, name, value) + + def cursor(self): + return FakeCursor(self) + + def commit(self): + self.commits += 1 + self.status = STATUS_READY + self._failed = False + + def rollback(self): + self.rollbacks += 1 + self.status = STATUS_READY + self._failed = False + + def close(self): + self.closed = True + + def _record(self, sql, params): + if self._failed: + raise psycopg2.errors.InFailedSqlTransaction( + "current transaction is aborted, commands ignored until end of transaction block" + ) + self.executed.append((sql, params)) + if self._fail_on and self._fail_on in sql: + self._failed = True + self.status = STATUS_IN_TRANSACTION + raise psycopg2.ProgrammingError(f"fake failure on: {self._fail_on}") + if not self.autocommit: + self.status = STATUS_IN_TRANSACTION + + def _pop_result(self): + """Returns (rows, rowcount) for the next queued entry.""" + entry = self._results.pop(0) if self._results else None + if isinstance(entry, int): + return [], entry + rows = list(entry or []) + return rows, len(rows) diff --git a/services/embed_docsite/tests/unit/test_docsite_indexer.py b/services/embed_docsite/tests/unit/test_docsite_indexer.py index db23b04a..8f2a203d 100644 --- a/services/embed_docsite/tests/unit/test_docsite_indexer.py +++ b/services/embed_docsite/tests/unit/test_docsite_indexer.py @@ -1,26 +1,24 @@ """Unit tests for the Postgres batch-lifecycle write path. Every DB call goes through an explicit `conn` parameter (never looked up -internally), so these tests use a MagicMock connection/cursor throughout — -no real Postgres needed. `register_vector` (which needs a live connection to -look up pgvector's type oid) and the OpenAI embeddings client are both -patched out. +internally), so these tests drive a FakeConn that models psycopg2's transaction +state. A MagicMock cannot, which is how the autocommit defects survived a green +suite. `register_vector` (which needs a live connection to look up pgvector's +type oid) and the OpenAI embeddings client are both patched out. """ from unittest.mock import MagicMock, patch -import embed_docsite.docsite_indexer as m +import psycopg2 +import pytest from pgvector import Vector - -def make_conn(): - conn = MagicMock() - cur = MagicMock() - conn.cursor.return_value.__enter__.return_value = cur - return conn, cur +import embed_docsite.docsite_indexer as m +from embed_docsite.tests.unit.fake_conn import FakeConn def test_register_vector_type_calls_pgvector_register(): + """Asserts a call, not transaction behaviour, so a MagicMock is right here.""" conn = MagicMock() with patch.object(m, "register_vector") as mock_register: m.register_vector_type(conn) @@ -34,20 +32,18 @@ def make_indexer(): def test_start_batch_inserts_row_and_returns_id(): - conn, cur = make_conn() - cur.fetchone.return_value = (7,) + conn = FakeConn(results=[[(7,)]]) indexer = make_indexer() batch_id = indexer.start_batch(conn, ["general_docs"]) assert batch_id == 7 - params = cur.execute.call_args[0][1] - assert params == (["general_docs"], 1000, 700, "fake-embedding-model") - conn.commit.assert_called_once() + assert conn.executed[0][1] == (["general_docs"], 1000, 700, "fake-embedding-model") + assert conn.commits == 1 def test_insert_documents_embeds_and_bulk_inserts(): - conn, cur = make_conn() + conn = FakeConn() indexer = make_indexer() indexer._embeddings.embed_documents.return_value = [[0.1, 0.2], [0.3, 0.4]] documents = [ @@ -63,45 +59,43 @@ def test_insert_documents_embeds_and_bulk_inserts(): rows = mock_execute_values.call_args[0][2] assert rows[0] == (7, "doc-a", "general_docs", 0, "chunk one", Vector([0.1, 0.2])) assert rows[1] == (7, "doc-a", "general_docs", 1, "chunk two", Vector([0.3, 0.4])) - conn.commit.assert_called_once() + assert conn.commits == 1 def test_insert_documents_returns_zero_for_empty_input(): - conn, _ = make_conn() + conn = FakeConn() indexer = make_indexer() count = indexer.insert_documents(conn, batch_id=7, documents=[], metadata_dict={}) assert count == 0 + assert conn.executed == [] indexer._embeddings.embed_documents.assert_not_called() def test_copy_forward_missing_docs_types_no_op_when_all_types_present(): - conn, cur = make_conn() + conn = FakeConn() indexer = make_indexer() copied = indexer.copy_forward_missing_docs_types(conn, batch_id=7, docs_types_present=m.ALL_DOCS_TYPES) assert copied == 0 - cur.execute.assert_not_called() + assert conn.executed == [] def test_copy_forward_missing_docs_types_copies_from_previous_batch(): - conn, cur = make_conn() - cur.fetchone.return_value = (3,) # previous complete batch id - cur.rowcount = 5 + # Queue: the SELECT finds batch 3, then the INSERT ... SELECT reports 5 rows. + conn = FakeConn(results=[[(3,)], 5]) indexer = make_indexer() copied = indexer.copy_forward_missing_docs_types(conn, batch_id=7, docs_types_present=["general_docs"]) assert copied == 5 - insert_call = cur.execute.call_args_list[-1] - assert insert_call[0][1] == (7, 3, ["adaptor_docs", "adaptor_functions"]) + assert conn.executed[-1][1] == (7, 3, ["adaptor_docs", "adaptor_functions"]) def test_copy_forward_missing_docs_types_returns_zero_when_no_previous_batch(): - conn, cur = make_conn() - cur.fetchone.return_value = None + conn = FakeConn(results=[None]) indexer = make_indexer() copied = indexer.copy_forward_missing_docs_types(conn, batch_id=7, docs_types_present=["general_docs"]) @@ -110,23 +104,67 @@ def test_copy_forward_missing_docs_types_returns_zero_when_no_previous_batch(): def test_promote_batch_updates_status_and_chunk_count(): - conn, cur = make_conn() + conn = FakeConn() indexer = make_indexer() indexer.promote_batch(conn, batch_id=7, chunk_count=42) - params = cur.execute.call_args[0][1] - assert params == (42, 7) - conn.commit.assert_called_once() + assert conn.executed[0][1] == (42, 7) + assert conn.commits == 1 def test_prune_old_batches_deletes_batches_beyond_keep_count(): - conn, cur = make_conn() - cur.fetchall.return_value = [(3,), (2,)] # older batches beyond keep_batches=2 + conn = FakeConn(results=[[(3,), (2,)]]) indexer = make_indexer() pruned = indexer.prune_old_batches(conn, keep_batches=2) assert pruned == [3, 2] - select_call = cur.execute.call_args_list[0] - assert select_call[0][1] == (2,) + assert conn.executed[0][1] == (2,) + + +def test_build_index_succeeds_when_a_prior_read_left_a_transaction_open(): + """C3: with no previous complete batch, copy_forward returns straight after + its SELECT without committing. psycopg2 then refuses `autocommit = True`, + build_index raises, and the batch strands in 'building' forever.""" + conn = FakeConn(results=[None]) + indexer = make_indexer() + indexer.copy_forward_missing_docs_types(conn, batch_id=7, docs_types_present=["general_docs"]) + + indexer.build_index(conn, batch_id=7) + + assert any("CREATE INDEX CONCURRENTLY" in sql for sql, _ in conn.executed) + assert conn.autocommit is False + + +def test_prune_old_batches_completes_after_its_own_select(): + """C2: the SELECT for old batch ids opens a transaction, so the DROP INDEX + step raised and pruning never ran — every re-index left its predecessor's + chunks and HNSW index behind.""" + conn = FakeConn(results=[[(3,), (2,)]]) + indexer = make_indexer() + + pruned = indexer.prune_old_batches(conn, keep_batches=2) + + assert pruned == [3, 2] + dropped = [sql for sql, _ in conn.executed if "DROP INDEX CONCURRENTLY" in sql] + assert len(dropped) == 2 + assert conn.autocommit is False + + +def test_fail_batch_recovers_an_aborted_transaction(): + """By the time fail_batch runs, the error that triggered it has already + aborted the transaction, so any statement would raise + InFailedSqlTransaction until something rolls back.""" + conn = FakeConn(fail_on="boom") + indexer = make_indexer() + with pytest.raises(psycopg2.ProgrammingError): + with conn.cursor() as cur: + cur.execute("boom") + + indexer.fail_batch(conn, batch_id=7) + + assert conn.rollbacks == 1 + sql, params = conn.executed[-1] + assert "status = 'failed'" in sql + assert params == (7,) diff --git a/services/embed_docsite/tests/unit/test_embed_docsite.py b/services/embed_docsite/tests/unit/test_embed_docsite.py index 9c72aeea..ade97ef6 100644 --- a/services/embed_docsite/tests/unit/test_embed_docsite.py +++ b/services/embed_docsite/tests/unit/test_embed_docsite.py @@ -127,3 +127,51 @@ def test_postgres_upload_runs_migrations_before_registering_vector_type(): m.main({"docs_to_upload": ["general_docs"], "target": "postgres"}) assert calls == ["migrate", "register"] + + +def test_postgres_upload_marks_batch_failed_and_reraises(): + """Schema allows status='failed' and nothing ever set it, so an interrupted + index run left a 'building' row that no later run could interpret.""" + import pytest + + fake_conn = MagicMock() + fake_indexer = MagicMock() + fake_indexer.start_batch.return_value = 7 + fake_indexer.insert_documents.side_effect = RuntimeError("embedding API down") + fake_processor = MagicMock() + fake_processor.get_preprocessed_docs.return_value = ([], {}) + + with patch.object(m, "get_db_connection", return_value=fake_conn), \ + patch.object(m, "run_migrations"), \ + patch.object(m, "register_vector_type"), \ + patch.object(m, "DocsiteProcessor", return_value=fake_processor), \ + patch.object(m, "DocsiteIndexer", return_value=fake_indexer), \ + patch.dict("os.environ", {"OPENAI_API_KEY": "sk-test"}): + with pytest.raises(RuntimeError, match="embedding API down"): + m.main({"docs_to_upload": ["general_docs"], "target": "postgres"}) + + fake_indexer.fail_batch.assert_called_once_with(fake_conn, 7) + fake_conn.close.assert_called_once() + + +def test_failed_marking_never_masks_the_original_error(): + """Bookkeeping must not become the reported failure — the operator needs to + see what actually broke.""" + import pytest + + fake_conn = MagicMock() + fake_indexer = MagicMock() + fake_indexer.start_batch.return_value = 7 + fake_indexer.insert_documents.side_effect = RuntimeError("embedding API down") + fake_indexer.fail_batch.side_effect = RuntimeError("connection already gone") + fake_processor = MagicMock() + fake_processor.get_preprocessed_docs.return_value = ([], {}) + + with patch.object(m, "get_db_connection", return_value=fake_conn), \ + patch.object(m, "run_migrations"), \ + patch.object(m, "register_vector_type"), \ + patch.object(m, "DocsiteProcessor", return_value=fake_processor), \ + patch.object(m, "DocsiteIndexer", return_value=fake_indexer), \ + patch.dict("os.environ", {"OPENAI_API_KEY": "sk-test"}): + with pytest.raises(RuntimeError, match="embedding API down"): + m.main({"docs_to_upload": ["general_docs"], "target": "postgres"}) From b5946bd910cdd7c35151c79e58e4aabd09f7e74a Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Wed, 29 Jul 2026 14:44:12 +0800 Subject: [PATCH 32/48] test: add Postgres docsite roundtrip integration suite --- .../tests/integration/__init__.py | 0 .../tests/integration/conftest.py | 38 ++++++++ .../tests/integration/helpers.py | 52 +++++++++++ .../test_postgres_docsite_roundtrip.py | 88 +++++++++++++++++++ 4 files changed, 178 insertions(+) create mode 100644 services/embed_docsite/tests/integration/__init__.py create mode 100644 services/embed_docsite/tests/integration/conftest.py create mode 100644 services/embed_docsite/tests/integration/helpers.py create mode 100644 services/embed_docsite/tests/integration/test_postgres_docsite_roundtrip.py diff --git a/services/embed_docsite/tests/integration/__init__.py b/services/embed_docsite/tests/integration/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/services/embed_docsite/tests/integration/conftest.py b/services/embed_docsite/tests/integration/conftest.py new file mode 100644 index 00000000..40d9c467 --- /dev/null +++ b/services/embed_docsite/tests/integration/conftest.py @@ -0,0 +1,38 @@ +"""Fixtures for the Postgres docsite integration suite. + +Requires POSTGRES_TEST_URL (not POSTGRES_URL) pointing at a database with the pgvector extension +available and a role permitted to CREATE EXTENSION: + + docker run -d --name apollo-pgvector-test -e POSTGRES_PASSWORD=postgres \ + -p 5433:5432 pgvector/pgvector:pg16 + export POSTGRES_TEST_URL=postgresql://postgres:postgres@localhost:5433/postgres + +The repo-root conftest blocks psycopg2.connect for `unit` tests only, so this +tier connects normally. +""" + +import psycopg2 +import pytest + +from embed_docsite.tests.integration.helpers import TEST_URL + + +@pytest.fixture +def clean_db(monkeypatch): + """An empty database: no docsite tables, no pgvector extension. + + Dropping the extension too means the reader's 'no extension' path is + reachable, and migrations have to prove they can recreate it. + """ + if not TEST_URL: + pytest.skip("POSTGRES_TEST_URL not set") + + conn = psycopg2.connect(TEST_URL) + conn.autocommit = True + with conn.cursor() as cur: + cur.execute("DROP TABLE IF EXISTS docsite_chunks, docsite_batches, _migrations_docs CASCADE") + cur.execute("DROP EXTENSION IF EXISTS vector CASCADE") + conn.close() + + monkeypatch.setenv("POSTGRES_URL", TEST_URL) + return TEST_URL diff --git a/services/embed_docsite/tests/integration/helpers.py b/services/embed_docsite/tests/integration/helpers.py new file mode 100644 index 00000000..02f207e5 --- /dev/null +++ b/services/embed_docsite/tests/integration/helpers.py @@ -0,0 +1,52 @@ +"""Shared helpers for the Postgres docsite integration suite.""" + +import hashlib +import os +import random + +import psycopg2 + +# Importing embed_docsite pulls in LegacyPineconeDocsiteIndexer, whose +# OpenAIEmbeddings() default arg validates credentials at construction. Dummy +# placeholders only — this suite makes no OpenAI call. +os.environ.setdefault("OPENAI_API_KEY", "sk-test-dummy") +os.environ.setdefault("PINECONE_API_KEY", "pc-test-dummy") + +TEST_URL = os.environ.get("POSTGRES_TEST_URL") + +EMBEDDING_DIMENSIONS = 1536 + + +class StubEmbeddings: + """Deterministic stand-in for OpenAIEmbeddings. + + Identical text yields an identical vector, so querying a chunk's exact text + puts that chunk at cosine distance 0 and therefore rank 1. Stubbing costs no + coverage here: `operator does not exist: vector <=> numeric[]` is raised + from the bound parameter's type, never its content. + """ + + model = "stub-embedding-model" + + def embed_documents(self, texts): + return [self._vector(text) for text in texts] + + def embed_query(self, text): + return self._vector(text) + + @staticmethod + def _vector(text): + seed = int.from_bytes(hashlib.sha256(text.encode("utf-8")).digest()[:8], "big") + rng = random.Random(seed) + return [rng.uniform(-1.0, 1.0) for _ in range(EMBEDDING_DIMENSIONS)] + + +def query(sql, params=None): + """Run a read against the test database on its own connection.""" + conn = psycopg2.connect(TEST_URL) + try: + with conn.cursor() as cur: + cur.execute(sql, params) + return cur.fetchall() + finally: + conn.close() diff --git a/services/embed_docsite/tests/integration/test_postgres_docsite_roundtrip.py b/services/embed_docsite/tests/integration/test_postgres_docsite_roundtrip.py new file mode 100644 index 00000000..c0f41a2c --- /dev/null +++ b/services/embed_docsite/tests/integration/test_postgres_docsite_roundtrip.py @@ -0,0 +1,88 @@ +"""End-to-end index-then-search against a real Postgres with pgvector. +""" + +from unittest.mock import patch + +import pytest + +import embed_docsite.docsite_indexer as indexer_module +from embed_docsite.embed_docsite import _upload_to_postgres +from embed_docsite.tests.integration.helpers import StubEmbeddings, query +from search_docsite.search_docsite import DocsiteSearch +from util import ApolloError + +DOCS = [ + { + "name": "webhook-guide.md", + "docs_type": "general_docs", + "doc_chunk": "Configure a webhook trigger to start a workflow when data arrives.", + }, + { + "name": "cron-guide.md", + "docs_type": "general_docs", + "doc_chunk": "Use a cron trigger to run a workflow on a fixed schedule.", + }, +] + + +def index_docs(keep_batches=2): + """Run the real Postgres upload path with stubbed embeddings.""" + with patch.object(indexer_module, "OpenAIEmbeddings", StubEmbeddings): + return _upload_to_postgres( + documents=[dict(doc) for doc in DOCS], + metadata_dict={}, + docs_to_upload=["general_docs"], + chunk_target_length=1000, + chunk_min_length=700, + keep_batches=keep_batches, + ) + + +def make_search(**kwargs): + search = DocsiteSearch(**kwargs) + search._embeddings = StubEmbeddings() + return search + + +def test_fresh_database_migrates_indexes_and_promotes(clean_db): + """C3: the first run on an empty database used to strand in 'building', + because copy_forward's SELECT left a transaction open.""" + result = index_docs() + + assert result["promoted"] is True + rows = query("SELECT status FROM docsite_batches WHERE id = %s", (result["batch_id"],)) + assert rows[0][0] == "complete" + + +@pytest.mark.parametrize("strategy", ["semantic", "keyword", "hybrid"]) +def test_search_returns_the_indexed_chunk(clean_db, strategy): + """C1: semantic and hybrid used to fail with + `operator does not exist: vector <=> numeric[]` on every query.""" + index_docs() + target = DOCS[0]["doc_chunk"] + + results = make_search().search(target, strategy=strategy, top_k=3) + + assert any(result.text == target for result in results) + + +def test_reindexing_prunes_the_previous_batch(clean_db): + """C2: pruning never ran, so every re-index permanently added a full + docsite copy and another HNSW index.""" + first = index_docs(keep_batches=1) + second = index_docs(keep_batches=1) + + assert first["batch_id"] in second["pruned_batches"] + assert query("SELECT count(*) FROM docsite_batches WHERE id = %s", (first["batch_id"],))[0][0] == 0 + index_name = f"idx_docsite_chunks_hnsw_{first['batch_id']}" + assert query("SELECT count(*) FROM pg_class WHERE relname = %s", (index_name,))[0][0] == 0 + + +def test_reader_without_a_schema_gets_a_clear_503(clean_db): + """C4: with migrations moved to the indexer, a reader on an un-indexed + database must explain itself rather than emit a psycopg2 traceback.""" + with pytest.raises(ApolloError) as exc: + make_search().search("anything") + + assert exc.value.code == 503 + assert "embed_docsite" in exc.value.message From 947e585cda460b7f9a7b7adb07a57cf3b0b20e50 Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Wed, 29 Jul 2026 15:43:08 +0800 Subject: [PATCH 33/48] fix: don't let a pruning failure invalidate the batch just promoted --- services/embed_docsite/embed_docsite.py | 48 ++++++++++++------- .../tests/unit/test_embed_docsite.py | 36 ++++++++++++++ 2 files changed, 68 insertions(+), 16 deletions(-) diff --git a/services/embed_docsite/embed_docsite.py b/services/embed_docsite/embed_docsite.py index 0ee35c82..d09896b3 100644 --- a/services/embed_docsite/embed_docsite.py +++ b/services/embed_docsite/embed_docsite.py @@ -102,6 +102,17 @@ def _mark_failed(indexer, conn, batch_id): logger.error(f"Could not mark batch {batch_id} failed: {exc}") +def _prune_old_batches(indexer, conn): + """Best-effort cleanup of OLDER, unrelated batches. Runs after the new + batch is already promoted, so a failure here must never retroactively + invalidate that batch or fail the whole call.""" + try: + return indexer.prune_old_batches(conn) + except Exception as exc: + logger.error(f"Could not prune old batches: {exc}") + return [] + + def _upload_to_postgres(documents, metadata_dict, docs_to_upload, chunk_target_length, chunk_min_length, keep_batches): indexer = DocsiteIndexer( chunk_target_length=chunk_target_length, @@ -110,19 +121,28 @@ def _upload_to_postgres(documents, metadata_dict, docs_to_upload, chunk_target_l ) conn = get_db_connection() - # Order matters: register_vector looks up pgvector's type oid, so the - # extension must exist before we register. - run_migrations(conn) - register_vector_type(conn) - - batch_id = None try: - batch_id = indexer.start_batch(conn, docs_to_upload) - chunk_count = indexer.insert_documents(conn, batch_id, documents, metadata_dict) - copied = indexer.copy_forward_missing_docs_types(conn, batch_id, docs_to_upload) - indexer.build_index(conn, batch_id) - indexer.promote_batch(conn, batch_id, chunk_count + copied) - pruned = indexer.prune_old_batches(conn) + # Order matters: register_vector looks up pgvector's type oid, so the + # extension must exist before we register. + run_migrations(conn) + register_vector_type(conn) + + batch_id = None + try: + batch_id = indexer.start_batch(conn, docs_to_upload) + chunk_count = indexer.insert_documents(conn, batch_id, documents, metadata_dict) + copied = indexer.copy_forward_missing_docs_types(conn, batch_id, docs_to_upload) + indexer.build_index(conn, batch_id) + indexer.promote_batch(conn, batch_id, chunk_count + copied) + except Exception: + if batch_id is not None: + _mark_failed(indexer, conn, batch_id) + raise + + # The new batch is already promoted and visible to readers. Pruning + # only touches older, unrelated batches, so its failure must not be + # attributed to the batch we just built. + pruned = _prune_old_batches(indexer, conn) return { "target": "postgres", @@ -133,10 +153,6 @@ def _upload_to_postgres(documents, metadata_dict, docs_to_upload, chunk_target_l "pruned_batches": pruned, "promoted": True, } - except Exception: - if batch_id is not None: - _mark_failed(indexer, conn, batch_id) - raise finally: conn.close() diff --git a/services/embed_docsite/tests/unit/test_embed_docsite.py b/services/embed_docsite/tests/unit/test_embed_docsite.py index ade97ef6..bd890df5 100644 --- a/services/embed_docsite/tests/unit/test_embed_docsite.py +++ b/services/embed_docsite/tests/unit/test_embed_docsite.py @@ -154,6 +154,42 @@ def test_postgres_upload_marks_batch_failed_and_reraises(): fake_conn.close.assert_called_once() +def test_prune_failure_after_promote_does_not_mark_batch_failed(): + """Pruning cleans up OLDER, unrelated batches. If it throws after the new + batch was already promoted, the new batch must stay 'complete' — it must + not be retroactively marked 'failed', and the call must still succeed.""" + fake_conn = MagicMock() + fake_indexer = MagicMock() + fake_indexer.start_batch.return_value = 7 + fake_indexer.insert_documents.return_value = 10 + fake_indexer.copy_forward_missing_docs_types.return_value = 3 + fake_indexer.prune_old_batches.side_effect = RuntimeError("lock conflict on DROP INDEX CONCURRENTLY") + fake_processor = MagicMock() + fake_processor.get_preprocessed_docs.return_value = ([{"name": "a.md", "docs_type": "general_docs", "doc_chunk": "x"}], {"a.md": {}}) + + with patch.object(m, "get_db_connection", return_value=fake_conn), \ + patch.object(m, "run_migrations"), \ + patch.object(m, "register_vector_type"), \ + patch.object(m, "DocsiteProcessor", return_value=fake_processor), \ + patch.object(m, "DocsiteIndexer", return_value=fake_indexer), \ + patch.dict("os.environ", {"OPENAI_API_KEY": "sk-test"}): + result = m.main({"docs_to_upload": ["general_docs"], "target": "postgres"}) + + fake_indexer.promote_batch.assert_called_once_with(fake_conn, 7, 13) + fake_indexer.fail_batch.assert_not_called() + fake_conn.close.assert_called_once() + + assert result == { + "target": "postgres", + "batch_id": 7, + "docs_types": ["general_docs"], + "chunk_count": 10, + "copied_forward": 3, + "pruned_batches": [], + "promoted": True, + } + + def test_failed_marking_never_masks_the_original_error(): """Bookkeeping must not become the reported failure — the operator needs to see what actually broke.""" From a303af8dbf002f98f8b583440216f3a4ed69c6a3 Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Wed, 29 Jul 2026 17:08:59 +0800 Subject: [PATCH 34/48] fix: change integration test to use 127.0.0.1 instead of localhost to speedup --- services/embed_docsite/tests/integration/conftest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/services/embed_docsite/tests/integration/conftest.py b/services/embed_docsite/tests/integration/conftest.py index 40d9c467..b45fe9e5 100644 --- a/services/embed_docsite/tests/integration/conftest.py +++ b/services/embed_docsite/tests/integration/conftest.py @@ -5,7 +5,7 @@ docker run -d --name apollo-pgvector-test -e POSTGRES_PASSWORD=postgres \ -p 5433:5432 pgvector/pgvector:pg16 - export POSTGRES_TEST_URL=postgresql://postgres:postgres@localhost:5433/postgres + export POSTGRES_TEST_URL=postgresql://postgres:postgres@127.0.0.1:5433/postgres The repo-root conftest blocks psycopg2.connect for `unit` tests only, so this tier connects normally. From 74889bf36940c3d8d088100e0c7c9e76b4dc16b6 Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Wed, 29 Jul 2026 22:41:35 +0800 Subject: [PATCH 35/48] fix: update env example to support docsite backend and postgres test url for integration testing --- .env.example | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/.env.example b/.env.example index 67f1bd48..2ee155f0 100644 --- a/.env.example +++ b/.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-... From 483de4e4ff1ef0847577587f6cfb64327b4eebde Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Wed, 29 Jul 2026 23:02:15 +0800 Subject: [PATCH 36/48] fix: run_eval compares postgres semantic, and loads dotenv --- .../search_docsite/pinecone_legacy_search.py | 5 ++++- services/search_docsite/tests/eval/run_eval.py | 18 ++++++++++++------ 2 files changed, 16 insertions(+), 7 deletions(-) diff --git a/services/search_docsite/pinecone_legacy_search.py b/services/search_docsite/pinecone_legacy_search.py index b1b0016b..11676dcd 100644 --- a/services/search_docsite/pinecone_legacy_search.py +++ b/services/search_docsite/pinecone_legacy_search.py @@ -93,8 +93,11 @@ def _get_most_recent_namespace(self): index_stats = index.describe_index_stats() namespaces = index_stats.get('namespaces', {}).keys() + # The indexer names namespaces docsite-%Y%m%d%H%M (20 chars); namespaces + # created by hand use docsite-%Y%m%d (16). Accept both, so a namespace the + # indexer just wrote is discoverable. valid_namespaces = sorted( - (ns for ns in namespaces if ns.startswith("docsite-") and ns[8:].isdigit() and len(ns) == 16), + (ns for ns in namespaces if ns.startswith("docsite-") and ns[8:].isdigit() and len(ns) in (16, 20)), reverse=True ) diff --git a/services/search_docsite/tests/eval/run_eval.py b/services/search_docsite/tests/eval/run_eval.py index 9044a19f..4e8704a0 100644 --- a/services/search_docsite/tests/eval/run_eval.py +++ b/services/search_docsite/tests/eval/run_eval.py @@ -2,16 +2,22 @@ Usage: poetry run python -m search_docsite.tests.eval.run_eval -Compares the Postgres-backed DocsiteSearch (strategy='hybrid') against -LegacyPineconeDocsiteSearch (strategy='semantic') over the golden query set. -Postgres should match or beat Pinecone's recall@5 and p95 latency before -DOCSITE_SEARCH_BACKEND is flipped to 'postgres' by default. +Compares the Postgres-backed DocsiteSearch against LegacyPineconeDocsiteSearch +over the golden query set. Both run strategy='semantic'. Postgres should match +or beat Pinecone's recall@5 and p95 latency before DOCSITE_SEARCH_BACKEND is +flipped to 'postgres' by default. """ import time from pathlib import Path import yaml +from dotenv import load_dotenv + +# Needed as LegacyPineconeDocsiteSearch evaluates OpenAIEmbeddings() as a +# default argument. +load_dotenv() + from search_docsite.pinecone_legacy_search import LegacyPineconeDocsiteSearch from search_docsite.search_docsite import DocsiteSearch @@ -101,10 +107,10 @@ def main(): with open(GOLDEN_QUERIES_PATH) as f: golden_queries = yaml.safe_load(f)["queries"] - postgres_report = run_eval(golden_queries, DocsiteSearch, strategy="hybrid") + postgres_report = run_eval(golden_queries, DocsiteSearch, strategy="semantic") pinecone_report = run_eval(golden_queries, LegacyPineconeDocsiteSearch, strategy="semantic") - print(f"Postgres (hybrid): recall@5={postgres_report['recall_at_k']} " + print(f"Postgres (semantic): recall@5={postgres_report['recall_at_k']} " f"(scored={postgres_report['queries_scored']}, skipped={postgres_report['queries_skipped']}) " f"p50={postgres_report['p50_latency_s']:.3f}s p95={postgres_report['p95_latency_s']:.3f}s") print(f"Pinecone (semantic): recall@5={pinecone_report['recall_at_k']} " From 590b80853d2e6495e58ceacc9bee13a12f53a11b Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Thu, 30 Jul 2026 08:56:53 +0800 Subject: [PATCH 37/48] Add results from running golden queries --- .../tests/eval/golden_queries.yaml | 34 +++++++++++-------- 1 file changed, 19 insertions(+), 15 deletions(-) diff --git a/services/search_docsite/tests/eval/golden_queries.yaml b/services/search_docsite/tests/eval/golden_queries.yaml index 812849d6..3a931d51 100644 --- a/services/search_docsite/tests/eval/golden_queries.yaml +++ b/services/search_docsite/tests/eval/golden_queries.yaml @@ -1,27 +1,31 @@ -# Starter set of representative job_chat/general-docs queries. Queries with an -# empty expected_doc_titles list are still run and reported (so this fixture -# is usable immediately), but excluded from the recall@k aggregate until a -# human curates the expected doc title(s) by inspecting real search results -# against the live docsite content (neither this script's author nor the -# implementer of this plan has that access). +# Representative job_chat/general-docs queries with expected doc titles. +# +# recall@k counts a query as a hit when ANY expected title appears in the top k, +# so each entry lists every doc that would be a reasonable answer, not only the +# single best one. Titles are docsite filenames without .md, matching +# docsite_chunks.doc_title. queries: - query: "how do I configure a webhook trigger" - expected_doc_titles: [] + expected_doc_titles: ["triggers", "webhook-auth"] - query: "how do I set up a cron trigger" - expected_doc_titles: [] + expected_doc_titles: ["triggers"] - query: "what is a run in OpenFn" - expected_doc_titles: [] + expected_doc_titles: ["terminology", "glossary"] - query: "how do I use collections to store state between runs" - expected_doc_titles: [] + expected_doc_titles: ["collections", "state", "cli-collections"] - query: "how do I configure a credential for an adaptor" - expected_doc_titles: [] + expected_doc_titles: ["credentials", "manage-credentials"] - query: "what is the difference between a job and a workflow" - expected_doc_titles: [] + expected_doc_titles: ["terminology", "glossary", "workflows"] - query: "how do I deploy a project using the CLI" - expected_doc_titles: [] + expected_doc_titles: ["cli-sync", "portability"] - query: "how do I debug a failed run" - expected_doc_titles: [] + expected_doc_titles: + ["troubleshooting", "rerunning-workflow", "inspect-runs"] + # Unscorable as run_eval filters to docs_type='general_docs', but the + # adaptor listing lives in adaptor_docs. - query: "what adaptors are available for HTTP requests" expected_doc_titles: [] - query: "how do I write a data transform function" - expected_doc_titles: [] + expected_doc_titles: + ["data-transformation", "job-writing-guide", "operations"] From 9fb9f1172c08c4478410579152aceef0bd0a933d Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Thu, 30 Jul 2026 09:44:57 +0800 Subject: [PATCH 38/48] fix: import ordering and streamline comments --- services/embed_docsite/README.md | 22 +++++++++++++------ services/embed_docsite/docsite_indexer.py | 16 +++++--------- services/embed_docsite/docsite_processor.py | 3 ++- services/embed_docsite/embed_docsite.py | 12 ++++------ .../embed_docsite/pinecone_legacy_indexer.py | 9 ++++---- .../tests/integration/conftest.py | 1 - .../test_postgres_docsite_roundtrip.py | 3 +-- services/embed_docsite/tests/unit/conftest.py | 3 +-- .../tests/unit/test_db_migrations.py | 12 +++++----- .../tests/unit/test_docsite_indexer.py | 18 +++------------ .../tests/unit/test_embed_docsite.py | 4 ++-- services/job_chat/retrieve_docs.py | 16 ++++++++------ .../integration/test_adaptor_docs_pipeline.py | 5 ++--- services/job_chat/tests/unit/conftest.py | 3 +-- .../job_chat/tests/unit/test_retrieve_docs.py | 2 -- .../search_docsite/pinecone_legacy_search.py | 9 ++++---- services/search_docsite/search_docsite.py | 13 +++-------- services/search_docsite/tests/conftest.py | 4 +--- .../tests/unit/test_docsite_search.py | 5 ++--- .../tests/unit/test_pinecone_legacy_search.py | 1 - .../tests/unit/test_search_docsite_main.py | 1 - .../search_documentation.py | 9 +++----- services/util.py | 6 ----- 23 files changed, 69 insertions(+), 108 deletions(-) diff --git a/services/embed_docsite/README.md b/services/embed_docsite/README.md index b3c2530b..71512072 100644 --- a/services/embed_docsite/README.md +++ b/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,6 +18,7 @@ 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 @@ -22,7 +26,11 @@ 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`. @@ -36,15 +44,15 @@ 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"], - "docs_to_ignore": ["job-examples.md", "release-notes.md"], + "docs_to_upload": ["adaptor_docs", "general_docs", "adaptor_functions"], // Select from 3 types of documentation to upload + "docs_to_ignore": ["job-examples.md", "release-notes.md"], // Titles of documents that should not be indexed "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", - "max_total_collections": 3, + "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 diff --git a/services/embed_docsite/docsite_indexer.py b/services/embed_docsite/docsite_indexer.py index bca24ba9..c16ea57e 100644 --- a/services/embed_docsite/docsite_indexer.py +++ b/services/embed_docsite/docsite_indexer.py @@ -1,9 +1,9 @@ from contextlib import contextmanager -from psycopg2.extras import execute_values from langchain_openai import OpenAIEmbeddings from pgvector import Vector from pgvector.psycopg2 import register_vector +from psycopg2.extras import execute_values from util import create_logger logger = create_logger("DocsiteIndexer") @@ -29,12 +29,7 @@ def autocommit(conn): def register_vector_type(conn): - """Register pgvector's psycopg2 adapters on this connection. - - Registers `Vector` and `numpy.ndarray` — *not* `list`. A bare list is sent - as `numeric[]`, which Postgres will cast on assignment but has no `<=>` - operator for, so writes succeed and searches fail. - """ + """Register pgvector's psycopg2 adapters on this connection.""" register_vector(conn) @@ -59,8 +54,7 @@ def __init__(self, chunk_target_length=1000, chunk_min_length=700, keep_batches= @property def embeddings(self): - """Lazily construct the OpenAI embeddings client, so importing this module - needs no credentials.""" + """Lazily construct the OpenAI embeddings client.""" if self._embeddings is None: self._embeddings = OpenAIEmbeddings() return self._embeddings @@ -115,7 +109,7 @@ def _embed_in_batches(self, texts, batch_size=100): return embeddings def copy_forward_missing_docs_types(self, conn, batch_id, docs_types_present): - """Copy chunks for docs_types NOT in this run from the previous complete batch, + """Copy chunks for docs_types not in this run from the previous complete batch, so every complete batch is a full snapshot across all docs_types. Returns rows copied.""" missing_types = [t for t in ALL_DOCS_TYPES if t not in docs_types_present] if not missing_types: @@ -159,7 +153,7 @@ def build_index(self, conn, batch_id): logger.info(f"Built HNSW index for batch {batch_id}") def promote_batch(self, conn, batch_id, chunk_count): - """Flip a batch to 'complete' — the moment it becomes visible to readers.""" + """Flip a batch to 'complete' the moment it becomes visible to readers.""" with conn.cursor() as cur: cur.execute( "UPDATE docsite_batches SET status = 'complete', completed_at = now(), chunk_count = %s WHERE id = %s", diff --git a/services/embed_docsite/docsite_processor.py b/services/embed_docsite/docsite_processor.py index c49f74a1..5a14e395 100644 --- a/services/embed_docsite/docsite_processor.py +++ b/services/embed_docsite/docsite_processor.py @@ -1,6 +1,7 @@ import json import os import re + import nltk from embed_docsite.github_utils import get_docs from util import create_logger @@ -139,7 +140,7 @@ def _accumulate_chunks(self, splits, target_length, overlap, min_length): if len(current_chunk) >= min_length: accumulated.append(current_chunk) # Store the completed chunk - # add overlap + # Add overlap if self.docs_type == "adaptor_functions": overlap_sections = " ".join(current_chunk.split("\n")[-overlap:]) else: diff --git a/services/embed_docsite/embed_docsite.py b/services/embed_docsite/embed_docsite.py index d09896b3..6e03c9c3 100644 --- a/services/embed_docsite/embed_docsite.py +++ b/services/embed_docsite/embed_docsite.py @@ -71,7 +71,7 @@ def main(data: dict) -> dict: def _upload_to_pinecone(data, documents, metadata_dict, docs_to_upload): - """Legacy write path. Deliberately opens no Postgres connection.""" + """Legacy write path to pinecone.""" pinecone_api_key = data.get("PINECONE_API_KEY") or os.environ.get("PINECONE_API_KEY") if not pinecone_api_key: msg = "Missing API key: PINECONE_API_KEY" @@ -103,9 +103,8 @@ def _mark_failed(indexer, conn, batch_id): def _prune_old_batches(indexer, conn): - """Best-effort cleanup of OLDER, unrelated batches. Runs after the new - batch is already promoted, so a failure here must never retroactively - invalidate that batch or fail the whole call.""" + """Best-effort cleanup of older, unrelated batches. Runs after the new + batch is already promoted.""" try: return indexer.prune_old_batches(conn) except Exception as exc: @@ -122,8 +121,6 @@ def _upload_to_postgres(documents, metadata_dict, docs_to_upload, chunk_target_l conn = get_db_connection() try: - # Order matters: register_vector looks up pgvector's type oid, so the - # extension must exist before we register. run_migrations(conn) register_vector_type(conn) @@ -140,8 +137,7 @@ def _upload_to_postgres(documents, metadata_dict, docs_to_upload, chunk_target_l raise # The new batch is already promoted and visible to readers. Pruning - # only touches older, unrelated batches, so its failure must not be - # attributed to the batch we just built. + # only touches older, unrelated batches. pruned = _prune_old_batches(indexer, conn) return { diff --git a/services/embed_docsite/pinecone_legacy_indexer.py b/services/embed_docsite/pinecone_legacy_indexer.py index 4094c79f..263991c7 100644 --- a/services/embed_docsite/pinecone_legacy_indexer.py +++ b/services/embed_docsite/pinecone_legacy_indexer.py @@ -1,12 +1,13 @@ import os import time from datetime import datetime + import pandas as pd -from pinecone import Pinecone, ServerlessSpec -from langchain_pinecone import PineconeVectorStore -from langchain_openai import OpenAIEmbeddings from langchain_community.document_loaders import DataFrameLoader -from util import create_logger, ApolloError +from langchain_openai import OpenAIEmbeddings +from langchain_pinecone import PineconeVectorStore +from pinecone import Pinecone, ServerlessSpec +from util import ApolloError, create_logger logger = create_logger("LegacyPineconeDocsiteIndexer") diff --git a/services/embed_docsite/tests/integration/conftest.py b/services/embed_docsite/tests/integration/conftest.py index b45fe9e5..a9551513 100644 --- a/services/embed_docsite/tests/integration/conftest.py +++ b/services/embed_docsite/tests/integration/conftest.py @@ -13,7 +13,6 @@ import psycopg2 import pytest - from embed_docsite.tests.integration.helpers import TEST_URL diff --git a/services/embed_docsite/tests/integration/test_postgres_docsite_roundtrip.py b/services/embed_docsite/tests/integration/test_postgres_docsite_roundtrip.py index c0f41a2c..4096d219 100644 --- a/services/embed_docsite/tests/integration/test_postgres_docsite_roundtrip.py +++ b/services/embed_docsite/tests/integration/test_postgres_docsite_roundtrip.py @@ -3,9 +3,8 @@ from unittest.mock import patch -import pytest - import embed_docsite.docsite_indexer as indexer_module +import pytest from embed_docsite.embed_docsite import _upload_to_postgres from embed_docsite.tests.integration.helpers import StubEmbeddings, query from search_docsite.search_docsite import DocsiteSearch diff --git a/services/embed_docsite/tests/unit/conftest.py b/services/embed_docsite/tests/unit/conftest.py index bb48d19b..1d3871be 100644 --- a/services/embed_docsite/tests/unit/conftest.py +++ b/services/embed_docsite/tests/unit/conftest.py @@ -5,8 +5,7 @@ default arg validates credentials at construction (openai 2.x / langchain-openai 1.x). A key must therefore exist at import time. -Dummy placeholders only: unit tests mock every real network call, so no real -key is ever used. `setdefault` means a real key (from services/.env) wins. +This test only sets dummy environment variables for unit tests. """ import os diff --git a/services/embed_docsite/tests/unit/test_db_migrations.py b/services/embed_docsite/tests/unit/test_db_migrations.py index 6468434a..682a78a8 100644 --- a/services/embed_docsite/tests/unit/test_db_migrations.py +++ b/services/embed_docsite/tests/unit/test_db_migrations.py @@ -1,9 +1,8 @@ """Unit tests for the Python migration runner. -Mirrors the TypeScript runner in platform/src/db/migrate.ts: lexical ordering, -already-applied files skipped, advisory lock taken. The connection and cursor -are MagicMocks — the repo-root conftest blocks real psycopg2.connect in unit -tests anyway. +Mirrors the TypeScript runner in platform/src/db/migrate.ts using lexical ordering, +already-applied files skipped and an advisory lock taken. The connection and cursor +are MagicMocks. """ from unittest.mock import MagicMock, patch @@ -43,9 +42,8 @@ def test_run_migrations_creates_tracking_table(): def test_migration_files_returns_sql_files_in_lexical_order(tmp_path): """The sort lives in _migration_files, so it must be tested against the real - filesystem — patching that function out (as the apply tests below do) would - bypass the very ordering being asserted. Files are created out of order and - a non-.sql file is included to prove it is filtered.""" + filesystem. Files are created out of order and a non-.sql file is included + to prove it is filtered.""" (tmp_path / "0002_second.sql").write_text("SELECT 2;", encoding="utf-8") (tmp_path / "0010_tenth.sql").write_text("SELECT 10;", encoding="utf-8") (tmp_path / "0001_first.sql").write_text("SELECT 1;", encoding="utf-8") diff --git a/services/embed_docsite/tests/unit/test_docsite_indexer.py b/services/embed_docsite/tests/unit/test_docsite_indexer.py index 8f2a203d..76ca2400 100644 --- a/services/embed_docsite/tests/unit/test_docsite_indexer.py +++ b/services/embed_docsite/tests/unit/test_docsite_indexer.py @@ -2,19 +2,16 @@ Every DB call goes through an explicit `conn` parameter (never looked up internally), so these tests drive a FakeConn that models psycopg2's transaction -state. A MagicMock cannot, which is how the autocommit defects survived a green -suite. `register_vector` (which needs a live connection to look up pgvector's -type oid) and the OpenAI embeddings client are both patched out. +state. """ from unittest.mock import MagicMock, patch +import embed_docsite.docsite_indexer as m import psycopg2 import pytest -from pgvector import Vector - -import embed_docsite.docsite_indexer as m from embed_docsite.tests.unit.fake_conn import FakeConn +from pgvector import Vector def test_register_vector_type_calls_pgvector_register(): @@ -124,9 +121,6 @@ def test_prune_old_batches_deletes_batches_beyond_keep_count(): def test_build_index_succeeds_when_a_prior_read_left_a_transaction_open(): - """C3: with no previous complete batch, copy_forward returns straight after - its SELECT without committing. psycopg2 then refuses `autocommit = True`, - build_index raises, and the batch strands in 'building' forever.""" conn = FakeConn(results=[None]) indexer = make_indexer() indexer.copy_forward_missing_docs_types(conn, batch_id=7, docs_types_present=["general_docs"]) @@ -138,9 +132,6 @@ def test_build_index_succeeds_when_a_prior_read_left_a_transaction_open(): def test_prune_old_batches_completes_after_its_own_select(): - """C2: the SELECT for old batch ids opens a transaction, so the DROP INDEX - step raised and pruning never ran — every re-index left its predecessor's - chunks and HNSW index behind.""" conn = FakeConn(results=[[(3,), (2,)]]) indexer = make_indexer() @@ -153,9 +144,6 @@ def test_prune_old_batches_completes_after_its_own_select(): def test_fail_batch_recovers_an_aborted_transaction(): - """By the time fail_batch runs, the error that triggered it has already - aborted the transaction, so any statement would raise - InFailedSqlTransaction until something rolls back.""" conn = FakeConn(fail_on="boom") indexer = make_indexer() with pytest.raises(psycopg2.ProgrammingError): diff --git a/services/embed_docsite/tests/unit/test_embed_docsite.py b/services/embed_docsite/tests/unit/test_embed_docsite.py index bd890df5..9651ff90 100644 --- a/services/embed_docsite/tests/unit/test_embed_docsite.py +++ b/services/embed_docsite/tests/unit/test_embed_docsite.py @@ -1,5 +1,5 @@ """Unit tests for embed_docsite's orchestration. DocsiteProcessor/DocsiteIndexer -and get_db_connection are all mocked — this only tests call order and wiring.""" +and get_db_connection are all mocked.""" from unittest.mock import MagicMock, patch @@ -95,8 +95,8 @@ def test_main_pinecone_target_does_not_require_postgres_url(): def test_main_rejects_unknown_target(): - from util import ApolloError import pytest + from util import ApolloError with patch.dict("os.environ", {"OPENAI_API_KEY": "sk-test"}): with pytest.raises(ApolloError) as exc: diff --git a/services/job_chat/retrieve_docs.py b/services/job_chat/retrieve_docs.py index 523310b8..0ec9ced1 100644 --- a/services/job_chat/retrieve_docs.py +++ b/services/job_chat/retrieve_docs.py @@ -1,21 +1,23 @@ -import os import json +import os + import anthropic +import sentry_sdk from anthropic import ( APIConnectionError, - BadRequestError, AuthenticationError, - PermissionDeniedError, + BadRequestError, + InternalServerError, NotFoundError, - UnprocessableEntityError, + PermissionDeniedError, RateLimitError, - InternalServerError, + UnprocessableEntityError, ) -import sentry_sdk from langfuse import observe -from util import ApolloError, create_logger from models import resolve_model from search_docsite.search_docsite import resolve_backend +from util import ApolloError, create_logger + from .rag_config_loader import ConfigLoader logger = create_logger("job_chat.retrieve_docs") diff --git a/services/job_chat/tests/integration/test_adaptor_docs_pipeline.py b/services/job_chat/tests/integration/test_adaptor_docs_pipeline.py index 8b668098..7f08e7aa 100644 --- a/services/job_chat/tests/integration/test_adaptor_docs_pipeline.py +++ b/services/job_chat/tests/integration/test_adaptor_docs_pipeline.py @@ -3,7 +3,6 @@ import psycopg2 import pytest from dotenv import load_dotenv - from job_chat.prompt import generate_system_message load_dotenv() @@ -147,7 +146,7 @@ def test_generate_queries_returns_valid_structure(): print("==================TEST==================") print("Description: Testing generate_queries returns valid JSON structure") - from job_chat.retrieve_docs import generate_queries, get_client, format_context + from job_chat.retrieve_docs import format_context, generate_queries, get_client # Step 1: Prepare test inputs print("\n1. Preparing test inputs...") @@ -192,7 +191,7 @@ def test_search_docs_returns_general_docs_only(): from job_chat.retrieve_docs import search_docs queries = [{"query": "http adaptor merge() function"}] - results = search_docs(queries, top_k=3) + results = search_docs(queries, top_k=3, threshold=0.5) print(results) assert isinstance(results, list), "Should return a list" diff --git a/services/job_chat/tests/unit/conftest.py b/services/job_chat/tests/unit/conftest.py index 1816dd69..ce4637da 100644 --- a/services/job_chat/tests/unit/conftest.py +++ b/services/job_chat/tests/unit/conftest.py @@ -4,8 +4,7 @@ module-level `OpenAIEmbeddings()` default arg validates credentials at construction (openai 2.x / langchain-openai 1.x). A key must therefore exist at import time. -Dummy placeholders only — unit tests mock every network seam, and the repo-root -conftest blocks real client construction. `setdefault` lets a real key win. +Dummy placeholders only. """ import os diff --git a/services/job_chat/tests/unit/test_retrieve_docs.py b/services/job_chat/tests/unit/test_retrieve_docs.py index d53387e1..45b57ddf 100644 --- a/services/job_chat/tests/unit/test_retrieve_docs.py +++ b/services/job_chat/tests/unit/test_retrieve_docs.py @@ -14,13 +14,11 @@ from unittest.mock import MagicMock, patch import pytest - from embeddings.embeddings import SearchResult from job_chat import retrieve_docs as rd from job_chat.retrieve_docs import search_docs from util import ApolloError - # --- generate_queries ---------------------------------------------------------- def test_generate_queries_truncates_to_four(): diff --git a/services/search_docsite/pinecone_legacy_search.py b/services/search_docsite/pinecone_legacy_search.py index 11676dcd..90d1fa01 100644 --- a/services/search_docsite/pinecone_legacy_search.py +++ b/services/search_docsite/pinecone_legacy_search.py @@ -1,9 +1,10 @@ import os -from pinecone import Pinecone -from langchain_pinecone import PineconeVectorStore -from langchain_openai import OpenAIEmbeddings -from util import create_logger, ApolloError + from embeddings.embeddings import SearchResult +from langchain_openai import OpenAIEmbeddings +from langchain_pinecone import PineconeVectorStore +from pinecone import Pinecone +from util import ApolloError, create_logger logger = create_logger("LegacyPineconeDocsiteSearch") diff --git a/services/search_docsite/search_docsite.py b/services/search_docsite/search_docsite.py index a066b1af..7638a55a 100644 --- a/services/search_docsite/search_docsite.py +++ b/services/search_docsite/search_docsite.py @@ -1,12 +1,13 @@ import os + import psycopg2 from dotenv import load_dotenv +from embeddings.embeddings import SearchResult from langchain_openai import OpenAIEmbeddings from pgvector import Vector from pgvector.psycopg2 import register_vector -from util import create_logger, ApolloError, get_db_connection -from embeddings.embeddings import SearchResult from search_docsite.pinecone_legacy_search import LegacyPineconeDocsiteSearch +from util import ApolloError, create_logger, get_db_connection logger = create_logger("DocsiteSearch") @@ -42,11 +43,6 @@ def embeddings(self): def _connect(self): """Open a connection with pgvector registered. - - A database the indexer has never touched fails here rather than at the - query: register_vector looks up pgvector's type oid and raises if the - extension is absent. 503, not 500 — it resolves by running the indexer, - with no redeploy. """ conn = get_db_connection() try: @@ -108,9 +104,6 @@ def _semantic_search(self, conn, batch_id, query, top_k, threshold, doc_title, d top_k = self.default_top_k max_k = top_k or 50 - # Vector, not list: psycopg2 registers adapters for Vector and ndarray - # only, and a bare list arrives as numeric[], for which `<=>` has no - # operator. query_embedding = Vector(self.embeddings.embed_query(query)) sql = """ diff --git a/services/search_docsite/tests/conftest.py b/services/search_docsite/tests/conftest.py index d066197b..eb256c39 100644 --- a/services/search_docsite/tests/conftest.py +++ b/services/search_docsite/tests/conftest.py @@ -5,9 +5,7 @@ credentials at construction time. That happens when the test module is imported, before any test runs — so a key must exist in the environment or import fails. -These are dummy placeholders only: unit tests inject mocks for every network -seam and the repo-root conftest additionally blocks real client construction, so -no real key is ever used. `setdefault` means a real key (from services/.env) wins. +We only set dummy environment variables for unit tests. """ import os diff --git a/services/search_docsite/tests/unit/test_docsite_search.py b/services/search_docsite/tests/unit/test_docsite_search.py index cc68093e..c283481f 100644 --- a/services/search_docsite/tests/unit/test_docsite_search.py +++ b/services/search_docsite/tests/unit/test_docsite_search.py @@ -1,8 +1,8 @@ """Unit tests for the Postgres-backed DocsiteSearch (semantic/keyword/hybrid). -get_db_connection and register_vector_type are mocked throughout — no real +get_db_connection and register_vector_type are mocked throughout, no real Postgres connection is made. The OpenAI embeddings client is mocked via the -lazy `_embeddings` attribute, matching the DocsiteIndexer test pattern. +`_embeddings` attribute, matching the DocsiteIndexer test pattern. """ import json @@ -11,7 +11,6 @@ import psycopg2 import pytest - import search_docsite.search_docsite as m from pgvector import Vector from util import ApolloError diff --git a/services/search_docsite/tests/unit/test_pinecone_legacy_search.py b/services/search_docsite/tests/unit/test_pinecone_legacy_search.py index 5d6b2bf2..ed67f6e6 100644 --- a/services/search_docsite/tests/unit/test_pinecone_legacy_search.py +++ b/services/search_docsite/tests/unit/test_pinecone_legacy_search.py @@ -16,7 +16,6 @@ from unittest.mock import MagicMock, patch import pytest - import search_docsite.pinecone_legacy_search as m from util import ApolloError diff --git a/services/search_docsite/tests/unit/test_search_docsite_main.py b/services/search_docsite/tests/unit/test_search_docsite_main.py index 19af97d9..c6724b84 100644 --- a/services/search_docsite/tests/unit/test_search_docsite_main.py +++ b/services/search_docsite/tests/unit/test_search_docsite_main.py @@ -7,7 +7,6 @@ from unittest.mock import patch import pytest - import search_docsite.search_docsite as m from util import ApolloError diff --git a/services/tools/search_documentation/search_documentation.py b/services/tools/search_documentation/search_documentation.py index 65697a20..7e771a2c 100644 --- a/services/tools/search_documentation/search_documentation.py +++ b/services/tools/search_documentation/search_documentation.py @@ -6,15 +6,15 @@ 2. As a tool by supervisor via search_documentation_tool() """ import sys +from dataclasses import dataclass from pathlib import Path from typing import Dict -from dataclasses import dataclass # Import utilities from services directory sys.path.append(str(Path(__file__).parent.parent.parent)) -from util import create_logger, ApolloError from search_docsite.search_docsite import resolve_backend +from util import ApolloError, create_logger logger = create_logger(__name__) @@ -46,10 +46,7 @@ def _search_implementation(query: str, num_results: int) -> Dict: logger.info(f"Searching documentation for: {query[:100]}...") # Both backends use semantic search with the same cosine-similarity cutoff, so - # results are directly comparable. Hybrid (RRF) is deliberately not used here: - # its score has no calibratable scale, so it can be neither thresholded nor - # rendered as a relevance figure. It stays available via the search_docsite - # service and run_eval for evaluation. + # results are directly comparable. # Initialize docsite search docsite_search = resolve_backend()() diff --git a/services/util.py b/services/util.py index b4164fe6..45d6d0f7 100644 --- a/services/util.py +++ b/services/util.py @@ -112,11 +112,6 @@ def apollo(name: str, payload: dict) -> dict: def get_db_connection() -> "psycopg2.extensions.connection": """Get database connection from POSTGRES_URL environment variable. - Returns a plain connection. Schema migrations belong to the indexer and are - run explicitly by embed_docsite: applying them here would put CREATE - EXTENSION in the path of every reader, including deployments that never use - the Postgres docsite backend and roles without the privilege to run it. - Returns: psycopg2.connection: Database connection @@ -126,7 +121,6 @@ def get_db_connection() -> "psycopg2.extensions.connection": db_url = os.environ.get("POSTGRES_URL") if not db_url: raise ApolloError(500, "Missing POSTGRES_URL environment variable", type="DATABASE_ERROR") - return psycopg2.connect(db_url) From 6cbaec16d1c0b2e2ce04c99795c3bb44496617da Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Thu, 30 Jul 2026 09:52:51 +0800 Subject: [PATCH 39/48] lint: fix B008 and B904 linting issues --- .../embed_docsite/pinecone_legacy_indexer.py | 6 +++--- services/job_chat/retrieve_docs.py | 20 +++++++++---------- .../search_docsite/pinecone_legacy_search.py | 5 +++-- .../search_documentation.py | 4 ++-- 4 files changed, 18 insertions(+), 17 deletions(-) diff --git a/services/embed_docsite/pinecone_legacy_indexer.py b/services/embed_docsite/pinecone_legacy_indexer.py index 263991c7..f9332e9c 100644 --- a/services/embed_docsite/pinecone_legacy_indexer.py +++ b/services/embed_docsite/pinecone_legacy_indexer.py @@ -25,10 +25,10 @@ class LegacyPineconeDocsiteIndexer: :param dimension: Embedding dimension (default: 1536 for OpenAI Embeddings) :param max_total_collections: Max total collections in index. Delete old collections by date if exceeded after a new upload (default: 50) """ - def __init__(self, collection_name=None, index_name="docsite", embeddings=OpenAIEmbeddings(), dimension=1536, max_total_collections=50): + def __init__(self, collection_name=None, index_name="docsite", embeddings=None, dimension=1536, max_total_collections=50): self.collection_name = collection_name if collection_name is not None else f"docsite-{datetime.now().strftime('%Y%m%d%H%M')}" self.index_name = index_name - self.embeddings = embeddings + self.embeddings = embeddings if embeddings is not None else OpenAIEmbeddings() self.dimension = dimension self.max_total_collections = max_total_collections self.pc = Pinecone(api_key=os.environ.get("PINECONE_API_KEY")) @@ -37,7 +37,7 @@ def __init__(self, collection_name=None, index_name="docsite", embeddings=OpenAI self.create_index() self.index = self.pc.Index(self.index_name) - self.vectorstore = PineconeVectorStore(index_name=index_name, namespace=self.collection_name, embedding=embeddings) + self.vectorstore = PineconeVectorStore(index_name=index_name, namespace=self.collection_name, embedding=self.embeddings) def insert_documents(self, inputs, metadata_dict): """ diff --git a/services/job_chat/retrieve_docs.py b/services/job_chat/retrieve_docs.py index 0ec9ced1..33d35a89 100644 --- a/services/job_chat/retrieve_docs.py +++ b/services/job_chat/retrieve_docs.py @@ -164,7 +164,7 @@ def generate_queries(content, client, user_context=""): "Failed to generate search queries - invalid response from AI service", type="INVALID_LLM_RESPONSE", details={"response_preview": text[:200]} - ) + ) from e if len(answer_parsed) >= 4: answer_parsed = answer_parsed[:4] @@ -255,10 +255,10 @@ def call_llm(model, temperature, system_prompt, user_prompt, client, output_sche "Unable to reach the AI service for documentation search", type="CONNECTION_ERROR", details=details, - ) + ) from e except AuthenticationError as e: logger.error(f"Authentication error during knowledge retrieval: {e}") - raise ApolloError(401, "Authentication failed with AI service", type="AUTH_ERROR") + raise ApolloError(401, "Authentication failed with AI service", type="AUTH_ERROR") from e except RateLimitError as e: logger.error(f"Rate limit error during knowledge retrieval: {e}") retry_after = int(e.response.headers.get('retry-after', 60)) if hasattr(e, 'response') else 60 @@ -267,22 +267,22 @@ def call_llm(model, temperature, system_prompt, user_prompt, client, output_sche "Rate limit exceeded for documentation search, please try again later", type="RATE_LIMIT", details={"retry_after": retry_after} - ) + ) from e except BadRequestError as e: logger.error(f"Bad request error during knowledge retrieval: {e}") - raise ApolloError(400, f"Invalid request to AI service: {str(e)}", type="BAD_REQUEST") + raise ApolloError(400, f"Invalid request to AI service: {str(e)}", type="BAD_REQUEST") from e except PermissionDeniedError as e: logger.error(f"Permission denied error during knowledge retrieval: {e}") - raise ApolloError(403, "Not authorized to perform this action", type="FORBIDDEN") + raise ApolloError(403, "Not authorized to perform this action", type="FORBIDDEN") from e except NotFoundError as e: logger.error(f"Not found error during knowledge retrieval: {e}") - raise ApolloError(404, "Resource not found", type="NOT_FOUND") + raise ApolloError(404, "Resource not found", type="NOT_FOUND") from e except UnprocessableEntityError as e: logger.error(f"Unprocessable entity error during knowledge retrieval: {e}") - raise ApolloError(422, str(e), type="INVALID_REQUEST") + raise ApolloError(422, str(e), type="INVALID_REQUEST") from e except InternalServerError as e: logger.error(f"Internal server error from AI service during knowledge retrieval: {e}") - raise ApolloError(500, "The AI service encountered an error", type="PROVIDER_ERROR") + raise ApolloError(500, "The AI service encountered an error", type="PROVIDER_ERROR") from e except Exception as e: logger.error(f"Unexpected error during LLM call for knowledge retrieval: {str(e)}") - raise ApolloError(500, f"Unexpected error during documentation search: {str(e)}", type="UNKNOWN_ERROR") \ No newline at end of file + raise ApolloError(500, f"Unexpected error during documentation search: {str(e)}", type="UNKNOWN_ERROR") from e \ No newline at end of file diff --git a/services/search_docsite/pinecone_legacy_search.py b/services/search_docsite/pinecone_legacy_search.py index 90d1fa01..384f7ef3 100644 --- a/services/search_docsite/pinecone_legacy_search.py +++ b/services/search_docsite/pinecone_legacy_search.py @@ -20,16 +20,17 @@ class LegacyPineconeDocsiteSearch: :param default_top_k: Default number of results to return (default: 5) :param embeddings: LangChain embedding type (default: OpenAIEmbeddings()) """ - def __init__(self, collection_name=None, index_name="docsite", default_top_k=5, embeddings=OpenAIEmbeddings()): + def __init__(self, collection_name=None, index_name="docsite", default_top_k=5, embeddings=None): self.index_client = index_name self.default_top_k = default_top_k + self.embeddings = embeddings if embeddings is not None else OpenAIEmbeddings() if collection_name is None: logger.info("Collection name not provided; retrieving the most recent collection name.") collection_name = self._get_most_recent_namespace() self.collection_name = collection_name - self.vectorstore = PineconeVectorStore(index_name=index_name, namespace=collection_name, embedding=embeddings) + self.vectorstore = PineconeVectorStore(index_name=index_name, namespace=collection_name, embedding=self.embeddings) def search(self, query, top_k=None, threshold=None, strategy='semantic', doc_title=None, docs_type=None): filters = self._build_filter(doc_title=doc_title, docs_type=docs_type) diff --git a/services/tools/search_documentation/search_documentation.py b/services/tools/search_documentation/search_documentation.py index 7e771a2c..018cc172 100644 --- a/services/tools/search_documentation/search_documentation.py +++ b/services/tools/search_documentation/search_documentation.py @@ -92,7 +92,7 @@ def main(data: Dict) -> Dict: raise except Exception as e: logger.exception("Error in search_documentation service") - raise ApolloError(500, f"Documentation search failed: {str(e)}") + raise ApolloError(500, f"Documentation search failed: {str(e)}") from e def search_documentation_tool(tool_input: Dict) -> str: @@ -141,4 +141,4 @@ def search_documentation_tool(tool_input: Dict) -> str: except Exception as e: logger.exception("Error in search_documentation tool") - raise ApolloError(500, f"Documentation search failed: {str(e)}") + raise ApolloError(500, f"Documentation search failed: {str(e)}") from e From ad6e2e82e523132cac298137225d356a2945617c Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Thu, 30 Jul 2026 09:58:51 +0800 Subject: [PATCH 40/48] lint: fix B905, RUF059, edited docs in helpers --- services/embed_docsite/docsite_indexer.py | 2 +- services/embed_docsite/tests/integration/helpers.py | 4 +--- .../job_chat/tests/integration/test_adaptor_docs_pipeline.py | 2 +- services/search_docsite/tests/eval/run_eval.py | 2 +- 4 files changed, 4 insertions(+), 6 deletions(-) diff --git a/services/embed_docsite/docsite_indexer.py b/services/embed_docsite/docsite_indexer.py index c16ea57e..c7a32480 100644 --- a/services/embed_docsite/docsite_indexer.py +++ b/services/embed_docsite/docsite_indexer.py @@ -83,7 +83,7 @@ def insert_documents(self, conn, batch_id, documents, metadata_dict): doc_title_indices = {} rows = [] - for doc, embedding in zip(documents, embeddings): + for doc, embedding in zip(documents, embeddings, strict=True): doc_title = doc["name"].removesuffix(".md") chunk_index = doc_title_indices.get(doc_title, 0) doc_title_indices[doc_title] = chunk_index + 1 diff --git a/services/embed_docsite/tests/integration/helpers.py b/services/embed_docsite/tests/integration/helpers.py index 02f207e5..0dae25fd 100644 --- a/services/embed_docsite/tests/integration/helpers.py +++ b/services/embed_docsite/tests/integration/helpers.py @@ -21,9 +21,7 @@ class StubEmbeddings: """Deterministic stand-in for OpenAIEmbeddings. Identical text yields an identical vector, so querying a chunk's exact text - puts that chunk at cosine distance 0 and therefore rank 1. Stubbing costs no - coverage here: `operator does not exist: vector <=> numeric[]` is raised - from the bound parameter's type, never its content. + puts that chunk at cosine distance 0 and therefore rank 1. """ model = "stub-embedding-model" diff --git a/services/job_chat/tests/integration/test_adaptor_docs_pipeline.py b/services/job_chat/tests/integration/test_adaptor_docs_pipeline.py index 7f08e7aa..f926316e 100644 --- a/services/job_chat/tests/integration/test_adaptor_docs_pipeline.py +++ b/services/job_chat/tests/integration/test_adaptor_docs_pipeline.py @@ -158,7 +158,7 @@ def test_generate_queries_returns_valid_structure(): # Step 2: Call generate_queries print("\n2. Calling generate_queries...") client = get_client() - queries, usage = generate_queries(question, client, user_context) + queries, _usage = generate_queries(question, client, user_context) print(f" Generated {len(queries)} queries") # Step 3: Validate structure diff --git a/services/search_docsite/tests/eval/run_eval.py b/services/search_docsite/tests/eval/run_eval.py index 4e8704a0..f352d6d1 100644 --- a/services/search_docsite/tests/eval/run_eval.py +++ b/services/search_docsite/tests/eval/run_eval.py @@ -36,7 +36,7 @@ def compute_agreement(report_a, report_b): golden_queries.yaml is curated. """ per_query = [] - for a, b in zip(report_a["per_query"], report_b["per_query"]): + for a, b in zip(report_a["per_query"], report_b["per_query"], strict=True): titles_a = {t for t in a["retrieved_titles"] if t is not None} titles_b = {t for t in b["retrieved_titles"] if t is not None} union = titles_a | titles_b From 4bf19d476c696bd7499b78da044ae7c57c863727 Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Thu, 30 Jul 2026 10:04:07 +0800 Subject: [PATCH 41/48] fix: lint F401 and docsite name accepting in pinecone --- services/embed_docsite/pinecone_legacy_indexer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/services/embed_docsite/pinecone_legacy_indexer.py b/services/embed_docsite/pinecone_legacy_indexer.py index f9332e9c..d0d2ae7c 100644 --- a/services/embed_docsite/pinecone_legacy_indexer.py +++ b/services/embed_docsite/pinecone_legacy_indexer.py @@ -7,7 +7,7 @@ from langchain_openai import OpenAIEmbeddings from langchain_pinecone import PineconeVectorStore from pinecone import Pinecone, ServerlessSpec -from util import ApolloError, create_logger +from util import create_logger logger = create_logger("LegacyPineconeDocsiteIndexer") @@ -113,7 +113,7 @@ def delete_old_collections(self, max_total_collections): index_stats = index.describe_index_stats() namespaces = index_stats.get('namespaces', {}).keys() valid_namespaces = sorted( - (ns for ns in namespaces if ns.startswith("docsite-") and ns[8:].isdigit() and len(ns) == 16), + (ns for ns in namespaces if ns.startswith("docsite-") and ns[8:].isdigit() and len(ns) in (16, 20)), reverse=False ) if len(valid_namespaces) > max_total_collections: From 3b78106d65f7cf7034cf7f79024b918292ca1bf5 Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Thu, 30 Jul 2026 11:05:50 +0800 Subject: [PATCH 42/48] feat: add results for each query for run_eval.py --- .../search_docsite/tests/eval/run_eval.py | 22 ++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/services/search_docsite/tests/eval/run_eval.py b/services/search_docsite/tests/eval/run_eval.py index f352d6d1..ecd955a6 100644 --- a/services/search_docsite/tests/eval/run_eval.py +++ b/services/search_docsite/tests/eval/run_eval.py @@ -3,9 +3,10 @@ Usage: poetry run python -m search_docsite.tests.eval.run_eval Compares the Postgres-backed DocsiteSearch against LegacyPineconeDocsiteSearch -over the golden query set. Both run strategy='semantic'. Postgres should match -or beat Pinecone's recall@5 and p95 latency before DOCSITE_SEARCH_BACKEND is -flipped to 'postgres' by default. +over the golden query set. Both run strategy='semantic'. + +Golden queries live in services/search_docsite/tests/eval/golden_queries.yaml +— edit that file to add queries or curate expected_doc_titles. """ import time @@ -103,6 +104,19 @@ def run_eval(golden_queries, backend_cls, strategy, top_k=5): } +def print_per_query_results(postgres_report, pinecone_report): + """Print each golden query's expected titles and both backends' hit/miss + retrieved titles.""" + print("\nPer-query results:") + for pg, pc in zip(postgres_report["per_query"], pinecone_report["per_query"], strict=True): + expected = ", ".join(pg["expected_titles"]) or "(unlabelled — excluded from recall)" + print(f" {pg['query']}") + print(f" expected: {expected}") + for name, report_item in (("postgres", pg), ("pinecone", pc)): + status = {True: "HIT ", False: "MISS", None: "-- "}[report_item["hit"]] + titles = ", ".join(t for t in report_item["retrieved_titles"] if t is not None) + print(f" {name:8} {status} {report_item['latency_s']:.3f}s {titles}") + + def main(): with open(GOLDEN_QUERIES_PATH) as f: golden_queries = yaml.safe_load(f)["queries"] @@ -123,6 +137,8 @@ def main(): for q in agreement["per_query"]: print(f" {q['jaccard']:.2f} overlap={q['overlap']}/{q['union']} {q['query']}") + print_per_query_results(postgres_report, pinecone_report) + if __name__ == "__main__": main() From b994b6ffd96f2fa99f2760d47c31440bf69eff3a Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Fri, 7 Aug 2026 17:52:42 +0800 Subject: [PATCH 43/48] refactor: use a timestamped migration filename --- ...d_chunks.sql => 20260728000000_docsite_batches_and_chunks.sql} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename services/migrations/{0001_docsite_batches_and_chunks.sql => 20260728000000_docsite_batches_and_chunks.sql} (100%) diff --git a/services/migrations/0001_docsite_batches_and_chunks.sql b/services/migrations/20260728000000_docsite_batches_and_chunks.sql similarity index 100% rename from services/migrations/0001_docsite_batches_and_chunks.sql rename to services/migrations/20260728000000_docsite_batches_and_chunks.sql From 1b401d85f0b559ee426a73fe9603f3111d82dd46 Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Fri, 7 Aug 2026 17:53:45 +0800 Subject: [PATCH 44/48] test: expand golden query set to 30 categorised queries --- .../tests/eval/golden_queries.yaml | 94 +++++++++++++++++-- 1 file changed, 86 insertions(+), 8 deletions(-) diff --git a/services/search_docsite/tests/eval/golden_queries.yaml b/services/search_docsite/tests/eval/golden_queries.yaml index 3a931d51..9f4cbad1 100644 --- a/services/search_docsite/tests/eval/golden_queries.yaml +++ b/services/search_docsite/tests/eval/golden_queries.yaml @@ -4,28 +4,106 @@ # so each entry lists every doc that would be a reasonable answer, not only the # single best one. Titles are docsite filenames without .md, matching # docsite_chunks.doc_title. +# +# category is 'conceptual' (natural-language questions) or 'keyword' (exact-term +# lookups). run_eval reports recall per category, because hybrid search is +# expected to help on keyword lookups and tie on conceptual ones — a blended +# figure would hide exactly that effect. +# +# Labels were derived from the general_docs corpus via full-text search on each +# query's key terms, deliberately not from vector-search output: labelling from +# the thing under test would be circular. queries: + # --- conceptual ------------------------------------------------------------- - query: "how do I configure a webhook trigger" + category: conceptual expected_doc_titles: ["triggers", "webhook-auth"] - query: "how do I set up a cron trigger" + category: conceptual expected_doc_titles: ["triggers"] - query: "what is a run in OpenFn" + category: conceptual expected_doc_titles: ["terminology", "glossary"] - query: "how do I use collections to store state between runs" + category: conceptual expected_doc_titles: ["collections", "state", "cli-collections"] - query: "how do I configure a credential for an adaptor" + category: conceptual expected_doc_titles: ["credentials", "manage-credentials"] - query: "what is the difference between a job and a workflow" + category: conceptual expected_doc_titles: ["terminology", "glossary", "workflows"] - query: "how do I deploy a project using the CLI" + category: conceptual expected_doc_titles: ["cli-sync", "portability"] - query: "how do I debug a failed run" - expected_doc_titles: - ["troubleshooting", "rerunning-workflow", "inspect-runs"] - # Unscorable as run_eval filters to docs_type='general_docs', but the - # adaptor listing lives in adaptor_docs. - - query: "what adaptors are available for HTTP requests" - expected_doc_titles: [] + category: conceptual + expected_doc_titles: ["troubleshooting", "rerunning-workflow", "inspect-runs"] - query: "how do I write a data transform function" - expected_doc_titles: - ["data-transformation", "job-writing-guide", "operations"] + category: conceptual + expected_doc_titles: ["data-transformation", "job-writing-guide", "operations"] + - query: "how do I control who can access a project" + category: conceptual + expected_doc_titles: ["collaboration", "user-roles-permissions"] + - query: "how do I set up a sandbox environment" + category: conceptual + expected_doc_titles: ["sandboxes"] + - query: "how long is my run data kept" + category: conceptual + expected_doc_titles: ["retention-periods", "io-data-storage", "security-for-devs"] + - query: "how do I get notified when a workflow fails" + category: conceptual + expected_doc_titles: ["notifications"] + - query: "how do I version control my project with GitHub" + category: conceptual + expected_doc_titles: ["link-to-gh", "cli-sync"] + - query: "what security measures does OpenFn have" + category: conceptual + expected_doc_titles: ["security", "security-compliance"] + + # --- keyword ---------------------------------------------------------------- + - query: "state.cursor" + category: keyword + expected_doc_titles: ["using-cursors"] + - query: "--force flag on project push" + category: keyword + expected_doc_titles: ["cli-sync"] + - query: "webhook auth method" + category: keyword + expected_doc_titles: ["webhook-auth"] + - query: "workflow snapshots" + category: keyword + expected_doc_titles: ["workflow-snapshots"] + - query: "lazy state operator" + category: keyword + expected_doc_titles: ["lazy-state-operator"] + - query: "dataValue" + category: keyword + expected_doc_titles: ["job-snippets"] + - query: "each operation" + category: keyword + expected_doc_titles: ["operations"] + - query: "rerun a workflow" + category: keyword + expected_doc_titles: ["rerunning-workflow"] + - query: "collections CLI commands" + category: keyword + expected_doc_titles: ["cli-collections"] + - query: "activity history" + category: keyword + expected_doc_titles: ["activity-history"] + - query: "project.yaml" + category: keyword + expected_doc_titles: ["portability-v3"] + - query: "step editor" + category: keyword + expected_doc_titles: ["step-editor"] + - query: "git branch" + category: keyword + expected_doc_titles: ["working-with-branches", "cli-sync"] + - query: "sandbox" + category: keyword + expected_doc_titles: ["sandboxes"] + - query: "security compliance" + category: keyword + expected_doc_titles: ["security-compliance"] From 4ae1fb1a81ed588b21567c2593858022ae3fee68 Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Fri, 7 Aug 2026 17:55:51 +0800 Subject: [PATCH 45/48] feat: report eval recall per query category --- .../search_docsite/tests/eval/run_eval.py | 15 ++++++- .../tests/unit/test_run_eval.py | 42 +++++++++++++++++++ 2 files changed, 55 insertions(+), 2 deletions(-) diff --git a/services/search_docsite/tests/eval/run_eval.py b/services/search_docsite/tests/eval/run_eval.py index ecd955a6..a0161461 100644 --- a/services/search_docsite/tests/eval/run_eval.py +++ b/services/search_docsite/tests/eval/run_eval.py @@ -53,18 +53,20 @@ def compute_agreement(report_a, report_b): return {"per_query": per_query, "mean_jaccard": mean_jaccard} -def run_eval(golden_queries, backend_cls, strategy, top_k=5): +def run_eval(golden_queries, make_backend, strategy, top_k=5): """Run every golden query against one backend/strategy and return a report dict.""" - backend = backend_cls() + backend = make_backend() per_query = [] latencies = [] scored = 0 skipped = 0 hits = 0 + by_category = {} for item in golden_queries: query = item["query"] expected_titles = item.get("expected_doc_titles", []) + category = item.get("category") start = time.time() results = backend.search(query, top_k=top_k, strategy=strategy, docs_type="general_docs") @@ -77,12 +79,17 @@ def run_eval(golden_queries, backend_cls, strategy, top_k=5): hit = compute_recall_at_k(retrieved_titles, expected_titles) scored += 1 hits += int(hit) + if category: + counts = by_category.setdefault(category, {"scored": 0, "hits": 0}) + counts["scored"] += 1 + counts["hits"] += int(hit) else: hit = None skipped += 1 per_query.append({ "query": query, + "category": category, "retrieved_titles": retrieved_titles, "expected_titles": expected_titles, "hit": hit, @@ -96,6 +103,10 @@ def run_eval(golden_queries, backend_cls, strategy, top_k=5): return { "recall_at_k": (hits / scored) if scored else None, + "recall_by_category": { + name: {"recall": c["hits"] / c["scored"], "scored": c["scored"]} + for name, c in sorted(by_category.items()) + }, "queries_scored": scored, "queries_skipped": skipped, "p50_latency_s": p50, diff --git a/services/search_docsite/tests/unit/test_run_eval.py b/services/search_docsite/tests/unit/test_run_eval.py index 83599cf6..873d2e22 100644 --- a/services/search_docsite/tests/unit/test_run_eval.py +++ b/services/search_docsite/tests/unit/test_run_eval.py @@ -112,3 +112,45 @@ def test_compute_agreement_means_across_queries(): agreement = compute_agreement(report_a, report_b) assert agreement["mean_jaccard"] == pytest.approx(0.5) + + +def test_run_eval_reports_recall_per_category(): + golden_queries = [ + {"query": "c1", "category": "conceptual", "expected_doc_titles": ["A"]}, + {"query": "c2", "category": "conceptual", "expected_doc_titles": ["B"]}, + {"query": "k1", "category": "keyword", "expected_doc_titles": ["C"]}, + {"query": "k2", "category": "keyword", "expected_doc_titles": ["MISSING"]}, + ] + backend_cls = make_fake_backend({"c1": ["A"], "c2": ["B"], "k1": ["C"], "k2": ["Z"]}) + + report = run_eval(golden_queries, backend_cls, strategy="semantic", top_k=5) + + assert report["recall_by_category"]["conceptual"] == {"recall": 1.0, "scored": 2} + assert report["recall_by_category"]["keyword"] == {"recall": 0.5, "scored": 2} + assert report["recall_at_k"] == 0.75 + + +def test_run_eval_tolerates_queries_without_a_category(): + golden_queries = [ + {"query": "q1", "expected_doc_titles": ["A"]}, + {"query": "q2", "category": "keyword", "expected_doc_titles": ["B"]}, + ] + backend_cls = make_fake_backend({"q1": ["A"], "q2": ["B"]}) + + report = run_eval(golden_queries, backend_cls, strategy="semantic", top_k=5) + + assert report["recall_at_k"] == 1.0 + assert report["recall_by_category"] == {"keyword": {"recall": 1.0, "scored": 1}} + + +def test_run_eval_excludes_unlabelled_queries_from_category_recall(): + golden_queries = [ + {"query": "k1", "category": "keyword", "expected_doc_titles": ["A"]}, + {"query": "k2", "category": "keyword", "expected_doc_titles": []}, + ] + backend_cls = make_fake_backend({"k1": ["A"], "k2": ["B"]}) + + report = run_eval(golden_queries, backend_cls, strategy="semantic", top_k=5) + + assert report["recall_by_category"]["keyword"] == {"recall": 1.0, "scored": 1} + assert report["queries_skipped"] == 1 From 2fed53a4218438d5d6c93832815a9ff3d2e69249 Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Fri, 7 Aug 2026 17:56:42 +0800 Subject: [PATCH 46/48] feat: resolve docsite eval batches by chunk size --- .../search_docsite/tests/eval/run_eval.py | 22 +++++++++++++ .../tests/unit/test_run_eval.py | 33 +++++++++++++++++++ 2 files changed, 55 insertions(+) diff --git a/services/search_docsite/tests/eval/run_eval.py b/services/search_docsite/tests/eval/run_eval.py index a0161461..004fa443 100644 --- a/services/search_docsite/tests/eval/run_eval.py +++ b/services/search_docsite/tests/eval/run_eval.py @@ -21,10 +21,32 @@ from search_docsite.pinecone_legacy_search import LegacyPineconeDocsiteSearch from search_docsite.search_docsite import DocsiteSearch +from util import get_db_connection GOLDEN_QUERIES_PATH = Path(__file__).parent / "golden_queries.yaml" +def resolve_batch_id(chunk_target_length): + """Newest complete batch indexed at the given chunk size, or None if there is none. + + Batch ids are environment-specific, so the eval looks them up by the indexing + configuration recorded on each batch rather than hardcoding them. + """ + conn = get_db_connection() + try: + with conn.cursor() as cur: + cur.execute( + "SELECT id FROM docsite_batches WHERE status = 'complete' " + "AND chunk_target_length = %s ORDER BY id DESC LIMIT 1", + (chunk_target_length,), + ) + row = cur.fetchone() + finally: + conn.close() + + return row[0] if row else None + + def compute_recall_at_k(retrieved_titles, expected_titles): """True if any expected title appears among the retrieved titles.""" return bool(set(retrieved_titles) & set(expected_titles)) diff --git a/services/search_docsite/tests/unit/test_run_eval.py b/services/search_docsite/tests/unit/test_run_eval.py index 873d2e22..02f62428 100644 --- a/services/search_docsite/tests/unit/test_run_eval.py +++ b/services/search_docsite/tests/unit/test_run_eval.py @@ -154,3 +154,36 @@ def test_run_eval_excludes_unlabelled_queries_from_category_recall(): assert report["recall_by_category"]["keyword"] == {"recall": 1.0, "scored": 1} assert report["queries_skipped"] == 1 + + +def test_resolve_batch_id_returns_newest_matching_batch(): + from unittest.mock import patch + + import search_docsite.tests.eval.run_eval as m + + conn = MagicMock() + cur = MagicMock() + conn.cursor.return_value.__enter__.return_value = cur + cur.fetchone.return_value = (11,) + + with patch.object(m, "get_db_connection", return_value=conn): + assert m.resolve_batch_id(2500) == 11 + + assert cur.execute.call_args[0][1] == (2500,) + conn.close.assert_called_once() + + +def test_resolve_batch_id_returns_none_when_no_batch_matches(): + from unittest.mock import patch + + import search_docsite.tests.eval.run_eval as m + + conn = MagicMock() + cur = MagicMock() + conn.cursor.return_value.__enter__.return_value = cur + cur.fetchone.return_value = None + + with patch.object(m, "get_db_connection", return_value=conn): + assert m.resolve_batch_id(1800) is None + + conn.close.assert_called_once() From f9e30d089456de4ada6f123181de34cfffed2233 Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Fri, 7 Aug 2026 17:59:59 +0800 Subject: [PATCH 47/48] feat: compare semantic and hybrid across chunk sizes in the eval --- .../search_docsite/tests/eval/run_eval.py | 100 +++++++++++++----- 1 file changed, 73 insertions(+), 27 deletions(-) diff --git a/services/search_docsite/tests/eval/run_eval.py b/services/search_docsite/tests/eval/run_eval.py index 004fa443..f20b6c10 100644 --- a/services/search_docsite/tests/eval/run_eval.py +++ b/services/search_docsite/tests/eval/run_eval.py @@ -2,8 +2,14 @@ Usage: poetry run python -m search_docsite.tests.eval.run_eval -Compares the Postgres-backed DocsiteSearch against LegacyPineconeDocsiteSearch -over the golden query set. Both run strategy='semantic'. +Compares Postgres semantic and hybrid search at each indexed chunk size against +the legacy Pinecone baseline, over the golden query set. Recall is reported per +query category, because hybrid is expected to help on keyword lookups and tie on +conceptual ones — a blended figure would hide that. + +Requires a Postgres batch per chunk size in CHUNK_SIZES; configurations without +one are skipped. Index them with embed_docsite, passing chunk_target_length, +chunk_min_length, and keep_batches >= 3 so earlier batches are not pruned. Golden queries live in services/search_docsite/tests/eval/golden_queries.yaml — edit that file to add queries or curate expected_doc_titles. @@ -25,6 +31,13 @@ GOLDEN_QUERIES_PATH = Path(__file__).parent / "golden_queries.yaml" +# Chunk sizes to evaluate, matching docsite_batches.chunk_target_length. +CHUNK_SIZES = [1000, 1800, 2500] +STRATEGIES = ["semantic", "hybrid"] + +# The migration-fidelity pairing: same strategy and chunk size, different store. +AGREEMENT_PAIR = ("Pinecone semantic", "Postgres semantic (1000)") + def resolve_batch_id(chunk_target_length): """Newest complete batch indexed at the given chunk size, or None if there is none. @@ -137,40 +150,73 @@ def run_eval(golden_queries, make_backend, strategy, top_k=5): } -def print_per_query_results(postgres_report, pinecone_report): +def print_per_query_results(label_a, report_a, label_b, report_b): """Print each golden query's expected titles and both backends' hit/miss + retrieved titles.""" print("\nPer-query results:") - for pg, pc in zip(postgres_report["per_query"], pinecone_report["per_query"], strict=True): - expected = ", ".join(pg["expected_titles"]) or "(unlabelled — excluded from recall)" - print(f" {pg['query']}") + for a, b in zip(report_a["per_query"], report_b["per_query"], strict=True): + expected = ", ".join(a["expected_titles"]) or "(unlabelled — excluded from recall)" + print(f" [{a['category'] or 'uncategorised'}] {a['query']}") print(f" expected: {expected}") - for name, report_item in (("postgres", pg), ("pinecone", pc)): - status = {True: "HIT ", False: "MISS", None: "-- "}[report_item["hit"]] - titles = ", ".join(t for t in report_item["retrieved_titles"] if t is not None) - print(f" {name:8} {status} {report_item['latency_s']:.3f}s {titles}") + for label, item in ((label_a, a), (label_b, b)): + status = {True: "HIT ", False: "MISS", None: "-- "}[item["hit"]] + titles = ", ".join(t for t in item["retrieved_titles"] if t is not None) + print(f" {label:26} {status} {item['latency_s']:.3f}s {titles}") -def main(): - with open(GOLDEN_QUERIES_PATH) as f: - golden_queries = yaml.safe_load(f)["queries"] +def build_configs(): + """Every (label, backend factory, strategy) to evaluate. - postgres_report = run_eval(golden_queries, DocsiteSearch, strategy="semantic") - pinecone_report = run_eval(golden_queries, LegacyPineconeDocsiteSearch, strategy="semantic") + Postgres configurations for a chunk size with no complete batch are skipped + with a notice, so the eval still runs when only some batches are indexed. + """ + configs = [("Pinecone semantic", LegacyPineconeDocsiteSearch, "semantic")] - print(f"Postgres (semantic): recall@5={postgres_report['recall_at_k']} " - f"(scored={postgres_report['queries_scored']}, skipped={postgres_report['queries_skipped']}) " - f"p50={postgres_report['p50_latency_s']:.3f}s p95={postgres_report['p95_latency_s']:.3f}s") - print(f"Pinecone (semantic): recall@5={pinecone_report['recall_at_k']} " - f"(scored={pinecone_report['queries_scored']}, skipped={pinecone_report['queries_skipped']}) " - f"p50={pinecone_report['p50_latency_s']:.3f}s p95={pinecone_report['p95_latency_s']:.3f}s") + for chunk_size in CHUNK_SIZES: + batch_id = resolve_batch_id(chunk_size) + if batch_id is None: + print(f"Skipping Postgres configs for chunk size {chunk_size} — no complete batch") + continue + for strategy in STRATEGIES: + label = f"Postgres {strategy} ({chunk_size})" + configs.append((label, lambda b=batch_id: DocsiteSearch(batch_id=b), strategy)) - agreement = compute_agreement(postgres_report, pinecone_report) - print(f"Backend agreement: mean doc-title Jaccard={agreement['mean_jaccard']:.3f} " - f"across {len(agreement['per_query'])} queries") - for q in agreement["per_query"]: - print(f" {q['jaccard']:.2f} overlap={q['overlap']}/{q['union']} {q['query']}") + return configs + + +def _fmt(recall): + """Render a recall figure, or a dash when the category had no scored queries.""" + return "-" if recall is None else f"{recall:.2f}" + + +def main(): + with open(GOLDEN_QUERIES_PATH) as f: + golden_queries = yaml.safe_load(f)["queries"] - print_per_query_results(postgres_report, pinecone_report) + reports = {} + for label, make_backend, strategy in build_configs(): + reports[label] = run_eval(golden_queries, make_backend, strategy) + + print(f"\n{'configuration':<28} {'conceptual':>11} {'keyword':>9} {'overall':>9} {'p50':>8} {'p95':>8}") + for label, report in reports.items(): + by_cat = report["recall_by_category"] + conceptual = by_cat.get("conceptual", {}).get("recall") + keyword = by_cat.get("keyword", {}).get("recall") + print(f"{label:<28} " + f"{_fmt(conceptual):>11} {_fmt(keyword):>9} {_fmt(report['recall_at_k']):>9} " + f"{report['p50_latency_s']:>7.3f}s {report['p95_latency_s']:>7.3f}s") + + scored = next(iter(reports.values())) + print(f"\nScored {scored['queries_scored']} queries, skipped {scored['queries_skipped']}") + + label_a, label_b = AGREEMENT_PAIR + if label_a in reports and label_b in reports: + agreement = compute_agreement(reports[label_a], reports[label_b]) + print(f"\nMigration fidelity ({label_a} vs {label_b}):") + print(f" mean doc-title Jaccard={agreement['mean_jaccard']:.3f} " + f"across {len(agreement['per_query'])} queries") + for q in agreement["per_query"]: + print(f" {q['jaccard']:.2f} overlap={q['overlap']}/{q['union']} {q['query']}") + print_per_query_results(label_a, reports[label_a], label_b, reports[label_b]) if __name__ == "__main__": From 10263d77dfdb7ee47847de7bdb0f5894b0d27276 Mon Sep 17 00:00:00 2001 From: IZO-Ong Date: Tue, 11 Aug 2026 20:23:34 +0800 Subject: [PATCH 48/48] docs: update comments in roundtrip test --- .../tests/integration/test_postgres_docsite_roundtrip.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/services/embed_docsite/tests/integration/test_postgres_docsite_roundtrip.py b/services/embed_docsite/tests/integration/test_postgres_docsite_roundtrip.py index 4096d219..23d67992 100644 --- a/services/embed_docsite/tests/integration/test_postgres_docsite_roundtrip.py +++ b/services/embed_docsite/tests/integration/test_postgres_docsite_roundtrip.py @@ -44,7 +44,7 @@ def make_search(**kwargs): def test_fresh_database_migrates_indexes_and_promotes(clean_db): - """C3: the first run on an empty database used to strand in 'building', + """The first run on an empty database used to strand in 'building', because copy_forward's SELECT left a transaction open.""" result = index_docs() @@ -55,7 +55,7 @@ def test_fresh_database_migrates_indexes_and_promotes(clean_db): @pytest.mark.parametrize("strategy", ["semantic", "keyword", "hybrid"]) def test_search_returns_the_indexed_chunk(clean_db, strategy): - """C1: semantic and hybrid used to fail with + """Semantic and hybrid used to fail with `operator does not exist: vector <=> numeric[]` on every query.""" index_docs() target = DOCS[0]["doc_chunk"] @@ -66,7 +66,7 @@ def test_search_returns_the_indexed_chunk(clean_db, strategy): def test_reindexing_prunes_the_previous_batch(clean_db): - """C2: pruning never ran, so every re-index permanently added a full + """Pruning never ran, so every re-index permanently added a full docsite copy and another HNSW index.""" first = index_docs(keep_batches=1) second = index_docs(keep_batches=1) @@ -78,7 +78,7 @@ def test_reindexing_prunes_the_previous_batch(clean_db): def test_reader_without_a_schema_gets_a_clear_503(clean_db): - """C4: with migrations moved to the indexer, a reader on an un-indexed + """With migrations moved to the indexer, a reader on an un-indexed database must explain itself rather than emit a psycopg2 traceback.""" with pytest.raises(ApolloError) as exc: make_search().search("anything")