Extension Points¶
Context: this document describes the extensibility surface of pydocs-mcp. The MCP surface (nine task-shaped tools, frozen in docs/tool-contracts.md), the storage Protocol layer, the sklearn-style retrieval pipeline, the reference graph, and the hybrid BM25 + dense retrieval stack have all shipped — every extension hook listed below is live in the current codebase.
Purpose: a reference menu for picking “test-the-architecture” work. Use it to propose small follow-up PRs that exercise one extension point end-to-end and validate that the abstractions hold up in practice.
Architecture snapshot¶
┌──────────────────────────────────────────────────────────────────┐
│ MCP handlers (server.py) + CLI (__main__.py) │
└──────────────────────────┬───────────────────────────────────────┘
│
▼
┌──────────────────────────────────────────────────────────────────┐
│ application/ — Application Services │
│ IndexingService, ProjectIndexer │
│ DocsSearch, ApiSearch, PackageLookup, ModuleInspector │
└──────────────────────────┬───────────────────────────────────────┘
│
▼
┌──────────────────────────────────────────────────────────────────┐
│ retrieval/ — sklearn-style Pipeline + Step ABC + formatters │
│ RetrieverStep (ABC), RetrieverPipeline, RetrieverState │
│ ChunkFetcherStep, BM25ScorerStep, MemberFetcherStep, │
│ DenseFetcherStep, DenseScorerStep, PreFilterStep, │
│ TopKFilterStep, MetadataPostFilterStep, LimitStep, │
│ TokenBudgetStep, RouteStep, ConditionalStep, ParallelStep, │
│ RRFFusionStep, GraphExpandStep, CentralityPriorStep, │
│ CommunityDiversityStep, LateInteractionScorerStep, │
│ LlmTreeReasoningStep, WeightedScoreInterpolationStep │
│ ResultFormatter, PredicateRegistry │
│ AppConfig (pydantic-settings + YAML) │
└──────────────────────────┬───────────────────────────────────────┘
│
▼
┌──────────────────────────────────────────────────────────────────┐
│ storage/ — Ports (Protocols) + Adapters (concrete) │
│ ChunkStore, PackageStore, ModuleMemberStore, │
│ DocumentTreeStore, ReferenceStore, Embedder, ResultFuser │
│ TextSearchable, VectorSearchable, HybridSearchable │
│ UnitOfWork, ConnectionProvider │
│ Filter tree + FilterFormat + FilterAdapter + MetadataSchema │
│ Sqlite* + TurboQuantStore adapters │
│ SearchBackend / SqliteCompositeBackend capability factory │
│ SqliteUnitOfWork + TurboQuantUnitOfWork + CompositeUnitOfWork │
└──────────────────────────┬───────────────────────────────────────┘
│
▼
┌──────────────────────────────────────────────────────────────────┐
│ db.py + src/lib.rs │
└──────────────────────────────────────────────────────────────────┘
Every layer boundary is a Protocol; every swappable component has a registry.
A. Storage backends¶
Each new backend implements the Protocol combinations it supports. Zero modification of existing code.
Extension |
What to implement |
Estimated size |
|---|---|---|
|
|
~400 LOC |
|
|
~300 LOC |
|
|
~400 LOC |
|
|
~400 LOC |
|
|
~300 LOC |
|
|
~400 LOC |
|
Alternative dense backend keeping vectors in SQLite (via |
~150 LOC |
|
|
~300 LOC |
Wiring: instantiate in server.py startup; reference by config path.
B. Filter system¶
Extension |
What to implement |
|---|---|
New filter operator in |
Extend |
New filter tree node type (e.g., |
Frozen dataclass + |
New user-facing filter format (e.g., |
Class implementing |
New backend filter adapter |
Class implementing |
C. Retrieval pipeline¶
Every primitive subclasses RetrieverStep (ABC) and uses registry + decorator. One class + one @step_registry.register("name") line + optional YAML preset. (step_registry holds RetrieverStep subclasses keyed by their YAML type: string. The extraction/ pipeline has its own separate stage_registry — see CLAUDE.md §”Naming: retrieval vs ingestion pipelines”.)
Extension |
Register via |
|---|---|
New pipeline step ( |
|
New fetcher / scorer step ( |
|
New formatter ( |
|
New conditional predicate ( |
|
New fusion algorithm ( |
|
New pipeline blueprint |
YAML file under |
D. Configuration¶
Extension |
Where |
|---|---|
New |
|
New metadata schema (retriever-category field allowlist) |
|
New pipeline route |
Add entry under |
New env-var override |
Already wired: |
E. Domain model¶
Extension |
Impact |
|---|---|
Add a field to |
Additive — old callers keep working; schema DDL in |
New |
Add to enum; |
New domain entity (e.g., |
New Protocol + concrete repository + new DDL in |
New operator in |
Parser branch + |
F. MCP / CLI surface¶
Extension |
Where |
|---|---|
New MCP tool |
|
New CLI subcommand |
|
Cost reference (LOC to add each extension)¶
Extension |
LOC |
Files touched |
Existing code changes |
|---|---|---|---|
New pipeline step |
~30 |
1 new class subclassing |
None |
New predicate |
~3 |
Existing |
None |
New formatter |
~30 |
1 new class |
None |
New fetcher / scorer step |
~60 |
1 new class subclassing |
None |
New pipeline preset (YAML) |
~20 |
1 new YAML file ( |
None |
New filter operator in |
~10 |
|
Backends that don’t support it keep raising |
New filter tree node type |
~10 |
|
Every |
New filter format ( |
~80 |
1 new class + registry entry |
None |
New vector-store backend |
~400 |
1 new file in |
|
New |
~80 |
1 new file |
|
New |
~3 |
|
None |
New MCP tool |
~40 |
|
None |
Known coupling limits (deliberate)¶
These are the walls — the current architecture accepts them as trade-offs rather than walling off.
Thing |
Why it’s coupled |
When to revisit |
|---|---|---|
Schema DDL |
Lives in |
A future PR may extract a |
Rust-side |
Defined in |
Only parser output; |
|
Cross-table JOIN filters would need a |
Tracked as a future extension; see the filter-adapter docstring |
Cross-backend atomicity |
Physics — no single backend supports all stores’ native transactions |
|
|
Multi-query (find A and B simultaneously) requires a different model |
Additive: add a |
Three entity types ( |
Adding a fundamentally new entity requires a new Protocol + repository + schema |
Standard extension work, just scope |
Suggested test extensions — picks for future follow-up work¶
Each is small enough to land in one focused PR and exercises the abstractions end-to-end. Ordered by complexity.
Tier 1 — tiny PRs (~50–150 LOC)¶
Add
CompactMarkdownFormatter— aResultFormatterthat renders matches more densely for token-constrained MCP responses. Exercises theformatter_registry+ YAML preset path.Add
is_code_like_querypredicate — heuristic that detectsself.,.__, parentheses in the query terms. Exercises the@predicatedecorator +RouteStepcomposition.Add
TimedStep(inner)— a decorator step that logs duration. Exercises the uniformRetrieverStepABC + the “compound step” pattern (wraps another step). Pipeline-IS-a-Step composition means it can wrap a whole sub-pipeline.Add
JsonResultFormatter— output as JSON for tooling that wants structured results. Pair with ajson_chunkYAML preset.[SHIPPED]
WeightedScoreInterpolationStep— seepython/pydocs_mcp/retrieval/steps/weighted_score_interpolation.py. Alternative fusion to the shippedRRFFusionStep. Normalizes each branch’s scores to[0, 1](min-max), then blends viaα·norm(bm25) + (1-α)·norm(dense). RRF discards score magnitude; this preserves it, which sometimes wins when one retriever is dramatically stronger than the other on a given query. Reads from the samestate.scratch[<branch>.ranked]keys RRF uses, so it drops in as a YAML swap. Pairs naturally withConditionalStepfor per-query-type routing (e.g., RRF for short queries, weighted interpolation for long).
Tier 2 — small PRs (~200–400 LOC)¶
Add
FilterTreeFormat— full dict-form filter with$and/$or/$not. Lights upAny_andNotin theSqliteFilterAdapter. First non-multifield format; validates the two-format architecture.Add
TryStep(inner, on_error=None)— the first error-tolerance primitive. Exercises step-wrapper composition + the “steps propagate exceptions” contract.Add
FieldRange(field, lo, hi)filter node — validates that the Filter tree extends cleanly;SqliteFilterAdaptergains aBETWEENtranslation. Would support queries like “chunks indexed in the last 7 days” once a timestamp field lands.Add
CachingStep(inner, cache)with a simple LRU — decorator that memoizes step output on query hash. Validates the compound-step pattern under real load.
Tier 3 — medium PRs (~400–800 LOC)¶
N. Add capability-aware ingestion (REQUIRES declarations + auto-derivation) — couple the ingestion pipeline to retrieval needs without coupling code. Each RetrieverStep gains a class-level REQUIRES: ClassVar[frozenset[str]] declaring what storage shapes it reads at query time (e.g., BM25ScorerStep → {"chunks", "chunks_fts"}; DenseScorerStep → {"chunks", "chunk_embeddings"}; LlmTreeReasoningStep → {"chunks", "document_trees"}). A new derive_ingestion_capabilities(retrieval_yaml_path) walks the active retrieval YAML, unions every step’s REQUIRES, and build_ingestion_pipeline() conditionally assembles stages based on the result — FlattenStage runs only when "chunks" is needed, EmbedChunksStage only when "chunk_embeddings" is needed, etc.
Wins:
- **Zero-mismatch by construction** — ingestion automatically produces exactly the storage shapes the active retrieval pipeline will read. Switching retrieval YAMLs triggers a re-index only when the new YAML's `REQUIRES` is a strict superset. No more "I set `tree_only.yaml` and BM25 returns nothing" support burden.
- **Tree-only deployments save real money** — skipping `EmbedChunksStage` zeros out the FastEmbed ONNX inference (or OpenAI embedding spend) at index time. On a 100-dep project this is the difference between "indexes in 30s" and "indexes in 5 minutes + costs ~$2 on OpenAI" — a meaningful unlock for users who go all-in on LLM tree reasoning.
- **Self-documenting** — reading any step's `REQUIRES = frozenset({...})` line tells you exactly which storage shapes it consumes. New contributors get capability wiring right without reading the ingestion pipeline.
- **Composes with `pipeline_hash`** (shipped in the chunk-cache work) — extend the hash input to include the capability set so switching retrieval profiles auto-invalidates the index when the new profile needs strictly more storage. No manual `--force` required after a profile flip.
- **Forces honest step design** — a step that secretly reads from `uow.chunks` but doesn't declare it in `REQUIRES` becomes a code-review issue. Encourages single-responsibility.
Implementation notes:
- Add `REQUIRES: ClassVar[frozenset[str]] = frozenset()` to the `RetrieverStep` ABC (default empty → backward-compatible).
- Override `REQUIRES` on every shipped step (~15 step classes, ~1 line each).
- Add `derive_ingestion_capabilities(yaml_path)` in `extraction/factories.py` (~50 LOC + tests).
- Refactor `build_ingestion_pipeline` to conditionally assemble stages from the capability set (~50 LOC).
- Extend `compute_ingestion_pipeline_hash` to include `sorted(capabilities)` in the hash input (~10 LOC).
- Lint rule (or test) that fails when a step's `run()` reads `state.scratch` / `uow.X` without declaring the matching capability (catches regressions).
Pairs well with the `LlmTreeReasoningStep` PR — once that lands, this is the natural follow-up that makes `tree_only.yaml` deployments stop paying for embeddings they never use.
Add
SqliteVecVectorStore— alternative dense backend that keeps vectors in SQLite (viasqlite-vec) instead of the shipped TurboQuant.tqsidecar. The dense plumbing (Chunk.embedding,EmbedderProtocol,FastEmbedEmbedder,OpenAIEmbedder,DenseScorerStep,DenseFetcherStep,TurboQuantVectorStorewired via theSearchBackendseam) already exists; this swaps the storage layer to make.dbself-contained. Useful for deployments that prefer a single-file index.Add
QdrantVectorStorewith the fullChunkStore + TextSearchable + VectorSearchable + HybridSearchablestack. First full backend swap — validates that nothing above the storage layer changes.Add
ChromaVectorStorewithChunkStore + VectorSearchableonly — tests that Option C (split Protocols) handles the “not all backends support every capability” case cleanly. Demonstrates the type-checker-level protection against wiring a BM25 fetcher step to a Chroma store.Add
LlmRerankStepwith an OpenAI/Anthropic/Cohere client — validates that an I/O-bound step composes with the pipeline. Tests predicate-guarded execution (skip rerank on short queries). The planned LLM-rerank step listed in §C “Retrieval pipeline”.[SHIPPED]
LlmTreeReasoningStep(vectorless RAG, PageIndex-style) — seepython/pydocs_mcp/retrieval/steps/llm_tree_reasoning.py. Vectorless RAG over__project__DocumentNodetrees. Three opt-in YAML presets ship underpython/pydocs_mcp/pipelines/:tree_only.yaml,chunk_search_with_tree_reasoning_parallel.yaml,chunk_search_with_tree_reasoning_after.yaml.LlmClientProtocol +OpenAiLlmClientconcrete atpython/pydocs_mcp/retrieval/llm_clients/. Two Jinja2 prompt templates (pageindex_v1 baseline + pydocs_v1 adapted) atpython/pydocs_mcp/retrieval/prompts/. ARetrieverStepthat uses an LLM to navigate the existingDocumentNodetrees and pick the nodes most likely to contain the answer, without embedding. Single-shot prompt: serialize the tree viaDocumentNode.to_pageindex_json()(strip body text, keep titles + summaries + node_ids), send(query, tree_json)to a configured LLM, parse{"thinking", "node_list": [node_id, ...]}from the response, then fetch the corresponding chunks viauow.chunksand emit them as the step’s result.Composes with the existing pipeline machinery:
In parallel with hybrid —
ParallelStepbranches: one runs BM25 + dense + RRF, the other runs tree reasoning; downstreamRRFFusionStep(orWeightedScoreInterpolationStep, registered asweighted_score_interpolation) fuses bybranch_keys=("hybrid.ranked", "tree.ranked").After hybrid —
ConditionalSteptriggers tree reasoning only on long or structural queries (is_long_querypredicate), using the hybrid result as a candidate filter.Standalone — a YAML preset that skips dense entirely; useful for long-doc corpora where TOC carries signal (PageIndex cites 98.7% on FinanceBench).
Reuses the storage layer already in place:
BuildContext.uow_factorythreads through, the step opensasync with uow_factory()and readsuow.trees(same pattern asLoadExistingChunkHashesStagefrom the chunk-cache work).New abstractions added by this PR:
LlmClientProtocol instorage/protocols.pymirroringEmbedderbut exposing BOTHasync chat()andchat_sync()— LLM calls surface in more contexts than embedding calls. First concrete:OpenAiLlmClientusingopenai>=1.40(already a required dep, no new dependencies). New providers (Anthropic, Gemini, LiteLLM) land as one-file additions per SOLID open/closed.LlmConfigsub-model inretrieval/config/embedder_models.pymirroringEmbeddingConfig(provider: Literal["openai", ...],model_name,temperature,max_tokens).Two Jinja2 prompt templates under
python/pydocs_mcp/retrieval/prompts/:tree_reasoning_pageindex_v1.j2(verbatim PageIndex baseline) andtree_reasoning_pydocs_v1.j2(adapted for code-doc queries). Prompts are versioned (_vNsuffix); selected at runtime via aprompt_templatedataclass field on the step. Versioned prompts make A/B comparison and rollback a YAML edit instead of a code change.
Inspired by VectifyAI/PageIndex (MIT). Re-implemented locally — no
pageindexpackage dependency. The single-shot algorithm is small (one Jinja2 prompt + onejson.loads+ one chunk fetch through the existinguow.chunks); vendoring the logic ≈30 LOC, avoids growing the install surface (already +90 MB after the FastEmbed/OpenAI promotion), and keeps prompt versioning + reproducibility in our repo instead of someone else’s release schedule. PageIndex isn’t on PyPI underpageindexanyway (pageindex-rsis an unrelated Rust port).
Tier 4 — larger PRs (~800+ LOC)¶
Add
CompositeUnitOfWork— coordinates multiple backend UoWs with best-effort rollback. Enables heterogeneous setups (e.g., Qdrant chunks + SQLite packages). Tests the heterogeneous-backends story.Add
HyDERetriever(Hypothetical Document Embeddings) — asks an LLM to draft a hypothetical answer, embeds it, then does vector search. Exercises retriever composition (LLM client + embedder + vector store).Add config-driven pipeline reloading — watch the YAML file; rebuild pipelines on change without restart. Tests the serialization boundary (dict round-trip).
How to propose a test-extension PR¶
Pick one extension from the list above (or a combination).
Brainstorm the spec in a session: clarifying questions → design sections → spec file under
docs/superpowers/specs/YYYY-MM-DD-<topic>.md. Match the shape of the existing specs in that directory.The spec template (see existing specs for reference) covers: Goal, Decisions, Scope, Domain components, Files touched, Risks, Acceptance criteria, Open items.
Acceptance criteria should include at least one test that verifies the extension integrates at a seam (e.g., “a pipeline YAML referencing the new stage round-trips through
to_dict/from_dictand executes end-to-end against a fixture”).