generalize position helper for parsing _position.yaml and sorting objects by name (#2803)

This commit is contained in:
Bowen Liang
2024-03-13 20:29:38 +08:00
committed by GitHub
parent 849dc0560b
commit 8b15b742ad
5 changed files with 95 additions and 46 deletions

View File

@@ -18,6 +18,7 @@ from core.model_runtime.entities.model_entities import (
)
from core.model_runtime.errors.invoke import InvokeAuthorizationError, InvokeError
from core.model_runtime.model_providers.__base.tokenizers.gpt2_tokenzier import GPT2Tokenizer
from core.utils.position_helper import get_position_map, sort_by_position_map
class AIModel(ABC):
@@ -148,15 +149,7 @@ class AIModel(ABC):
]
# get _position.yaml file path
position_file_path = os.path.join(provider_model_type_path, '_position.yaml')
# read _position.yaml file
position_map = {}
if os.path.exists(position_file_path):
with open(position_file_path, encoding='utf-8') as f:
positions = yaml.safe_load(f)
# convert list to dict with key as model provider name, value as index
position_map = {position: index for index, position in enumerate(positions)}
position_map = get_position_map(provider_model_type_path)
# traverse all model_schema_yaml_paths
for model_schema_yaml_path in model_schema_yaml_paths:
@@ -206,8 +199,7 @@ class AIModel(ABC):
model_schemas.append(model_schema)
# resort model schemas by position
if position_map:
model_schemas.sort(key=lambda x: position_map.get(x.model, 999))
model_schemas = sort_by_position_map(position_map, model_schemas, lambda x: x.model)
# cache model schemas
self.model_schemas = model_schemas

View File

@@ -1,10 +1,8 @@
import importlib
import logging
import os
from collections import OrderedDict
from typing import Optional
import yaml
from pydantic import BaseModel
from core.model_runtime.entities.model_entities import ModelType
@@ -12,6 +10,7 @@ from core.model_runtime.entities.provider_entities import ProviderConfig, Provid
from core.model_runtime.model_providers.__base.model_provider import ModelProvider
from core.model_runtime.schema_validators.model_credential_schema_validator import ModelCredentialSchemaValidator
from core.model_runtime.schema_validators.provider_credential_schema_validator import ProviderCredentialSchemaValidator
from core.utils.position_helper import get_position_map, sort_to_dict_by_position_map
logger = logging.getLogger(__name__)
@@ -200,7 +199,6 @@ class ModelProviderFactory:
if self.model_provider_extensions:
return self.model_provider_extensions
model_providers = {}
# get the path of current classes
current_path = os.path.abspath(__file__)
@@ -215,17 +213,10 @@ class ModelProviderFactory:
]
# get _position.yaml file path
position_file_path = os.path.join(model_providers_path, '_position.yaml')
# read _position.yaml file
position_map = {}
if os.path.exists(position_file_path):
with open(position_file_path, encoding='utf-8') as f:
positions = yaml.safe_load(f)
# convert list to dict with key as model provider name, value as index
position_map = {position: index for index, position in enumerate(positions)}
position_map = get_position_map(model_providers_path)
# traverse all model_provider_dir_paths
model_providers: list[ModelProviderExtension] = []
for model_provider_dir_path in model_provider_dir_paths:
# get model_provider dir name
model_provider_name = os.path.basename(model_provider_dir_path)
@@ -256,14 +247,13 @@ class ModelProviderFactory:
logger.warning(f"Missing {model_provider_name}.yaml file in {model_provider_dir_path}, Skip.")
continue
model_providers[model_provider_name] = ModelProviderExtension(
model_providers.append(ModelProviderExtension(
name=model_provider_name,
provider_instance=model_provider_class(),
position=position_map.get(model_provider_name)
)
))
sorted_items = sorted(model_providers.items(), key=lambda x: (x[1].position is None, x[1].position))
sorted_extensions = OrderedDict(sorted_items)
sorted_extensions = sort_to_dict_by_position_map(position_map, model_providers, lambda x: x.name)
self.model_provider_extensions = sorted_extensions