diff --git a/vector_store.py b/vector_store.py new file mode 100644 index 0000000..65a2652 --- /dev/null +++ b/vector_store.py @@ -0,0 +1,21 @@ +from langchain_qdrant import QdrantVectorStore +from langchain_ollama import OllamaEmbeddings +from langchain.schema import Document + +class QdrantStore: + def __init__(self, host="localhost", port=6333, collection_name="rag_collection"): + self.client = QdrantVectorStore( + url=f"http://{host}:{port}", + collection_name=collection_name, + embeddings=OllamaEmbeddings(model="nomic-embed-text") + ) + # ensure collection exists + if not self.client.collection_exists: + self.client.create_collection() + + def add_documents(self, docs): + # docs: list of Document + self.client.add_documents(docs) + + def search(self, query, limit=5): + return self.client.similarity_search(query, k=limit)