62 lines
1.9 KiB
Python
62 lines
1.9 KiB
Python
from pathlib import Path
|
|
from uuid import uuid4
|
|
|
|
from langchain_chroma import Chroma, Settings
|
|
from langchain_ollama import OllamaEmbeddings
|
|
from langchain_core.documents import Document
|
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
|
|
|
CHROMA_DIR = "./chroma_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 = CHROMA_DIR):
|
|
"""Return a Chroma collection configured for the local knowledge base.
|
|
|
|
The collection is persisted in ``persist_directory`` and named
|
|
``COLLECTION_NAME``.
|
|
"""
|
|
client = Chroma(Settings(persist_directory=persist_directory))
|
|
collection = client.get_or_create_collection(name=COLLECTION_NAME)
|
|
return collection
|
|
|
|
|
|
def load_documents(directory: str, collection) -> 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]
|
|
embeddings = OllamaEmbeddings(model=EMBED_MODEL, base_url=OLLAMA_BASE_URL)
|
|
vectors = embeddings.embed_documents([doc.page_content for doc in chunks])
|
|
collection.add(
|
|
ids=ids,
|
|
documents=[doc.page_content for doc in chunks],
|
|
embeddings=vectors,
|
|
metadatas=[doc.metadata for doc in chunks],
|
|
)
|
|
return len(chunks)
|