Feature (change?) request: Allow searching with embedding

[joan-teriihoania]

Hello @techiejd,

Note

First time here. Looking for ways to contribute. Also, I'm not that familiar with embeddings. I know the basics, but do let me know if what I'm saying isn't feasible. Or just plain doesn't work.

At the moment, vectorSearch only takes a query string which is itself converted into an embedding. Would it make sense to allow embedding search instead of only string, at least for internal uses (I can see why passing a vector in a REST call may not be the best). So maybe a query type like this:

type Query = string | ReturnType<EmbedQueryFn>

And instead of:

const queryEmbedding = await (async () => {
  const qE = await poolConfig.embeddingConfig.queryFn(query)
  return Array.isArray(qE) ? qE : Array.from(qE)
})()

We have:

const queryEmbedding = Array.isArray(query)
  ? await query
  : await (async () => {
      const qE = await poolConfig.embeddingConfig.queryFn(query)
      return Array.isArray(qE) ? qE : Array.from(qE)
    })()

We could even allow passing an embedding ID altogether since knowledge pools are collections themselves, so retrieving a document within the vector search would fall under their scope. Although, we would need to expose a getEmbedding function to make the process easier.

Let me know your thoughts!

[techiejd]

Hey @joan-teriihoania, thanks for opening this! It's genuinely appreciated.

First, I'd love to hear what use cases pushed you toward this:

• Are you doing more-like-this on stored docs?
• Reusing one embedding across multiple pools?
• HyDE or some other query transform?

Knowing the actual shape of what you're trying to do helps me make sure whatever API we add actually fits (and, also, because it's a morale booster for me!).

