Multiple datasets for a ContentRetriever and EmbeddingVectorStore

[thomas-reimonn]

I'd like to be able to create multiple datasets (or topics) in an EmbeddingVectorStore and retrieve content filtered by the topic. I've read through the examples and interfaces of ContentRetriever and EmbeddingVectorStore and I don't see much, except maybe a Filter on metadata... Do you have any suggestions?

[langchain4j]

Hi @thomas-reimonn yes, Filter is probably what you are looking for, if you want to have everything in the same EmbeddingStore.

Another option would be to have a separate EmbeddingStore/EmbeddingStoreContentRetriever for each topic and use QueryRouter to route to the appropriate EmbeddingStoreContentRetriever depending on a user query. Here is an example.

[thomas-reimonn]

Thanks for sharing those examples. I think it would work but wouldn't be quite right. The use case I'm trying to solve for is: I have many customers and I have (a smaller amount) of data about each customer. I want to generate customer-specific responses and it's critical that nothing be shared between customers.

Would you be open to a contribution of something like an EmbeddingStoreContentRetrieverFactory that produces EmbeddingStoreContentRetrievers, each scoped to its own fields, e.g. customer_id, tenant_id, etc?

[sfotex]

I'm facing the same issue. I would be happy to pair up and get something knocked out.

[langchain4j]

If you store customer_id and tenant_id in the metadata of each TextSegment, you can filter by it using dynamicFilter configuration of the EmbeddingStoreContentRetriever. There is an example here: #610

You could resolve customer_id/tenant_id from memoryId. Would that work for you?

[thomas-reimonn]

I think so! I think this approach should work. I have an example below. I'm still getting cross-retrieval (e.g. If I ask the Dinosaur chat about planets, it gets planet facts and if I ask the Planet topic about Dinosaurs I get Dino facts). I like the DynamicFilter suggestion and I think this could be refactored to use one.

It looks like the EmbeddingStores may need some enhancements to handle these filters well in the future -- indexing on the Metadata column in Postgres, either including it in the primary key or cluster column in Cassandra.

@Transactional
    private void embedTextsForTopic(UUID topicId) {
        // Map all topics.texts to TextSegments and add them to the embedding store if they are not
        // already present. Include topic_id and text_id as metadata.
        logger.debug("Embedding texts for topic: {}", topicId);
        Topic topic =
                topicRepository
                        .findById(topicId)
                        .orElseThrow(
                                () ->
                                        new ResourceNotFoundException(
                                                "Topic not found for ID: " + topicId));

        topic.getTexts()
                .forEach(
                        text -> {
                            if (!text.isEmbedded()) {
                                Map<String, String> metadata = new HashMap<>();
                                metadata.put("topic_id", topic.getId().toString());
                                metadata.put("text_id", text.getId().toString());
                                TextSegment textSegment =
                                        TextSegment.from(text.getContent(), new Metadata(metadata));
                                Response<Embedding> embedding = embeddingModel.embed(textSegment);
                                embeddingStore.add(embedding.content(), textSegment);
                                text.setEmbedded(true);
                                logger.debug(
                                        "Embedded Text(id={}, title={}, content={})",
                                        text.getId(),
                                        text.getTitle(),
                                        text.getContent());
                                textRepository.save(text);
                            }
                        });
    }

    public String answerQuery(UUID topicId, UUID memoryId, String query) {
        // Handles the route /topics/{topicId}/chats/{memoryId}?message={query}
        logger.info("Answering query: {} for topic: {} using memory: {}", query, topicId, memoryId);
        embedTextsForTopic(topicId); // Embed any texts about a topic if they are not already embedded
        
        Filter topicFilter = metadataKey("topic_id").isEqualTo(topicId.toString());
        ContentRetriever contentRetriever =
                EmbeddingStoreContentRetriever.builder()
                        .embeddingStore(embeddingStore) // A PgVectorEmbeddingStore
                        .embeddingModel(embeddingModel) // A BgeSmallEnV15QuantizedEmbeddingModel
                        .filter(topicFilter)
                        .build();

        ChatMemoryProvider chatMemoryProvider =
                memId ->
                        MessageWindowChatMemory.builder()
                                .id(memId)
                                .maxMessages(10)
                                .chatMemoryStore(persistentChatMemoryStore) // Hibernate/Postgres store
                                .build();

        Assistant assistant =
                AiServices.builder(Assistant.class)
                        .chatLanguageModel(
                                OpenAiChatModel.withApiKey(System.getenv("OPENAI_API_KEY")))
                        .chatMemoryProvider(chatMemoryProvider)
                        .contentRetriever(contentRetriever)
                        .build();

        return assistant.chat(memoryId, query);
    }

[langchain4j]

@thomas-reimonn I am afraid PgVectorEmbeddingStore does not support Filtering in the current (0.30.0) version. This is probably the reason why you see cross-retrieval. #851 is merged and will be released next week in 0.31.0 version. But you can use it already in 0.31.0-SNAPSHOT. Or use Milvus/Elasticsearch which support Filtering in 0.30.0.

Regarding indexing, it is already supported in that PR.