diff --git a/vector_store.py b/vector_store.py index ba0a8b3..e18955b 100644 --- a/vector_store.py +++ b/vector_store.py @@ -1,42 +1,26 @@ -from langchain_qdrant import QdrantVectorStore +from langchain_community.vectorstores import Chroma from langchain_ollama import OllamaEmbeddings -from langchain_core.documents import Document -from qdrant_client import QdrantClient -from qdrant_client.models import Distance, VectorParams +from langchain.schema import Document from langchain.text_splitter import RecursiveCharacterTextSplitter -import uuid -COLLECTION_NAME = "knowledge_base" +COLLECTION_NAME = "rag_collection" EMBEDDING_MODEL = "nomic-embed-text" -VECTOR_SIZE = 768 +CHROMA_PERSIST_DIR = "./chroma_db" -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), +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 add_documents(self, docs: list[Document]) -> None: + self.client.add_documents(docs) -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 search(self, query: str, limit: int = 5) -> list[Document]: + return self.client.similarity_search(query, k=limit) def add_documents(content: str, title: str) -> int: @@ -53,21 +37,20 @@ def add_documents(content: str, title: str) -> int: ) for i, chunk in enumerate(chunks) ] - store = get_vector_store() + store = ChromaStore() store.add_documents(docs) return len(docs) def search_documents(query: str, max_results: int = 5) -> list[dict]: - store = get_vector_store() - results = store.similarity_search_with_relevance_scores(query, k=max_results) + store = ChromaStore() + results = store.search(query, limit=max_results) output = [] - for doc, score in results: + for doc in results: output.append( { "content": doc.page_content, "metadata": doc.metadata, - "score": round(score, 4), } ) return output \ No newline at end of file