125 lines
3.5 KiB
Python
125 lines
3.5 KiB
Python
import os
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
from langchain_ollama import OllamaEmbeddings
|
|
from langchain_qdrant import QdrantVectorStore
|
|
from qdrant_client import QdrantClient
|
|
|
|
|
|
def create_vectorstore(
|
|
persist_directory: str = "./qdrant_db",
|
|
collection_name: str = "documents",
|
|
embedding_model: str = "nomic-embed-text",
|
|
) -> QdrantVectorStore:
|
|
"""
|
|
Create a Qdrant vector store backed by Ollama embeddings.
|
|
|
|
Parameters
|
|
----------
|
|
persist_directory : str
|
|
Path to the directory where Qdrant will store its data.
|
|
collection_name : str
|
|
Name of the collection to use.
|
|
embedding_model : str
|
|
Ollama model name for embeddings.
|
|
|
|
Returns
|
|
-------
|
|
QdrantVectorStore
|
|
Initialized vector store.
|
|
"""
|
|
# Ensure the persistence directory exists
|
|
Path(persist_directory).mkdir(parents=True, exist_ok=True)
|
|
|
|
# Create a local Qdrant client that stores data on disk
|
|
client = QdrantClient(path=persist_directory)
|
|
|
|
# Initialize embeddings
|
|
embeddings = OllamaEmbeddings(model=embedding_model)
|
|
|
|
# Create or load the collection
|
|
vectorstore = QdrantVectorStore(
|
|
client=client,
|
|
embeddings=embeddings,
|
|
collection_name=collection_name,
|
|
)
|
|
return vectorstore
|
|
|
|
|
|
def load_documents(
|
|
directory: str,
|
|
vectorstore: QdrantVectorStore,
|
|
chunk_size: int = 1000,
|
|
chunk_overlap: int = 200,
|
|
) -> None:
|
|
"""
|
|
Load all .txt and .md files from a directory, split them into chunks,
|
|
embed, and add to the vector store.
|
|
|
|
Parameters
|
|
----------
|
|
directory : str
|
|
Path to the directory containing documents.
|
|
vectorstore : QdrantVectorStore
|
|
The vector store to populate.
|
|
chunk_size : int
|
|
Maximum size of each chunk in characters.
|
|
chunk_overlap : int
|
|
Number of overlapping characters between chunks.
|
|
"""
|
|
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
|
from langchain.docstore.document import Document
|
|
|
|
# Prepare text splitter
|
|
splitter = RecursiveCharacterTextSplitter(
|
|
chunk_size=chunk_size,
|
|
chunk_overlap=chunk_overlap,
|
|
)
|
|
|
|
# Collect all documents
|
|
docs = []
|
|
for root, _, files in os.walk(directory):
|
|
for file in files:
|
|
if file.lower().endswith((".txt", ".md")):
|
|
file_path = os.path.join(root, file)
|
|
with open(file_path, "r", encoding="utf-8") as f:
|
|
content = f.read()
|
|
# Split content into chunks
|
|
chunks = splitter.split_text(content)
|
|
# Create Document objects with metadata
|
|
for i, chunk in enumerate(chunks):
|
|
doc = Document(
|
|
page_content=chunk,
|
|
metadata={
|
|
"source": file_path,
|
|
"chunk_index": i,
|
|
},
|
|
)
|
|
docs.append(doc)
|
|
|
|
if docs:
|
|
vectorstore.add_documents(docs)
|
|
else:
|
|
print("No documents found to load.")
|
|
|
|
|
|
def collection_exists(vectorstore: QdrantVectorStore) -> bool:
|
|
"""
|
|
Check if the collection already exists in Qdrant.
|
|
|
|
Parameters
|
|
----------
|
|
vectorstore : QdrantVectorStore
|
|
The vector store to check.
|
|
|
|
Returns
|
|
-------
|
|
bool
|
|
True if collection exists, False otherwise.
|
|
"""
|
|
try:
|
|
vectorstore.client.get_collection(vectorstore.collection_name)
|
|
return True
|
|
except Exception:
|
|
return False |