Публикация решения: update vectorstore.py
This commit is contained in:
+15
-16
@@ -1,35 +1,37 @@
|
||||
"""
|
||||
RAG vector store using Qdrant and Ollama embeddings.
|
||||
RAG vector store using ChromaDB and Ollama embeddings.
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
|
||||
import chromadb # kept for compatibility if needed
|
||||
import chromadb
|
||||
from langchain.embeddings.ollama import OllamaEmbeddings
|
||||
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
||||
from langchain.schema.document import Document
|
||||
from langchain.vectorstores import Qdrant
|
||||
from langchain.vectorstores import Chroma
|
||||
|
||||
CHROMA_DIR = "./chroma_db"
|
||||
EMBED_MODEL = "nomic-embed-text"
|
||||
|
||||
|
||||
def create_vectorstore(persist_directory: str = CHROMA_DIR):
|
||||
"""Create or load a Qdrant vector store.
|
||||
def create_vectorstore(persist_directory: str = CHROMA_DIR) -> Chroma:
|
||||
"""Create or load a Chroma vector store.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
persist_directory:
|
||||
Directory where the Qdrant database is stored. If it does not exist, it will be created.
|
||||
persist_directory : str, optional
|
||||
Directory where the Chroma database is stored. If it does not exist,
|
||||
it will be created automatically by Chroma.
|
||||
"""
|
||||
embeddings = OllamaEmbeddings(model=EMBED_MODEL)
|
||||
# Qdrant can use a local file store via `path` argument
|
||||
client = Qdrant(persist_directory=persist_directory, embedding_function=embeddings)
|
||||
return client
|
||||
# Chroma can use a local directory for persistence
|
||||
client = chromadb.PersistentClient(path=persist_directory)
|
||||
collection = client.get_or_create_collection(name="documents", embedding_function=embeddings)
|
||||
return Chroma(collection=collection, embedding_function=embeddings)
|
||||
|
||||
|
||||
def load_documents(directory: str, vectorstore) -> None:
|
||||
def load_documents(directory: str, vectorstore: Chroma) -> None:
|
||||
"""Load all .txt and .md files from *directory*, chunk them and add to the vector store.
|
||||
|
||||
The function does not return anything; it mutates the provided collection.
|
||||
@@ -44,8 +46,5 @@ def load_documents(directory: str, vectorstore) -> None:
|
||||
docs.extend(text_splitter.create_documents([content], metadata={"source": str(path)}))
|
||||
|
||||
if docs:
|
||||
# Qdrant expects texts and metadatas lists
|
||||
vectorstore.add_texts(
|
||||
texts=[doc.page_content for doc in docs],
|
||||
metadatas=[doc.metadata for doc in docs],
|
||||
)
|
||||
# Chroma expects a list of documents via add_documents
|
||||
vectorstore.add_documents(docs)
|
||||
|
||||
Reference in New Issue
Block a user