rag-chatbot-2

1
from pathlib import Path
2

3
import chromadb
4
from bot.memory.embedder import EmbedderHuggingFace
5
from bot.memory.vector_memory import VectorMemory
6
from helpers.prettier import prettify_source
7
from langchain_community.vectorstores.chroma import Chroma
8

9
if __name__ == "__main__":
10
    root_folder = Path(__file__).resolve().parent.parent
11
    # Contains an extract of documents uploaded to the RAG bot;
12
    declarative_vector_store_path = root_folder / "vector_store" / "docs_index"
13
    # Contains an extract of things the user said in the past;
14
    episodic_vector_store_path = root_folder / "vector_store" / "episodic_index"
15

16
    embedding = EmbedderHuggingFace().get_embedding()
17
    index = VectorMemory(vector_store_path=str(declarative_vector_store_path), embedding=embedding)
18

19
    # query = "<write_your_query_here>"
20
    query = "tell me a joke about ClearML"
21

22
    matched_docs, sources = index.similarity_search(query)
23

24
    for source in sources:
25
        print(prettify_source(source))
26

27
    persistent_client = chromadb.PersistentClient(path=str(episodic_vector_store_path))
28
    collection = persistent_client.get_or_create_collection("episodic_memory")
29
    collection.add(ids=["1", "2", "3"], documents=["a", "b", "c"])
30
    langchain_chroma = Chroma(
31
        client=persistent_client,
32
        collection_name="episodic_memory",
33
        embedding_function=embedding,
34
    )
35
    docs = langchain_chroma.similarity_search("a")
36
    docs_with_score = langchain_chroma.similarity_search_with_score("a")
37
    docs_with_relevance_score = langchain_chroma.similarity_search_with_relevance_scores("a")
38
    matched_doc = max(docs_with_relevance_score, key=lambda x: x[1])
39

40
    # The returned distance score is cosine distance. Therefore, a lower score is better.
41
    results = collection.query(
42
        query_texts=["a"],
43
        n_results=2,
44
        # where={"metadata_field": "is_equal_to_this"}, # optional filter
45
        # where_document={"$contains":"search_string"}  # optional filter
46
    )
47
    print(results)
48