File size: 1,452 Bytes
6f179e7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
import logging
from langchain_community.vectorstores.chroma import Chroma
from src.logging import logging_info
from .BaseDB import BaseDB
# TODO 数据库持久化 和 用户进入的加载。
class ChromaDB(BaseDB):
def __init__(self, embedding_name: str = None, persist_dir=None) -> None:
super().__init__(embedding_name, persist_dir)
# logging_info(self.embedding)
def init_db(self):
self.client = Chroma(
persist_directory=self.persist_dir, embedding_function=self.embedding
)
def addStories(self, stories: str, metas: dict = None):
logging_info(self.text_splitter(stories)[-1])
split_stories = self.text_splitter(stories)
self.client.add_texts(
texts=split_stories, metadatas=[metas] * len(split_stories)
)
def searchBySim(
self, query, n_results=5, metas: dict = None, only_return_document=True
):
result = self.client.similarity_search(query, k=n_results, filter=metas)
# print(result)
if only_return_document:
return [i.page_content for i in result]
return result
def deleteStoriesByMeta(self, metas):
ids = self.searchByMeta(metas=metas)["ids"]
if ids:
self.client.delete(ids)
def searchByMeta(self, metas=None, include: list[str] = None) -> dict[str, any]:
return self.client.get(where=metas, include=include)
|