from pathlib import Path
from phi.assistant import Assistant
from phi.knowledge.langchain import LangChainKnowledgeBase

from langchain.embeddings import OpenAIEmbeddings
from langchain.document_loaders import TextLoader
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores import Chroma

cookbook_dir = Path("__file__").parent
chroma_db_dir = cookbook_dir.joinpath("storage/chroma_db")


def load_vector_store():
    state_of_the_union = cookbook_dir.joinpath("data/demo/state_of_the_union.txt")
    # -*- Load the document
    raw_documents = TextLoader(str(state_of_the_union)).load()
    # -*- Split it into chunks
    text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
    documents = text_splitter.split_documents(raw_documents)
    # -*- Embed each chunk and load it into the vector store
    Chroma.from_documents(documents, OpenAIEmbeddings(), persist_directory=str(chroma_db_dir))


# -*- Load the vector store
load_vector_store()
# -*- Get the vectordb
db = Chroma(embedding_function=OpenAIEmbeddings(), persist_directory=str(chroma_db_dir))
# -*- Create a retriever from the vector store
retriever = db.as_retriever()

# -*- Create a knowledge base from the vector store
knowledge_base = LangChainKnowledgeBase(retriever=retriever)

conv = Assistant(knowledge_base=knowledge_base, debug_mode=True, add_references_to_prompt=True)
conv.print_response("What did the president say about technology?", markdown=True)

langchain.text_splitter.CharacterTextSplitter langchain.embeddings.OpenAIEmbeddings
[((1254, 1297), 'phi.knowledge.langchain.LangChainKnowledgeBase', 'LangChainKnowledgeBase', ([], {'retriever': 'retriever'}), '(retriever=retriever)\n', (1276, 1297), False, 'from phi.knowledge.langchain import LangChainKnowledgeBase\n'), ((1306, 1398), 'phi.assistant.Assistant', 'Assistant', ([], {'knowledge_base': 'knowledge_base', 'debug_mode': '(True)', 'add_references_to_prompt': '(True)'}), '(knowledge_base=knowledge_base, debug_mode=True,\n    add_references_to_prompt=True)\n', (1315, 1398), False, 'from phi.assistant import Assistant\n'), ((337, 353), 'pathlib.Path', 'Path', (['"""__file__"""'], {}), "('__file__')\n", (341, 353), False, 'from pathlib import Path\n'), ((672, 727), 'langchain.text_splitter.CharacterTextSplitter', 'CharacterTextSplitter', ([], {'chunk_size': '(1000)', 'chunk_overlap': '(0)'}), '(chunk_size=1000, chunk_overlap=0)\n', (693, 727), False, 'from langchain.text_splitter import CharacterTextSplitter\n'), ((887, 905), 'langchain.embeddings.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {}), '()\n', (903, 905), False, 'from langchain.embeddings import OpenAIEmbeddings\n'), ((1049, 1067), 'langchain.embeddings.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {}), '()\n', (1065, 1067), False, 'from langchain.embeddings import OpenAIEmbeddings\n')]