Custom Search Tools#

Each search tool is backed by a collection — a named set of documents that has been embedded and stored in pgvector. Adding a custom search tool means defining a new collection and making the server aware of it; it then appears as a search_<name> tool alongside the built-in ones.

There are two ways to add one:

Approach

Use it when

Effort

Config only (markdown)

Your content is plain Markdown (.md) files and you want standard semantic search over them — the common “add our own docs” case.

No Python — a config entry.

Python class

You need something config can’t express: a non-Markdown loader (PDF, HTML, a database, an API), a custom retriever (hybrid/ensemble/rerank, different k or metadata filters), custom result formatting, or replacing a built-in collection.

Write a small CollectionBase subclass.

Both produce an identical search_<name> tool; they differ only in how much control you need over loading and retrieval. Start with the config approach and reach for a class only when you hit its limits.

Option 1 — Config only (markdown directories)#

If your documentation is plain Markdown, you do not need to write any Python. Declare the collection in the CUSTOM_COLLECTIONS setting: a map of tool name -> { path, description, base_url }. Each entry reuses the same Markdown loader and vector retriever the built-in documentation collections use.

Fields per entry:

path (required)

Directory of .md files, relative to the docs source path (the --source_path given to build_index). May be a single string or a list of directories to combine several into one collection.

description (required)

The tool description the model reads when deciding whether to call search_<name>. Write it for the model — it directly affects retrieval quality.

base_url (optional)

URL prefix used to build a citation link for each document. Omit it (or set null) for internal docs with no public home, and no source URL is emitted.

Steps#

  1. Put your Markdown under the docs source directory, e.g. document_sources/client-docs/.

  2. Declare the collection and add its name to COLLECTIONS so it is indexed and served. In your .env (values are JSON):

    CUSTOM_COLLECTIONS={"client-docs": {"path": "client-docs", "description": "ACME internal OneTick docs: data onboarding, symbol conventions, entitlements. Use for ACME-specific setup and data.", "base_url": null}}
COLLECTIONS=["python-docs", "python-examples", "cloud", "sql-docs", "client-docs"]

  3. Index the new collection (leaves the others untouched):

    python -m onetick_llmbase.build_index -c onetick_mcp.config \
    --source_path ./document_sources --only_collections client-docs

  4. Restart the server. search_client-docs now appears in the tool list.

To combine several directories into one collection, make path a list, e.g. "path": ["client-docs", "extra-notes"]. To add more collections, add more entries to the map and more names to COLLECTIONS.

Option 2 — Python class#

For anything beyond plain-Markdown search, define a collection by subclassing CollectionBase and registering it with the register_collection decorator.

CollectionBase#

All collections inherit from onetick_llmbase.collection.base.CollectionBase.

Key members:

description (class attribute)

The tool description the model sees when deciding which tool to call. Set this to something specific — it directly affects retrieval quality.

get_retriever(vectorstore, llm) (abstract)

Return a LangChain BaseRetriever for this collection. Called at server startup. This is where a custom retriever (hybrid/ensemble, a reranker, custom k or metadata filters) goes.

get_loader(**kwargs) (optional)

Return a document loader used at index time. This is where a non-Markdown source (PDF, HTML, a database, an API) goes. If omitted, the collection cannot be indexed.

build_mcp_tool(vectorstore, llm)

Builds the StructuredTool registered with FastMCP. The default implementation wraps the retriever with contextual compression and formats results with source URLs.

Registering a custom collection#

from onetick_llmbase.collection import register_collection
from onetick_llmbase.collection.base import CollectionBase

@register_collection("my-collection")
class MyCollection(CollectionBase):
    description = "Search internal research notes for relevant context."

    def get_retriever(self, vectorstore, llm):
        return vectorstore.as_retriever(search_kwargs={"k": 5})

    def get_loader(self, **kwargs):
        # Return a LangChain document loader for your content
        ...

The @register_collection("my-collection") decorator adds the class to the global CollectionRegistry. Import the module that defines it from your config module so the decorator runs at startup, then add "my-collection" to COLLECTIONS (and it becomes selectable via LLMBASE_MCP_TOOLS).

Note

If a collection name is already registered, it will be overwritten with a warning. This allows downstream packages to replace built-in collections.

OneTick document collections#