Решение готово к публикации: update vectorstore.py
This commit is contained in:
+12
-3
@@ -2,7 +2,6 @@ from pathlib import Path
|
||||
from uuid import uuid4
|
||||
|
||||
from langchain_chroma import Chroma, Settings
|
||||
|
||||
from langchain_ollama import OllamaEmbeddings
|
||||
from langchain_core.documents import Document
|
||||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||||
@@ -14,12 +13,17 @@ OLLAMA_BASE_URL = "http://127.0.0.1:11434"
|
||||
|
||||
|
||||
def create_vectorstore(persist_directory: str = CHROMA_DIR):
|
||||
"""Return a Chroma collection configured for the local knowledge base.
|
||||
|
||||
The collection is persisted in ``persist_directory`` and named
|
||||
``COLLECTION_NAME``.
|
||||
"""
|
||||
client = Chroma(Settings(persist_directory=persist_directory))
|
||||
collection = client.get_or_create_collection(name=COLLECTION_NAME)
|
||||
return collection
|
||||
|
||||
|
||||
def load_documents(directory: str, vectorstore: Chroma) -> int:
|
||||
def load_documents(directory: str, collection) -> int:
|
||||
splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=1000,
|
||||
chunk_overlap=200,
|
||||
@@ -48,5 +52,10 @@ def load_documents(directory: str, vectorstore: Chroma) -> int:
|
||||
ids = [str(uuid4()) for _ in chunks]
|
||||
embeddings = OllamaEmbeddings(model=EMBED_MODEL, base_url=OLLAMA_BASE_URL)
|
||||
vectors = embeddings.embed_documents([doc.page_content for doc in chunks])
|
||||
vectorstore.add(ids=ids, documents=[doc.page_content for doc in chunks], embeddings=vectors, metadatas=[doc.metadata for doc in chunks])
|
||||
collection.add(
|
||||
ids=ids,
|
||||
documents=[doc.page_content for doc in chunks],
|
||||
embeddings=vectors,
|
||||
metadatas=[doc.metadata for doc in chunks],
|
||||
)
|
||||
return len(chunks)
|
||||
|
||||
Reference in New Issue
Block a user