Skip to content

Commit

Permalink
feat(improve-api-endpoints): Added available model API
Browse files Browse the repository at this point in the history
This API endpoint shares the same API token with Dataset API.
It is essential for patching the knowledge base in terms of embedding model.

Related to: #11727
  • Loading branch information
jasonfish568 committed Dec 31, 2024
1 parent a47b514 commit d7854a5
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 0 deletions.
1 change: 1 addition & 0 deletions api/controllers/service_api/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,3 +8,4 @@
from . import index
from .app import annotation, app, audio, completion, conversation, file, message, workflow
from .dataset import dataset, document, hit_testing, segment
from .workspace import models
22 changes: 22 additions & 0 deletions api/controllers/service_api/workspace/models.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,22 @@
import logging

from flask_login import current_user # type: ignore
from flask_restful import Resource, reqparse # type: ignore
from werkzeug.exceptions import Forbidden

from controllers.service_api import api
from controllers.service_api.wraps import validate_dataset_token
from core.model_runtime.utils.encoders import jsonable_encoder
from services.model_provider_service import ModelProviderService

class ModelProviderAvailableModelApi(Resource):
@validate_dataset_token
def get(self, _, model_type):
tenant_id = current_user.current_tenant_id

model_provider_service = ModelProviderService()
models = model_provider_service.get_models_by_model_type(tenant_id=tenant_id, model_type=model_type)

return jsonable_encoder({"data": models})

api.add_resource(ModelProviderAvailableModelApi, "/workspaces/current/models/model-types/<string:model_type>")

0 comments on commit d7854a5

Please sign in to comment.