22 lines
775 B
Python
22 lines
775 B
Python
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)
|