import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.embedding.AllMiniLmL6V2EmbeddingModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.MemoryId;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.filter.Filter;
import dev.langchain4j.store.embedding.filter.builder.sql.LanguageModelSqlFilterBuilder;
import dev.langchain4j.store.embedding.filter.builder.sql.TableDefinition;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import org.junit.jupiter.api.Test;

import java.util.function.Function;

import static dev.langchain4j.data.document.Metadata.metadata;
import static dev.langchain4j.store.embedding.filter.MetadataFilterBuilder.metadataKey;
import static org.assertj.core.api.Assertions.assertThat;

class _06_Metadata_Filtering {

    /**
     * More information can be found here: https://github.com/langchain4j/langchain4j/pull/610
     */

    ChatLanguageModel chatLanguageModel = OpenAiChatModel.builder()
            .apiKey(System.getenv("OPENAI_API_KEY"))
            .build();

    EmbeddingModel embeddingModel = new AllMiniLmL6V2EmbeddingModel();

    interface Assistant {

        String answer(String query);
    }

    @Test
    void Static_Metadata_Filter_Example() {

        // given
        TextSegment dogsSegment = TextSegment.from("Article about dogs ...", metadata("animal", "dog"));
        TextSegment birdsSegment = TextSegment.from("Article about birds ...", metadata("animal", "bird"));

        EmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
        embeddingStore.add(embeddingModel.embed(dogsSegment).content(), dogsSegment);
        embeddingStore.add(embeddingModel.embed(birdsSegment).content(), birdsSegment);
        // embeddingStore contains segments about both dogs and birds

        Filter onlyDogs = metadataKey("animal").isEqualTo("dog");

        ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .filter(onlyDogs) // by specifying the static filter, we limit the search to segments only about dogs
                .build();

        Assistant assistant = AiServices.builder(Assistant.class)
                .chatLanguageModel(chatLanguageModel)
                .contentRetriever(contentRetriever)
                .build();

        // when
        String answer = assistant.answer("Which animal?");

        // then
        assertThat(answer)
                .containsIgnoringCase("dog")
                .doesNotContainIgnoringCase("bird");
    }


    interface PersonalizedAssistant {

        String chat(@MemoryId String userId, @dev.langchain4j.service.UserMessage String userMessage);
    }

    @Test
    void Dynamic_Metadata_Filter_Example() {

        // given
        TextSegment user1Info = TextSegment.from("My favorite color is green", metadata("userId", "1"));
        TextSegment user2Info = TextSegment.from("My favorite color is red", metadata("userId", "2"));

        EmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
        embeddingStore.add(embeddingModel.embed(user1Info).content(), user1Info);
        embeddingStore.add(embeddingModel.embed(user2Info).content(), user2Info);
        // embeddingStore contains information about both first and second user

        Function<Query, Filter> filterByUserId =
                (query) -> metadataKey("userId").isEqualTo(query.metadata().chatMemoryId().toString());

        ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                // by specifying the dynamic filter, we limit the search to segments that belong only to the current user
                .dynamicFilter(filterByUserId)
                .build();

        PersonalizedAssistant personalizedAssistant = AiServices.builder(PersonalizedAssistant.class)
                .chatLanguageModel(chatLanguageModel)
                .contentRetriever(contentRetriever)
                .build();

        // when
        String answer1 = personalizedAssistant.chat("1", "Which color would be best for a dress?");

        // then
        assertThat(answer1)
                .containsIgnoringCase("green")
                .doesNotContainIgnoringCase("red");

        // when
        String answer2 = personalizedAssistant.chat("2", "Which color would be best for a dress?");

        // then
        assertThat(answer2)
                .containsIgnoringCase("red")
                .doesNotContainIgnoringCase("green");
    }

    @Test
    void LLM_generated_Metadata_Filter_Example() {

        // given
        TextSegment forrestGump = TextSegment.from("Forrest Gump", metadata("genre", "drama").put("year", 1994));
        TextSegment groundhogDay = TextSegment.from("Groundhog Day", metadata("genre", "comedy").put("year", 1993));
        TextSegment dieHard = TextSegment.from("Die Hard", metadata("genre", "action").put("year", 1998));

        // describe metadata keys as if they were columns in the SQL table
        TableDefinition tableDefinition = TableDefinition.builder()
                .name("movies")
                .addColumn("genre", "VARCHAR", "one of: [comedy, drama, action]")
                .addColumn("year", "INT")
                .build();

        LanguageModelSqlFilterBuilder sqlFilterBuilder = new LanguageModelSqlFilterBuilder(chatLanguageModel, tableDefinition);

        EmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
        embeddingStore.add(embeddingModel.embed(forrestGump).content(), forrestGump);
        embeddingStore.add(embeddingModel.embed(groundhogDay).content(), groundhogDay);
        embeddingStore.add(embeddingModel.embed(dieHard).content(), dieHard);

        ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .dynamicFilter(query -> sqlFilterBuilder.build(query)) // LLM will generate the filter dynamically
                .build();

        Assistant assistant = AiServices.builder(Assistant.class)
                .chatLanguageModel(chatLanguageModel)
                .contentRetriever(contentRetriever)
                .build();

        // when
        String answer = assistant.answer("Recommend me a good drama from 90s");

        // then
        assertThat(answer)
                .containsIgnoringCase("Forrest Gump")
                .doesNotContainIgnoringCase("Groundhog Day")
                .doesNotContainIgnoringCase("Die Hard");
    }
}

dev.langchain4j.store.embedding.filter.builder.sql.TableDefinition.builder dev.langchain4j.service.AiServices.builder dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever.builder dev.langchain4j.model.openai.OpenAiChatModel.builder
[((1356, 1455), 'dev.langchain4j.model.openai.OpenAiChatModel.builder'), ((1356, 1434), 'dev.langchain4j.model.openai.OpenAiChatModel.builder'), ((2328, 2607), 'dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever.builder'), ((2328, 2498), 'dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever.builder'), ((2328, 2464), 'dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever.builder'), ((2328, 2416), 'dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever.builder'), ((2640, 2806), 'dev.langchain4j.service.AiServices.builder'), ((2640, 2781), 'dev.langchain4j.service.AiServices.builder'), ((2640, 2729), 'dev.langchain4j.service.AiServices.builder'), ((3993, 4323), 'dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever.builder'), ((3993, 4298), 'dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever.builder'), ((3993, 4129), 'dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever.builder'), ((3993, 4081), 'dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever.builder'), ((4380, 4558), 'dev.langchain4j.service.AiServices.builder'), ((4380, 4533), 'dev.langchain4j.service.AiServices.builder'), ((4380, 4481), 'dev.langchain4j.service.AiServices.builder'), ((5624, 5830), 'dev.langchain4j.store.embedding.filter.builder.sql.TableDefinition.builder'), ((5624, 5805), 'dev.langchain4j.store.embedding.filter.builder.sql.TableDefinition.builder'), ((5624, 5763), 'dev.langchain4j.store.embedding.filter.builder.sql.TableDefinition.builder'), ((5624, 5681), 'dev.langchain4j.store.embedding.filter.builder.sql.TableDefinition.builder'), ((6344, 6620), 'dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever.builder'), ((6344, 6551), 'dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever.builder'), ((6344, 6480), 'dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever.builder'), ((6344, 6432), 'dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever.builder'), ((6653, 6819), 'dev.langchain4j.service.AiServices.builder'), ((6653, 6794), 'dev.langchain4j.service.AiServices.builder'), ((6653, 6742), 'dev.langchain4j.service.AiServices.builder')]