diff --git a/vectorstore.py b/vectorstore.py index b68ea1b..46018df 100644 --- a/vectorstore.py +++ b/vectorstore.py @@ -1,45 +1,26 @@ import os -from langchain_chroma import Chroma -from langchain_ollama import OllamaEmbeddings -from langchain_text_splitters import RecursiveCharacterTextSplitter -from langchain_community.document_loaders import TextLoader, DirectoryLoader +from pathlib import Path +from typing import List +from langchain_ollama import OllamaEmbeddings +from langchain_chroma import Chroma +from langchain_text_splitters import RecursiveCharacterTextSplitter +from langchain.docstore.document import Document def create_vectorstore(persist_directory: str = "./chroma_db"): - """Create a ChromaDB vectorstore with Ollama embeddings.""" + if not os.path.exists(persist_directory): + os.makedirs(persist_directory, exist_ok=True) embeddings = OllamaEmbeddings(model="nomic-embed-text") - vectorstore = Chroma( - persist_directory=persist_directory, - embedding_function=embeddings - ) - return vectorstore + return Chroma(persist_directory=persist_directory, embedding_function=embeddings) - -def load_documents(directory: str, vectorstore: Chroma): - """Load documents from directory, chunk them, and add to vectorstore.""" - # Load all .txt and .md files from directory - loader = DirectoryLoader( - directory, - glob="**/*.{txt,md}", - loader_cls=TextLoader, - loader_kwargs={"encoding": "utf-8"} - ) - documents = loader.load() - - # Chunk documents - text_splitter = RecursiveCharacterTextSplitter( - chunk_size=1000, - chunk_overlap=200, - length_function=len, - is_separator_regex=False - ) - chunks = text_splitter.split_documents(documents) - - # Add to vectorstore - if chunks: - vectorstore.add_documents(chunks) - print(f"Loaded {len(chunks)} chunks from {len(documents)} documents") - else: - print("No documents found to load") - - return vectorstore \ No newline at end of file +def load_documents(directory: str, vectorstore) -> None: + splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) + docs: List[Document] = [] + for path in Path(directory).rglob("*.txt"): + text = path.read_text(encoding="utf-8") + docs.extend(splitter.split_documents([Document(page_content=text)])) + for path in Path(directory).rglob("*.md"): + text = path.read_text(encoding="utf-8") + docs.extend(splitter.split_documents([Document(page_content=text)])) + if docs: + vectorstore.add_documents(docs)