weiweiw
/
Langchain-Chatchat
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

update format

This commit is contained in:
imClumsyPanda 2023-08-27 20:23:07 +08:00
parent 58d5787bdf
commit 8faafa6a32
1 changed files with 16 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
server/knowledge_base/utils.py
View File

@ -25,27 +25,34 @@ def validate_kb_name(knowledge_base_id: str) -> bool:
return False return False
return True return True
def get_kb_path(knowledge_base_name: str): def get_kb_path(knowledge_base_name: str):
return os.path.join(KB_ROOT_PATH, knowledge_base_name) return os.path.join(KB_ROOT_PATH, knowledge_base_name)
def get_doc_path(knowledge_base_name: str): def get_doc_path(knowledge_base_name: str):
return os.path.join(get_kb_path(knowledge_base_name), "content") return os.path.join(get_kb_path(knowledge_base_name), "content")
def get_vs_path(knowledge_base_name: str): def get_vs_path(knowledge_base_name: str):
return os.path.join(get_kb_path(knowledge_base_name), "vector_store") return os.path.join(get_kb_path(knowledge_base_name), "vector_store")
def get_file_path(knowledge_base_name: str, doc_name: str): def get_file_path(knowledge_base_name: str, doc_name: str):
return os.path.join(get_doc_path(knowledge_base_name), doc_name) return os.path.join(get_doc_path(knowledge_base_name), doc_name)
def list_kbs_from_folder(): def list_kbs_from_folder():
return [f for f in os.listdir(KB_ROOT_PATH) return [f for f in os.listdir(KB_ROOT_PATH)
if os.path.isdir(os.path.join(KB_ROOT_PATH, f))] if os.path.isdir(os.path.join(KB_ROOT_PATH, f))]
def list_docs_from_folder(kb_name: str): def list_docs_from_folder(kb_name: str):
doc_path = get_doc_path(kb_name) doc_path = get_doc_path(kb_name)
return [file for file in os.listdir(doc_path) return [file for file in os.listdir(doc_path)
if os.path.isfile(os.path.join(doc_path, file))] if os.path.isfile(os.path.join(doc_path, file))]
@lru_cache(1) @lru_cache(1)
def load_embeddings(model: str, device: str): def load_embeddings(model: str, device: str):
if model == "text-embedding-ada-002": # openai text-embedding-ada-002 if model == "text-embedding-ada-002": # openai text-embedding-ada-002
@ -61,7 +68,6 @@ def load_embeddings(model: str, device: str):
return embeddings return embeddings
LOADER_DICT = {"UnstructuredHTMLLoader": ['.html'], LOADER_DICT = {"UnstructuredHTMLLoader": ['.html'],
"UnstructuredMarkdownLoader": ['.md'], "UnstructuredMarkdownLoader": ['.md'],
"CustomJSONLoader": [".json"], "CustomJSONLoader": [".json"],
@ -81,12 +87,12 @@ class CustomJSONLoader(langchain.document_loaders.JSONLoader):
''' '''
def __init__( def __init__(
self, self,
file_path: Union[str, Path], file_path: Union[str, Path],
content_key: Optional[str] = None, content_key: Optional[str] = None,
metadata_func: Optional[Callable[[Dict, Dict], Dict]] = None, metadata_func: Optional[Callable[[Dict, Dict], Dict]] = None,
text_content: bool = True, text_content: bool = True,
json_lines: bool = False, json_lines: bool = False,
): ):
"""Initialize the JSONLoader. """Initialize the JSONLoader.
@ -141,6 +147,7 @@ class CustomJSONLoader(langchain.document_loaders.JSONLoader):
text = self._get_text(sample=sample, metadata=metadata) text = self._get_text(sample=sample, metadata=metadata)
docs.append(Document(page_content=text, metadata=metadata)) docs.append(Document(page_content=text, metadata=metadata))
langchain.document_loaders.CustomJSONLoader = CustomJSONLoader langchain.document_loaders.CustomJSONLoader = CustomJSONLoader
@ -186,9 +193,9 @@ class KnowledgeFile:
elif self.document_loader_name == "CustomJSONLoader": elif self.document_loader_name == "CustomJSONLoader":
loader = DocumentLoader(self.filepath, text_content=False) loader = DocumentLoader(self.filepath, text_content=False)
elif self.document_loader_name == "UnstructuredMarkdownLoader": elif self.document_loader_name == "UnstructuredMarkdownLoader":
loader = DocumentLoader(self.filepath, mode="elements") # TODO: 需要在实践中测试`elements`是否优于`single` loader = DocumentLoader(self.filepath, mode="elements") # TODO: 需要在实践中测试`elements`是否优于`single`
elif self.document_loader_name == "UnstructuredHTMLLoader": elif self.document_loader_name == "UnstructuredHTMLLoader":
loader = DocumentLoader(self.filepath, mode="elements") # TODO: 需要在实践中测试`elements`是否优于`single` loader = DocumentLoader(self.filepath, mode="elements") # TODO: 需要在实践中测试`elements`是否优于`single`
else: else:
loader = DocumentLoader(self.filepath) loader = DocumentLoader(self.filepath)