On the design: your request surfaced a confusing API naming (I'll deal with this some other time).

The current search is a convenience workhorse. It does a lot more than searching:

1. It embeds the query string,
2. runs the vector search via adapter.search,
3. and optionally reranks.

That bundling is great for the common case and I want to keep it. What's missing is an escape hatch for callers who already have an embedding, because adapter.search is buried inside the pipeline with no public way to reach it.

Overloading search (widening the query field to mean "string OR vector") is confusing. If you're searching by embedding, you will skip embedding the query string. If you're searching by embedding, you cannot rerank, because cross-encoder rerankers (Cohere, BGE, Jina) need text, not vectors. So searching by embedding skips 2/3rds of the search path and it's not obvious.

Rather than overloading, let's add a sibling method which can be exposed on the (internal only) VectorizedPayload:

payload.searchByEmbedding({
  knowledgePool,
  embedding,
  where,
  limit,
})

It's a thin wrapper over adapter.search.

That naming also opens the door to a searchBy* family. Your second idea (passing an embedding ID) becomes a natural sibling: searchById({ knowledgePool, ... }).

What do you think? And please do share the use cases.

[joan-teriihoania]

Thanks for the response!

First, I'd love to hear what use cases pushed you toward this: Are you doing more-like-this on stored docs?

Yes. I am trying to build a knowledge base from documentation files (mostly PDF files). I'm thinking of having a media collection, extracting the text within uploaded documents, and using that for the embedding generation. The idea is that AI agents and users can make searches through the KB using keywords, but also by selecting documents and searching for related ones.

The usecase goal is to give contextual information to the users and/or agents "passively" (without proper keywords).

Example: An agent has the emergency procedures documentation for a desalination plant in Town A open. Next to the document viewer, we might see recommendations for the plans of the plant, evacuation zones of Town A area, etc... Basically, whatever vector surrounds the document in the vector space.

I was initially thinking of extracting keywords from the "source" document and making multiple search queries and retrieving the results. But I was curious to see if it'd be better to make a vector search using the source itself rather than with snippets of it.

Overloading search (widening the query field to mean "string OR vector") is confusing. If you're searching by embedding, you will skip embedding the query string. If you're searching by embedding, you cannot rerank, because cross-encoder rerankers (Cohere, BGE, Jina) need text, not vectors. So searching by embedding skips 2/3rds of the search path and it's not obvious.

Rather than overloading, let's add a sibling method which can be exposed on the (internal only) VectorizedPayload:

I didn't think of that. That's true.

That naming also opens the door to a searchBy* family. Your second idea (passing an embedding ID) becomes a natural sibling: searchById({ knowledgePool, ... }).

That's actually way better than overloading the search method. And it's way cleaner.

What do you think? And please do share the use cases.

I do agree. This seems to fulfil the requirements I have. Plus, it makes the search API clear. I am curious of how this would be implemented and the quality of the results (but that's DB-side). Quick question: Is there already a built-in way to retrieve an existing embedding using their ID? Or the ID of their "source" (eg. the ID in the posts collection). If so, then maybe I could watch how you implement searchByEmbedding and follow up with an implementation of searchById?

[techiejd]

Quick question: Is there already a built-in way to retrieve an existing embedding using their ID?

Not generically, which is what makes this so complicated. This getEmbedding is much more similar to 'payload.find' than it is to an actual search given that it's a predicate based retrieval.

• If you're using postgres, you could simply use a payload.findByID.
• If you're using mongodb, the embeddings live in a raw MongoDB collection (not a Payload one), so payload.findByID won't see them. You'd have to reach into payload.db and query the configured collectionName directly via the native driver.
• If you're using cf, the vectors don't even live in your database. They're in Cloudflare's Vectorize binding. There's a Payload mapping collection (CF_MAPPINGS_SLUG) that ties (pool, sourceCollection, docId) to vector ids, so you'd query that first, then call vectorizeBinding.getByIds([...]) against the binding to fetch the actual vectors.

Or the ID of their "source" (eg. the ID in the posts collection).

No. For the same reason why there is no built-in way to retrieve an existing embedding using their ID. Each db adapter is distinct in how they store their vectors. Each db adapter will need to expose a way to do a predicate based retrieval.

Given what we just discussed, I propose the following:

• We keep searchByEmbedding just like we discussed--it's a simple wrapper over adapter.search.
• We ditch searchById. Instead we add a new function DbAdapter.find and a new function VectorizedPayload.find which is a simple wrapper over adapter.find. All db adapters will have to implement find which is a predicate based retrieval. Find will have a similar shape to payload's find/findById:
  adapter.findByID({ poolName, id }) → returns one EmbeddingRecord | null
adapter.find({ poolName, sourceCollection?, docId?, where?, limit? }) → returns EmbeddingRecord[]
EmbeddingRecord is the same as VectorSearchResult without score.

If so, then maybe I could watch how you implement searchByEmbedding and follow up with an implementation of searchById?

searchByEmbedding is a much easier implementation because you don't need to change the adapter API. If you want a simple contribution, I strongly suggest you work on searchByEmbedding.

searchByEmbedding would be a simple wrapper over adapter.search, which you can see used here. You can add a new VectorSearchEmbeddingQuery which is refactoring these parameters, similar to VectorSearchQuery. You would add searchByEmbedding to the VectorizedPayload on the next line after search. You can see VectorizedPayload being built here. You can add a 'searchByEmbedding.spec.ts' to the specs. Which would simply embed a doc and then you do a searchByEmbedding to retrieve it; see this spec for a similar example.

There's also a natural second PR after that one:

• add populateEmbedding?: Boolean to VectorSearchQuery and the VectorSearchEmbeddingQuery that you built. populateEmbedding? is a request to populate embeddings in the results. Add embedding?: number[] to VectorSearchResult
  which holds the embedding if asked to populate.

If you're confused I'm willing to do a short video chat to go over this. If you'd like me for me to take over, please do let me know.

I am not blocked on .find and .findById and can try to find some time next week. Will have an update in 3 weeks or less.

[joan-teriihoania]

Thanks for the answers and the pointers! I'll try out an implementation and see if I have more questions. I'll probably have some time if not this week, then next week. 👍