wanyaokun
121 lines
3.8 KiB
Python
121 lines
3.8 KiB
Python
import os
|
|
import logging
|
|
from typing import Dict
|
|
|
|
from llama_index.core.readers.base import BaseReader
|
|
from llama_index.core.readers.json import JSONReader
|
|
from llama_parse import LlamaParse
|
|
from pydantic import BaseModel, validator
|
|
from app.engine.loaders.markdownReader import ChunkMarkdownReader
|
|
from app.engine.loaders.projectJson import ProjectJson
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class FileLoaderConfig(BaseModel):
|
|
data_dir: str = "data"
|
|
use_llama_parse: bool = False
|
|
|
|
@validator("data_dir")
|
|
def data_dir_must_exist(cls, v):
|
|
if not os.path.isdir(v):
|
|
raise ValueError(f"Directory '{v}' does not exist")
|
|
return v
|
|
|
|
def llama_parse_parser():
|
|
if os.getenv("LLAMA_CLOUD_API_KEY") is None:
|
|
raise ValueError(
|
|
"LLAMA_CLOUD_API_KEY environment variable is not set. "
|
|
"Please set it in .env file or in your shell environment then run again!"
|
|
)
|
|
parser = LlamaParse(
|
|
result_type="markdown",
|
|
verbose=True,
|
|
language="en",
|
|
ignore_errors=False,
|
|
)
|
|
return parser
|
|
|
|
def llama_parse_extractor() -> Dict[str, LlamaParse]:
|
|
from llama_parse.utils import SUPPORTED_FILE_TYPES
|
|
|
|
parser = llama_parse_parser()
|
|
return {file_type: parser for file_type in SUPPORTED_FILE_TYPES}
|
|
|
|
def llama_local_extractor() -> Dict[str, BaseReader]:
|
|
parser = {
|
|
".json" : JSONReader(clean_json=False,levels_back=0),
|
|
".md" : ChunkMarkdownReader(),
|
|
}
|
|
return parser
|
|
|
|
def get_file_documents(config: FileLoaderConfig,childPath: str):
|
|
from llama_index.core.readers import SimpleDirectoryReader
|
|
|
|
try:
|
|
file_extractor = None
|
|
if config.use_llama_parse:
|
|
# LlamaParse is async first,
|
|
# so we need to use nest_asyncio to run it in sync mode
|
|
import nest_asyncio
|
|
|
|
nest_asyncio.apply()
|
|
|
|
file_extractor = llama_parse_extractor()
|
|
else:
|
|
file_extractor = llama_local_extractor()
|
|
|
|
reader = SimpleDirectoryReader(
|
|
os.path.join(config.data_dir,childPath.replace('_','\\')),
|
|
recursive=True,
|
|
filename_as_id=True,
|
|
raise_on_error=True,
|
|
file_extractor=file_extractor,
|
|
)
|
|
return reader.load_data()
|
|
except Exception as e:
|
|
import sys
|
|
import traceback
|
|
|
|
# Catch the error if the data dir is empty
|
|
# and return as empty document list
|
|
_, _, exc_traceback = sys.exc_info()
|
|
function_name = traceback.extract_tb(exc_traceback)[-1].name
|
|
if function_name == "_add_files":
|
|
logger.warning(
|
|
f"Failed to load file documents, error message: {e} . Return as empty document list."
|
|
)
|
|
return []
|
|
else:
|
|
# Raise the error if it is not the case of empty data dir
|
|
raise e
|
|
|
|
def prjFileSuffix(dir:str):
|
|
entries = os.listdir(dir)
|
|
file_names = [entry for entry in entries if os.path.isfile(os.path.join(dir, entry))]
|
|
if len(file_names) > 0:
|
|
return os.path.splitext(file_names[0])[1]
|
|
return ''
|
|
|
|
def getProjectName(dir:str):
|
|
suffix = prjFileSuffix(dir)
|
|
if suffix== '.json':
|
|
prjJson = ProjectJson(dir)
|
|
prjJson.parse()
|
|
tb = prjJson.table('工程属性')
|
|
records = tb.records()
|
|
for record in records:
|
|
name = record.value('名称')
|
|
if name == '工程名称':
|
|
return record.value('值')
|
|
elif suffix == '.md':
|
|
md_files = [f for f in os.listdir(dir) if f.endswith('.md')]
|
|
for md_file in md_files:
|
|
prjPath = os.path.join(dir, md_file)
|
|
basename = os.path.splitext(md_file)[0]
|
|
if basename =='工程属性':
|
|
rd = ChunkMarkdownReader()
|
|
rd.load_data(prjPath)
|
|
return rd.findValue("名称=='工程名称'",'值')
|
|
return '' |