56 lines
1.6 KiB
Python
56 lines
1.6 KiB
Python
from pathlib import Path
|
|
from uuid import uuid4
|
|
|
|
from langchain_qdrant import Qdrant
|
|
from langchain_core.documents import Document
|
|
from langchain_ollama import OllamaEmbeddings
|
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
|
|
|
QDRANT_DIR = "./qdrant_db"
|
|
COLLECTION_NAME = "local_kb"
|
|
EMBED_MODEL = "nomic-embed-text"
|
|
OLLAMA_BASE_URL = "http://127.0.0.1:11434"
|
|
|
|
|
|
def create_vectorstore(persist_directory: str = QDRANT_DIR) -> Qdrant:
|
|
embeddings = OllamaEmbeddings(
|
|
model=EMBED_MODEL,
|
|
base_url=OLLAMA_BASE_URL,
|
|
)
|
|
return Qdrant(
|
|
collection_name=COLLECTION_NAME,
|
|
embedding_function=embeddings,
|
|
persist_directory=persist_directory,
|
|
)
|
|
|
|
|
|
def load_documents(directory: str, vectorstore: Qdrant) -> int:
|
|
splitter = RecursiveCharacterTextSplitter(
|
|
chunk_size=1000,
|
|
chunk_overlap=200,
|
|
add_start_index=True,
|
|
)
|
|
|
|
base_path = Path(directory)
|
|
if not base_path.exists():
|
|
raise FileNotFoundError(f"Directory not found: {directory}")
|
|
|
|
raw_docs: list[Document] = []
|
|
for pattern in ("*.txt", "*.md"):
|
|
for path in sorted(base_path.rglob(pattern)):
|
|
content = path.read_text(encoding="utf-8")
|
|
raw_docs.append(
|
|
Document(
|
|
page_content=content,
|
|
metadata={"source": str(path)},
|
|
)
|
|
)
|
|
|
|
if not raw_docs:
|
|
return 0
|
|
|
|
chunks = splitter.split_documents(raw_docs)
|
|
ids = [str(uuid4()) for _ in chunks]
|
|
vectorstore.add_documents(documents=chunks, ids=ids)
|
|
return len(chunks)
|