소스 검색

Refa: add provider info while adding model. (#9273)

### What problem does this PR solve?
#9248

### Type of change

- [x] Refactoring
tags/v0.20.1
Kevin Hu 2 달 전
부모
커밋
9ca86d801e
No account linked to committer's email address
3개의 변경된 파일33개의 추가작업 그리고 33개의 파일을 삭제
  1. 7
    9
      api/apps/llm_app.py
  2. 1
    1
      rag/svr/task_executor.py
  3. 25
    23
      rag/utils/s3_conn.py

+ 7
- 9
api/apps/llm_app.py 파일 보기

@@ -81,7 +81,7 @@ def set_api_key():
raise Exception(m)
chat_passed = True
except Exception as e:
msg += f"\nFail to access model({llm.llm_name}) using this api key." + str(
msg += f"\nFail to access model({llm.fid}/{llm.llm_name}) using this api key." + str(
e)
elif not rerank_passed and llm.model_type == LLMType.RERANK:
assert factory in RerankModel, f"Re-rank model from {factory} is not supported yet."
@@ -94,7 +94,7 @@ def set_api_key():
rerank_passed = True
logging.debug(f'passed model rerank {llm.llm_name}')
except Exception as e:
msg += f"\nFail to access model({llm.llm_name}) using this api key." + str(
msg += f"\nFail to access model({llm.fid}/{llm.llm_name}) using this api key." + str(
e)
if any([embd_passed, chat_passed, rerank_passed]):
msg = ''
@@ -229,7 +229,7 @@ def add_llm():
if not tc and m.find("**ERROR**:") >= 0:
raise Exception(m)
except Exception as e:
msg += f"\nFail to access model({mdl_nm})." + str(
msg += f"\nFail to access model({factory}/{mdl_nm})." + str(
e)
elif llm["model_type"] == LLMType.RERANK:
assert factory in RerankModel, f"RE-rank model from {factory} is not supported yet."
@@ -243,9 +243,9 @@ def add_llm():
if len(arr) == 0:
raise Exception("Not known.")
except KeyError:
msg += f"{factory} dose not support this model({mdl_nm})"
msg += f"{factory} dose not support this model({factory}/{mdl_nm})"
except Exception as e:
msg += f"\nFail to access model({mdl_nm})." + str(
msg += f"\nFail to access model({factory}/{mdl_nm})." + str(
e)
elif llm["model_type"] == LLMType.IMAGE2TEXT.value:
assert factory in CvModel, f"Image to text model from {factory} is not supported yet."
@@ -260,7 +260,7 @@ def add_llm():
if not m and not tc:
raise Exception(m)
except Exception as e:
msg += f"\nFail to access model({mdl_nm})." + str(e)
msg += f"\nFail to access model({factory}/{mdl_nm})." + str(e)
elif llm["model_type"] == LLMType.TTS:
assert factory in TTSModel, f"TTS model from {factory} is not supported yet."
mdl = TTSModel[factory](
@@ -270,7 +270,7 @@ def add_llm():
for resp in mdl.tts("Hello~ Ragflower!"):
pass
except RuntimeError as e:
msg += f"\nFail to access model({mdl_nm})." + str(e)
msg += f"\nFail to access model({factory}/{mdl_nm})." + str(e)
else:
# TODO: check other type of models
pass
@@ -358,8 +358,6 @@ def my_llms():
return server_error_response(e)




@manager.route('/list', methods=['GET']) # noqa: F821
@login_required
def list_app():

+ 1
- 1
rag/svr/task_executor.py 파일 보기

@@ -231,7 +231,7 @@ async def get_storage_binary(bucket, name):
return await trio.to_thread.run_sync(lambda: STORAGE_IMPL.get(bucket, name))


@timeout(60*40, 1)
@timeout(60*80, 1)
async def build_chunks(task, progress_callback):
if task["size"] > DOC_MAXIMUM_SIZE:
set_progress(task["id"], prog=-1, msg="File size exceeds( <= %dMb )" %

+ 25
- 23
rag/utils/s3_conn.py 파일 보기

@@ -30,6 +30,7 @@ class RAGFlowS3:
self.s3_config = settings.S3
self.access_key = self.s3_config.get('access_key', None)
self.secret_key = self.s3_config.get('secret_key', None)
self.session_token = self.s3_config.get('session_token', None)
self.region = self.s3_config.get('region', None)
self.endpoint_url = self.s3_config.get('endpoint_url', None)
self.signature_version = self.s3_config.get('signature_version', None)
@@ -73,31 +74,32 @@ class RAGFlowS3:
s3_params = {
'aws_access_key_id': self.access_key,
'aws_secret_access_key': self.secret_key,
'aws_session_token': self.session_token,
}
if self.region in self.s3_config:
if self.region:
s3_params['region_name'] = self.region
if 'endpoint_url' in self.s3_config:
if self.endpoint_url:
s3_params['endpoint_url'] = self.endpoint_url
if 'signature_version' in self.s3_config:
config_kwargs['signature_version'] = self.signature_version
if 'addressing_style' in self.s3_config:
config_kwargs['addressing_style'] = self.addressing_style
if self.signature_version:
s3_params['signature_version'] = self.signature_version
if self.addressing_style:
s3_params['addressing_style'] = self.addressing_style
if config_kwargs:
s3_params['config'] = Config(**config_kwargs)
self.conn = boto3.client('s3', **s3_params)
self.conn = [boto3.client('s3', **s3_params)]
except Exception:
logging.exception(f"Fail to connect at region {self.region} or endpoint {self.endpoint_url}")

def __close__(self):
del self.conn
del self.conn[0]
self.conn = None

@use_default_bucket
def bucket_exists(self, bucket):
def bucket_exists(self, bucket, *args, **kwargs):
try:
logging.debug(f"head_bucket bucketname {bucket}")
self.conn.head_bucket(Bucket=bucket)
self.conn[0].head_bucket(Bucket=bucket)
exists = True
except ClientError:
logging.exception(f"head_bucket error {bucket}")
@@ -109,10 +111,10 @@ class RAGFlowS3:
fnm = "txtxtxtxt1"
fnm, binary = f"{self.prefix_path}/{fnm}" if self.prefix_path else fnm, b"_t@@@1"
if not self.bucket_exists(bucket):
self.conn.create_bucket(Bucket=bucket)
self.conn[0].create_bucket(Bucket=bucket)
logging.debug(f"create bucket {bucket} ********")

r = self.conn.upload_fileobj(BytesIO(binary), bucket, fnm)
r = self.conn[0].upload_fileobj(BytesIO(binary), bucket, fnm)
return r

def get_properties(self, bucket, key):
@@ -123,14 +125,14 @@ class RAGFlowS3:

@use_prefix_path
@use_default_bucket
def put(self, bucket, fnm, binary, **kwargs):
def put(self, bucket, fnm, binary, *args, **kwargs):
logging.debug(f"bucket name {bucket}; filename :{fnm}:")
for _ in range(1):
try:
if not self.bucket_exists(bucket):
self.conn.create_bucket(Bucket=bucket)
self.conn[0].create_bucket(Bucket=bucket)
logging.info(f"create bucket {bucket} ********")
r = self.conn.upload_fileobj(BytesIO(binary), bucket, fnm)
r = self.conn[0].upload_fileobj(BytesIO(binary), bucket, fnm)

return r
except Exception:
@@ -140,18 +142,18 @@ class RAGFlowS3:

@use_prefix_path
@use_default_bucket
def rm(self, bucket, fnm, **kwargs):
def rm(self, bucket, fnm, *args, **kwargs):
try:
self.conn.delete_object(Bucket=bucket, Key=fnm)
self.conn[0].delete_object(Bucket=bucket, Key=fnm)
except Exception:
logging.exception(f"Fail rm {bucket}/{fnm}")

@use_prefix_path
@use_default_bucket
def get(self, bucket, fnm, **kwargs):
def get(self, bucket, fnm, *args, **kwargs):
for _ in range(1):
try:
r = self.conn.get_object(Bucket=bucket, Key=fnm)
r = self.conn[0].get_object(Bucket=bucket, Key=fnm)
object_data = r['Body'].read()
return object_data
except Exception:
@@ -162,9 +164,9 @@ class RAGFlowS3:

@use_prefix_path
@use_default_bucket
def obj_exist(self, bucket, fnm, **kwargs):
def obj_exist(self, bucket, fnm, *args, **kwargs):
try:
if self.conn.head_object(Bucket=bucket, Key=fnm):
if self.conn[0].head_object(Bucket=bucket, Key=fnm):
return True
except ClientError as e:
if e.response['Error']['Code'] == '404':
@@ -174,10 +176,10 @@ class RAGFlowS3:

@use_prefix_path
@use_default_bucket
def get_presigned_url(self, bucket, fnm, expires, **kwargs):
def get_presigned_url(self, bucket, fnm, expires, *args, **kwargs):
for _ in range(10):
try:
r = self.conn.generate_presigned_url('get_object',
r = self.conn[0].generate_presigned_url('get_object',
Params={'Bucket': bucket,
'Key': fnm},
ExpiresIn=expires)

Loading…
취소
저장