diff --git a/vector_store.py b/vector_store.py index e18955b..ba0a8b3 100644 --- a/vector_store.py +++ b/vector_store.py @@ -1,26 +1,42 @@ -from langchain_community.vectorstores import Chroma +from langchain_qdrant import QdrantVectorStore from langchain_ollama import OllamaEmbeddings -from langchain.schema import Document +from langchain_core.documents import Document +from qdrant_client import QdrantClient +from qdrant_client.models import Distance, VectorParams from langchain.text_splitter import RecursiveCharacterTextSplitter +import uuid -COLLECTION_NAME = "rag_collection" +COLLECTION_NAME = "knowledge_base" EMBEDDING_MODEL = "nomic-embed-text" -CHROMA_PERSIST_DIR = "./chroma_db" +VECTOR_SIZE = 768 -class ChromaStore: - def __init__(self, collection_name: str = COLLECTION_NAME): - self.client = Chroma( - embedding_function=OllamaEmbeddings(model=EMBEDDING_MODEL), - collection_name=collection_name, - persist_directory=CHROMA_PERSIST_DIR, +def get_embeddings(): + return OllamaEmbeddings(model=EMBEDDING_MODEL) + + +def get_qdrant_client(): + return QdrantClient(host="localhost", port=6333) + + +def init_collection(client: QdrantClient): + collections = [c.name for c in client.get_collections().collections] + if COLLECTION_NAME not in collections: + client.create_collection( + collection_name=COLLECTION_NAME, + vectors_config=VectorParams(size=VECTOR_SIZE, distance=Distance.COSINE), ) - def add_documents(self, docs: list[Document]) -> None: - self.client.add_documents(docs) - def search(self, query: str, limit: int = 5) -> list[Document]: - return self.client.similarity_search(query, k=limit) +def get_vector_store() -> QdrantVectorStore: + client = get_qdrant_client() + init_collection(client) + embeddings = get_embeddings() + return QdrantVectorStore( + client=client, + collection_name=COLLECTION_NAME, + embedding=embeddings, + ) def add_documents(content: str, title: str) -> int: @@ -37,20 +53,21 @@ def add_documents(content: str, title: str) -> int: ) for i, chunk in enumerate(chunks) ] - store = ChromaStore() + store = get_vector_store() store.add_documents(docs) return len(docs) def search_documents(query: str, max_results: int = 5) -> list[dict]: - store = ChromaStore() - results = store.search(query, limit=max_results) + store = get_vector_store() + results = store.similarity_search_with_relevance_scores(query, k=max_results) output = [] - for doc in results: + for doc, score in results: output.append( { "content": doc.page_content, "metadata": doc.metadata, + "score": round(score, 4), } ) return output \ No newline at end of file