diff --git a/vectorstore.py b/vectorstore.py new file mode 100644 index 0000000..2deb990 --- /dev/null +++ b/vectorstore.py @@ -0,0 +1,34 @@ +import os +from langchain_community.embeddings 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: + embeddings = OllamaEmbeddings(model="nomic-embed-text") + return Chroma( + collection_name="knowledge", + embedding_function=embeddings, + persist_directory=persist_directory, + ) + +def load_documents(directory: str, vectorstore: Chroma) -> None: + splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) + documents = [] + for root, _, files in os.walk(directory): + for file in files: + if file.endswith((".txt", ".md")): + path = os.path.join(root, file) + with open(path, "r", encoding="utf-8") as f: + content = f.read() + docs = splitter.split_text(content) + for i, chunk in enumerate(docs): + documents.append( + Document( + page_content=chunk, + metadata={"source": path, "chunk": i}, + ) + ) + if documents: + vectorstore.add_documents(documents) + vectorstore.persist() \ No newline at end of file