chore: Add new session-level service for getting embeddings of a specific collection MCP-246 #626
Conversation
While we do our best to make sure we don't break anything, there might be situations where users want to disable the validation and insert documents as they please.
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
| try { | ||
| await provider.getSearchIndexes("test", "test"); | ||
| return true; | ||
| } catch { | ||
| return false; | ||
| } |
There was a problem hiding this comment.
Should we consider that atlas local deployments might take time to boot up search management service. This might become a common case as soon as we provide local-atlas tools.
|
|
||
| switch (definition.quantization) { | ||
| case "none": | ||
| return true; |
There was a problem hiding this comment.
Its worth adding this as a comment.
| } else { | ||
| return definition; | ||
| } |
There was a problem hiding this comment.
| } else { | |
| return definition; | |
| } | |
| } | |
| return definition; |
| constructor( | ||
| private readonly config: UserConfig, | ||
| private readonly embeddings: Map<EmbeddingNamespace, VectorFieldIndexDefinition[]> = new Map(), | ||
| private readonly atlasSearchStatus: Map<string, boolean> = new Map() |
There was a problem hiding this comment.
We probably want to listen to connection changes (here or in Session) to clear the orphan entries.
| const result = await provider.insertMany(database, collection, documents); | ||
|
|
||
| const embeddingValidations = new Set( | ||
| ...(await Promise.all( |
There was a problem hiding this comment.
A small note about Promise here, the only reason findFieldsWithWrongEmbeddings returns a Promise is because of isAtlasSearchAvailable (directly / indirectly).
When we're already in a tool call, is there a possibility to assume that isAtlasSearchAvailable is already populated and work off from that?
Just to be clear - what we have is not a problem. I see a small scope of improvement because I had to re-look what async task findFieldsWithWrongEmbeddings was doing.
| function extractInsertedIds(content: string): ObjectId[] { | ||
| expect(content).toContain("Documents were inserted successfully."); | ||
| expect(content).toContain("Inserted IDs:"); | ||
|
|
||
| const match = content.match(/Inserted IDs:\s(.*)/); | ||
| const group = match?.[1]; | ||
| return ( | ||
| group | ||
| ?.split(",") | ||
| .map((e) => e.trim()) | ||
| .map((e) => ObjectId.createFromHexString(e)) ?? [] | ||
| ); | ||
| } |
There was a problem hiding this comment.
Should we move this to helpers instead? I noticed myself sometimes re-inventing the helpers even though they exist in the codebase but only because they were in specific test files.
2 participants
Proposed changes
This PR adds a new session level service called VectorSearchEmbeddings, that is responsible of:
Given that the embedding combinations and detection can be inaccurate, we also provide a new configuration option called "disableEmbeddingsValidation" that can be set up by CLI/Env and when true, the validation is bypassed.
This PR also introduces the embedding validation in the insertMany tool, so users can not randomly add data that can break
existing models or indexes unknowingly.
We depend on #628 to be merged, as it implements a method to detect if Atlas Search is available. Whenever the PR is merged, I'll refactor the method introduced there and use VectorSearchEmbeddings, so we have only one single place for search detection.
Checklist