34 lines
1.3 KiB
Python
34 lines
1.3 KiB
Python
"""ChromaDB + OllamaEmbeddings (nomic-embed-text)."""
|
|
from pathlib import Path
|
|
from langchain_ollama import OllamaEmbeddings
|
|
from langchain_chroma import Chroma
|
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
|
from langchain_core.documents import Document
|
|
|
|
|
|
def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
|
|
"""Создать/открыть ChromaDB с Ollama-эмбеддингами."""
|
|
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
|
return Chroma(
|
|
collection_name="rag_kb",
|
|
embedding_function=embeddings,
|
|
persist_directory=persist_directory,
|
|
)
|
|
|
|
|
|
def load_documents(directory: str, vectorstore: Chroma) -> int:
|
|
"""Загрузить .txt/.md файлы из директории в ChromaDB с чанкингом."""
|
|
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
|
|
docs = []
|
|
for path in Path(directory).glob("**/*"):
|
|
if path.suffix.lower() not in (".txt", ".md"):
|
|
continue
|
|
text = path.read_text(encoding="utf-8", errors="replace")
|
|
chunks = splitter.split_text(text)
|
|
docs.extend(
|
|
[Document(page_content=c, metadata={"source": str(path)}) for c in chunks]
|
|
)
|
|
if docs:
|
|
vectorstore.add_documents(docs)
|
|
return len(docs)
|