Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
101 changes: 91 additions & 10 deletions packages/dbgpt-serve/src/dbgpt_serve/agent/hub/controller.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,10 @@
import logging
from abc import ABC
from typing import List
from functools import cache
from typing import List, Optional

from fastapi import APIRouter, Body, File, UploadFile
from fastapi import APIRouter, Body, Depends, File, HTTPException, UploadFile
from fastapi.security.http import HTTPAuthorizationCredentials, HTTPBearer

from dbgpt.agent.resource.tool.autogpt.plugins_util import scan_plugins
from dbgpt.agent.resource.tool.pack import AutoGPTPluginToolPack
Expand All @@ -24,6 +26,57 @@
router = APIRouter()
logger = logging.getLogger(__name__)

get_bearer_token = HTTPBearer(auto_error=False)
_api_keys_config: Optional[str] = None


def set_api_keys_config(api_keys: Optional[str]):
global _api_keys_config
_api_keys_config = api_keys


@cache
def _parse_api_keys(api_keys: str) -> List[str]:
if not api_keys:
return []
return [key.strip() for key in api_keys.split(",")]


async def check_api_key(
auth: Optional[HTTPAuthorizationCredentials] = Depends(get_bearer_token),
) -> Optional[str]:
global _api_keys_config
if _api_keys_config:
api_keys = _parse_api_keys(_api_keys_config)
if auth is None or (token := auth.credentials) not in api_keys:
raise HTTPException(
status_code=401,
detail={
"error": {
"message": "",
"type": "invalid_request_error",
"param": None,
"code": "invalid_api_key",
}
},
)
return token
return None


def _validate_user_access(
requested_user: Optional[str], authenticated_user: Optional[str]
) -> str:
global _api_keys_config
if _api_keys_config and authenticated_user:
if requested_user and requested_user != authenticated_user:
raise HTTPException(
status_code=403,
detail="Access denied: Cannot access another user's resources",
)
return requested_user or authenticated_user
return requested_user or "default"


class ModulePlugin(BaseComponent, ABC):
name = ComponentType.PLUGIN_HUB
Expand All @@ -34,6 +87,16 @@ def __init__(self):

def init_app(self, system_app: SystemApp):
system_app.app.include_router(router, prefix="/api", tags=["Agent"])
try:
global_api_keys = system_app.config.get("dbgpt.app.global.api_keys")
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

hi @geckosecurity global_api_keys variable maybe always None, it seem like ModulePlugin is initialized before the"dbgpt.app.global.api_keys"

if global_api_keys:
set_api_keys_config(global_api_keys)
elif hasattr(system_app, "config") and hasattr(
system_app.config, "API_KEYS"
):
set_api_keys_config(system_app.config.API_KEYS)
except Exception:
pass

def refresh_plugins(self):
self.plugins = scan_plugins(PLUGINS_DIR)
Expand Down Expand Up @@ -89,18 +152,26 @@ async def get_agent_list(filter: PagenationFilter[PluginHubFilter] = Body()):


@router.post("/v1/agent/my", response_model=Result[List[MyPluginVO]])
async def my_agents(user: str = None):
async def my_agents(
user: str = None, authenticated_user: Optional[str] = Depends(check_api_key)
):
logger.info(f"my_agents:{user}")
agents = plugin_hub.get_my_plugin(user)
effective_user = _validate_user_access(user, authenticated_user)
agents = plugin_hub.get_my_plugin(effective_user)
agent_dicts = MyPluginEntity.to_vo(agents)
return Result.succ(agent_dicts)


@router.post("/v1/agent/install", response_model=Result[str])
async def agent_install(plugin_name: str, user: str = None):
async def agent_install(
plugin_name: str,
user: str = None,
authenticated_user: Optional[str] = Depends(check_api_key),
):
logger.info(f"agent_install:{plugin_name},{user}")
effective_user = _validate_user_access(user, authenticated_user)
try:
plugin_hub.install_plugin(plugin_name, user)
plugin_hub.install_plugin(plugin_name, effective_user)

module_plugin.refresh_plugins()

Expand All @@ -111,10 +182,15 @@ async def agent_install(plugin_name: str, user: str = None):


@router.post("/v1/agent/uninstall", response_model=Result[str])
async def agent_uninstall(plugin_name: str, user: str = None):
async def agent_uninstall(
plugin_name: str,
user: str = None,
authenticated_user: Optional[str] = Depends(check_api_key),
):
logger.info(f"agent_uninstall:{plugin_name},{user}")
effective_user = _validate_user_access(user, authenticated_user)
try:
plugin_hub.uninstall_plugin(plugin_name, user)
plugin_hub.uninstall_plugin(plugin_name, effective_user)

module_plugin.refresh_plugins()
return Result.succ(None)
Expand All @@ -124,10 +200,15 @@ async def agent_uninstall(plugin_name: str, user: str = None):


@router.post("/v1/personal/agent/upload", response_model=Result[str])
async def personal_agent_upload(doc_file: UploadFile = File(...), user: str = None):
async def personal_agent_upload(
doc_file: UploadFile = File(...),
user: str = None,
authenticated_user: Optional[str] = Depends(check_api_key),
):
logger.info(f"personal_agent_upload:{doc_file.filename},{user}")
effective_user = _validate_user_access(user, authenticated_user)
try:
await plugin_hub.upload_my_plugin(doc_file, user)
await plugin_hub.upload_my_plugin(doc_file, effective_user)
module_plugin.refresh_plugins()
return Result.succ(None)
except Exception as e:
Expand Down