diff --git a/vectorstore.py b/vectorstore.py new file mode 100644 index 0000000..b68ea1b --- /dev/null +++ b/vectorstore.py @@ -0,0 +1,45 @@ +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 + + +def create_vectorstore(persist_directory: str = "./chroma_db"): + """Create a ChromaDB vectorstore with Ollama embeddings.""" + embeddings = OllamaEmbeddings(model="nomic-embed-text") + vectorstore = Chroma( + persist_directory=persist_directory, + embedding_function=embeddings + ) + return vectorstore + + +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