import os
import re
import stat
import sys
from urllib.parse import unquote, urlparse
import jsonlines
from agent.log import logger
from agent.utils.doc_parser import parse_doc, parse_html_bs
from agent.utils.utils import print_traceback, save_text_to_file
from schema import Record
from b2sdk.v2 import B2Api
from b2sdk.v2 import InMemoryAccountInfo
import hashlib
import datetime
from io import BytesIO
def _fix_secure_write_for_code_interpreter(code_interpreter_ws):
if 'linux' in sys.platform.lower():
fname = os.path.join(code_interpreter_ws, 'test_file_permission.txt')
if os.path.exists(fname):
os.remove(fname)
with os.fdopen(
os.open(fname, os.O_CREAT | os.O_WRONLY | os.O_TRUNC, 0o0600),
'w') as f:
f.write('test')
file_mode = stat.S_IMODE(os.stat(fname).st_mode) & 0o6677
if file_mode != 0o0600:
os.environ['JUPYTER_ALLOW_INSECURE_WRITES'] = '1'
work_space_root = "./workspace"
cache_root = f"{work_space_root}/browser_cache/"
download_root = f"{work_space_root}/download/"
code_interpreter_ws = f"{work_space_root}/ci_workspace/"
cache_file_popup_url = os.path.join(cache_root, 'popup_url.jsonl')
cache_file = os.path.join(cache_root, 'browse.jsonl')
max_ref_token = 4000
max_days = 7
os.makedirs(work_space_root, exist_ok=True)
os.makedirs(cache_root, exist_ok=True)
os.makedirs(download_root, exist_ok=True)
os.makedirs(code_interpreter_ws, exist_ok=True)
code_interpreter_work_dir = code_interpreter_ws
os.environ['M6_CODE_INTERPRETER_WORK_DIR'] = code_interpreter_work_dir
os.environ['M6_CODE_INTERPRETER_STATIC_URL'] = f'{os.getenv("DOMAIN")}/static'
os.environ["HF_HOME"] = ".cache/huggingface/"
os.environ["MPLCONFIGDIR"] = ".cache/huggingface/"
_fix_secure_write_for_code_interpreter(code_interpreter_ws)
class B2Manager():
def __init__(self):
info = InMemoryAccountInfo()
b2_api = B2Api(info)
application_key_id = os.environ.get("b2_key_id")
application_key = os.environ.get("b2_key")
b2_api.authorize_account("production", application_key_id, application_key)
self.b2_bucket = b2_api.get_bucket_by_name(os.environ.get("b2_bucket_name"))
self.b2_api = b2_api
self.file_name = None
def gen_file_name(self, access_token, url, need_md5):
url_md5 = hashlib.md5(b'%s' % url.encode(encoding='UTF-8')).hexdigest()
return f"{access_token}/{url_md5}" if need_md5 else f"{access_token}/{url}"
def get(self, access_token, url, need_md5=True):
in_memory_file = BytesIO()
self.b2_bucket.download_file_by_name(self.gen_file_name(access_token, url, need_md5)).save(in_memory_file)
# export_file = self.b2_bucket.download_file_by_name(self.file_name)
# export_file.save(in_memory_file)
in_memory_file.seek(0)
return str(in_memory_file.read(), "utf-8")
def upsert(self, access_token, url, content, need_md5=True):
self.b2_bucket.upload_bytes(content.encode('utf-8'), self.gen_file_name(access_token, url, need_md5), file_infos=None)
# self.b2_bucket.upload()
def delete(self, access_token, url, need_md5=True):
file_version_info = self.b2_bucket.get_file_info_by_name(self.gen_file_name(access_token, url, need_md5))
self.b2_bucket.hide_file(file_version_info.file_name)
# for version in self.b2_bucket.list_file_versions(self.file_name):
# self.b2_bucket.delete_file_version(version.id_, version.file_name)
def list_files(self, access_token):
files = []
for file_version_info, folder_name in self.b2_bucket.ls(folder_to_list=f"{access_token}/"):
# The upload timestamp is in milliseconds, so we divide by 1000 to convert it to seconds
upload_timestamp = datetime.datetime.fromtimestamp(file_version_info.upload_timestamp / 1000.0)
files.append(f"File Name: {file_version_info.file_name}, \nUpload timestamp: {upload_timestamp}, \nMetadata: {file_version_info.file_info}")
return files
def exists(self, access_token, url=None, need_md5=True):
try:
self.b2_bucket.get_file_info_by_name(self.gen_file_name(access_token, url, need_md5))
return True
except:
return False
def update_pop_url(data, cache_file_popup_url, access_token):
new_line = {'url': data['url'], "access_token": access_token}
lines = []
for line in jsonlines.open(cache_file_popup_url):
if line['access_token'] == access_token and line['url'] != data['url']:
lines.append(line)
lines.append(new_line)
with jsonlines.open(cache_file_popup_url, mode='w') as writer:
for new_line in lines:
writer.write(new_line)
response = 'Update URL'
return response
def change_checkbox_state(text, cache_file, access_token):
if not os.path.exists(cache_file):
return {'result': 'no file'}
lines = []
for line in jsonlines.open(cache_file):
if line['access_token'] == access_token and line['url'] == text[3:]:
if line['checked']:
line['checked'] = False
else:
line['checked'] = True
lines.append(line)
with jsonlines.open(cache_file, mode='w') as writer:
for new_line in lines:
writer.write(new_line)
return {'result': 'changed'}
def is_local_path(path):
if path.startswith('https://') or path.startswith('http://'):
return False
return True
def sanitize_chrome_file_path(file_path: str) -> str:
# For Linux and macOS.
if os.path.exists(file_path):
return file_path
# For native Windows, drop the leading '/' in '/C:/'
win_path = file_path
if win_path.startswith('/'):
win_path = win_path[1:]
if os.path.exists(win_path):
return win_path
# For Windows + WSL.
if re.match(r'^[A-Za-z]:/', win_path):
wsl_path = f'/mnt/{win_path[0].lower()}/{win_path[3:]}'
if os.path.exists(wsl_path):
return wsl_path
# For native Windows, replace / with \.
win_path = win_path.replace('/', '\\')
if os.path.exists(win_path):
return win_path
return file_path
def extract_and_cache_document(data, cache_root, access_token):
logger.info('Starting cache pages...')
if data['url'].split('.')[-1].lower() in ['pdf', 'docx', 'pptx']:
date1 = datetime.datetime.now()
# generate one processing record
new_record = Record(url=data['url'],
time='',
type=data['type'],
raw=[],
extract='',
access_token=access_token,
topic='',
checked=False,
session=[])
service.upsert(access_token, data['url'], new_record.model_dump_json())
if data['url'].startswith('https://') or data['url'].startswith( 'http://'):
pdf_path = data['url']
else:
parsed_url = urlparse(data['url'])
pdf_path = unquote(parsed_url.path)
pdf_path = sanitize_chrome_file_path(pdf_path)
try:
pdf_content = parse_doc(pdf_path)
except Exception:
print_traceback()
# del the processing record
service.delete(access_token, data['url'])
return 'failed'
date2 = datetime.datetime.now()
logger.info('Parsing pdf time: ' + str(date2 - date1))
data['content'] = pdf_content
data['type'] = 'pdf'
extract = pdf_path.split('/')[-1].split('\\')[-1].split('.')[0]
elif data['content'] and data['type'] == 'html':
new_record = Record(url=data['url'],
time='',
type=data['type'],
raw=[],
extract='',
access_token=access_token,
topic='',
checked=False,
session=[])
service.upsert(access_token, data['url'], new_record.model_dump_json())
try:
tmp_html_file = os.path.join(cache_root, 'tmp.html')
save_text_to_file(tmp_html_file, data['content'])
data['content'] = parse_html_bs(tmp_html_file)
except Exception:
print_traceback()
extract = data['content'][0]['metadata']['title']
else:
logger.error(
'Only Support the Following File Types: [\'.html\', \'.pdf\', \'.docx\', \'.pptx\']'
)
raise NotImplementedError
today = datetime.date.today()
new_record = Record(url=data['url'],
time=str(today),
type=data['type'],
raw=data['content'],
extract=extract,
access_token=access_token,
topic='',
checked=True,
session=[])
service.upsert(access_token, data['url'], new_record.model_dump_json())
response = 'Cached'
return response
service = B2Manager()
if __name__ == '__main__':
# print(service.gen_file_name("test", "settings.xml"))
# print(service.get("test", "settings.xml"))
# print(service.upsert("test", "settings.xml", b"1111"))
print(service.list_files("test"))
print(service.exists("test", "https://tree-iad1-0003.secure.backblaze.com/b2_browse_files2.htm1"))