Update qdrant_store.py
This commit is contained in:
+24
-37
@@ -1,50 +1,37 @@
|
|||||||
"""Module for interacting with Qdrant vector store using Ollama embeddings.
|
"""Wrapper around QdrantVectorStore using Ollama embeddings.
|
||||||
|
|
||||||
|
This module is intentionally lightweight and can be imported from any
|
||||||
|
context without requiring the caller to set up a package structure.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from typing import List, Dict, Any
|
import os
|
||||||
|
import sys
|
||||||
|
|
||||||
|
# Ensure the current directory is in sys.path so that relative imports work
|
||||||
|
sys.path.append(os.path.dirname(__file__))
|
||||||
|
|
||||||
from langchain_ollama import OllamaEmbeddings
|
from langchain_ollama import OllamaEmbeddings
|
||||||
from langchain_qdrant import QdrantVectorStore
|
from langchain_qdrant import QdrantVectorStore
|
||||||
from qdrant_client import QdrantClient
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||||||
|
from typing import List, Dict
|
||||||
|
|
||||||
class QdrantStore:
|
class QdrantStore:
|
||||||
"""Wrapper around QdrantVectorStore.
|
"""Wrapper around QdrantVectorStore using Ollama embeddings."""
|
||||||
|
|
||||||
Parameters
|
def __init__(self, host: str = "localhost", port: int = 6333, collection_name: str = "knowledge_base"):
|
||||||
----------
|
|
||||||
collection_name: str
|
|
||||||
Name of the Qdrant collection.
|
|
||||||
host: str
|
|
||||||
Qdrant host URL.
|
|
||||||
port: int
|
|
||||||
Qdrant port.
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(self, collection_name: str = "rag_collection", host: str = "localhost", port: int = 6333):
|
|
||||||
self.collection_name = collection_name
|
|
||||||
self.client = QdrantClient(host=host, port=port)
|
|
||||||
self.embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
self.embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
||||||
# Create collection if not exists
|
self.store = QdrantVectorStore(
|
||||||
if collection_name not in self.client.get_collections().collections:
|
url=f"http://{host}:{port}",
|
||||||
self.client.recreate_collection(collection_name=collection_name, vectors_config={"size": 512, "distance": "Cosine"})
|
|
||||||
self.store = QdrantVectorStore.from_existing_collection(
|
|
||||||
collection_name=collection_name,
|
collection_name=collection_name,
|
||||||
embedding=self.embeddings,
|
embeddings=self.embeddings,
|
||||||
client=self.client,
|
|
||||||
)
|
)
|
||||||
|
|
||||||
def add_documents(self, documents: List[str], titles: List[str] | None = None, metadatas: List[Dict[str, Any]] | None = None) -> None:
|
def add_document(self, content: str, title: str) -> None:
|
||||||
"""Add documents to the collection.
|
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
|
||||||
|
chunks = splitter.split_text(content)
|
||||||
|
metadatas = [{"title": title, "source": title} for _ in chunks]
|
||||||
|
self.store.add_texts(chunks, metadatas=metadatas)
|
||||||
|
|
||||||
Each document is added as a separate vector. If titles or metadatas are provided, they are attached.
|
def search(self, query: str, k: int = 5) -> List[Dict]:
|
||||||
"""
|
results = self.store.similarity_search(query, k=k)
|
||||||
if titles is None:
|
return [{"content": doc.page_content, "metadata": doc.metadata} for doc in results]
|
||||||
titles = [f"doc_{i}" for i in range(len(documents))]
|
|
||||||
if metadatas is None:
|
|
||||||
metadatas = [{} for _ in range(len(documents))]
|
|
||||||
self.store.add_texts(texts=documents, metadatas=metadatas, ids=titles)
|
|
||||||
|
|
||||||
def search(self, query: str, k: int = 5) -> List[Dict[str, Any]]:
|
|
||||||
"""Semantic search returning list of dicts with 'content' and 'metadata'."""
|
|
||||||
results = self.store.similarity_search_with_score(query, k=k)
|
|
||||||
return [{"content": r[0].page_content, "metadata": r[0].metadata, "score": r[1]} for r in results]
|
|
||||||
Reference in New Issue
Block a user