make examples/search_kb.py work

This commit is contained in:
seehi 2023-11-22 14:25:07 +08:00
parent b17846401e
commit 8f24808c8a
4 changed files with 44 additions and 24 deletions

View file

@ -140,4 +140,3 @@ class SearchAndSummarize(Action):
logger.debug(prompt)
logger.debug(result)
return result

View file

@ -5,6 +5,7 @@
@Author : alexanderwu
@File : faiss_store.py
"""
import asyncio
import pickle
from pathlib import Path
from typing import Optional
@ -20,7 +21,7 @@ from metagpt.logs import logger
class FaissStore(LocalStore):
def __init__(self, raw_data: Path, cache_dir=None, meta_col='source', content_col='output'):
def __init__(self, raw_data: Path, cache_dir=None, meta_col="source", content_col="output"):
self.meta_col = meta_col
self.content_col = content_col
super().__init__(raw_data, cache_dir)
@ -50,7 +51,7 @@ class FaissStore(LocalStore):
pickle.dump(store, f)
store.index = index
def search(self, query, expand_cols=False, sep='\n', *args, k=5, **kwargs):
def search(self, query, expand_cols=False, sep="\n", *args, k=5, **kwargs):
rsp = self.store.similarity_search(query, k=k, **kwargs)
logger.debug(rsp)
if expand_cols:
@ -58,6 +59,9 @@ class FaissStore(LocalStore):
else:
return str(sep.join([f"{x.page_content}" for x in rsp]))
async def asearch(self, *args, **kwargs):
return await asyncio.to_thread(self.search, *args, **kwargs)
def write(self):
"""Initialize the index and library based on the Document (JSON / XLSX, etc.) file provided by the user."""
if not self.raw_data.exists():
@ -78,8 +82,8 @@ class FaissStore(LocalStore):
raise NotImplementedError
if __name__ == '__main__':
faiss_store = FaissStore(DATA_PATH / 'qcs/qcs_4w.json')
logger.info(faiss_store.search('Oily Skin Facial Cleanser'))
faiss_store.add([f'Oily Skin Facial Cleanser-{i}' for i in range(3)])
logger.info(faiss_store.search('Oily Skin Facial Cleanser'))
if __name__ == "__main__":
faiss_store = FaissStore(DATA_PATH / "qcs/qcs_4w.json")
logger.info(faiss_store.search("Oily Skin Facial Cleanser"))
faiss_store.add([f"Oily Skin Facial Cleanser-{i}" for i in range(3)])
logger.info(faiss_store.search("Oily Skin Facial Cleanser"))

View file

@ -12,24 +12,23 @@ from metagpt.tools import SearchEngineType
class Sales(Role):
def __init__(
self,
name="Xiaomei",
profile="Retail sales guide",
desc="I am a sales guide in retail. My name is Xiaomei. I will answer some customer questions next, and I "
"will answer questions only based on the information in the knowledge base."
"If I feel that you can't get the answer from the reference material, then I will directly reply that"
" I don't know, and I won't tell you that this is from the knowledge base,"
"but pretend to be what I know. Note that each of my replies will be replied in the tone of a "
"professional guide",
store=None
self,
name="Xiaomei",
profile="Retail sales guide",
desc="I am a sales guide in retail. My name is Xiaomei. I will answer some customer questions next, and I "
"will answer questions only based on the information in the knowledge base."
"If I feel that you can't get the answer from the reference material, then I will directly reply that"
" I don't know, and I won't tell you that this is from the knowledge base,"
"but pretend to be what I know. Note that each of my replies will be replied in the tone of a "
"professional guide",
store=None,
):
super().__init__(name, profile, desc=desc)
self._set_store(store)
def _set_store(self, store):
if store:
action = SearchAndSummarize("", engine=SearchEngineType.CUSTOM_ENGINE, search_func=store.search)
action = SearchAndSummarize("", engine=SearchEngineType.CUSTOM_ENGINE, search_func=store.asearch)
else:
action = SearchAndSummarize()
self._init_actions([action])