from pathlib import Path

import chromadb
from llama_index.core import VectorStoreIndex
from llama_index.vector_stores.chroma import ChromaVectorStore

from keyword_embedding import KeywordEmbedding


persist_dir = Path("chroma_support_store")
client = chromadb.PersistentClient(path=str(persist_dir))
collection = client.get_collection("support_runbooks")
vector_store = ChromaVectorStore(chroma_collection=collection)
index = VectorStoreIndex.from_vector_store(
    vector_store,
    embed_model=KeywordEmbedding(),
)
retriever = index.as_retriever(similarity_top_k=1)