add simple access control

.env.example

@@ -6,6 +6,8 @@ APP_SERVER_HOST=0.0.0.0
 APP_SERVER_PORT=8086
 APP_SERVER_WORKERS=1
 APP_API_PREFIX=/api
+APP_AUTH_ENABLE=False
+APP_AUTH_ADMIN_TOKEN=admin
 
 LOG_LEVEL=DEBUG
 

README.md

@@ -101,6 +101,15 @@ Before running, you need to run `pip install openai` to install the Python `open
 python tests/e2e/index.py
 ```
 
+### Permissions
+Simple user isolation is provided based on tokens to meet SaaS deployment requirements. It can be enabled by configuring `APP_AUTH_ENABLE`.
+
+![](docs/imgs/user.png)
+
+1. The authentication method is Bearer token. You can include `Authorization: Bearer ***` in the header for authentication.
+2. Token management is described in the token section of the API documentation. Relevant APIs need to be authenticated with an admin token, which is configured as `APP_AUTH_ADMIN_TOKEN` and defaults to "admin".
+3. When creating a token, you need to provide the base URL and API key of the large model. The created assistant will use the corresponding configuration to access the large model.
+
 ## Community and Support
 
 - Join the [Slack](https://join.slack.com/t/openassistant-qbu7007/shared_invite/zt-29t8j9y12-9og5KZL6GagXTEvbEDf6UQ)

README_CN.md

@@ -97,6 +97,16 @@ Api Base URL: http://127.0.0.1:8086/api/v1
 python tests/e2e/index.py
 ```
 
+### 权限
+基于 token 提供简单用户隔离，满足 SaaS 部署需求，可通过配置 ```APP_AUTH_ENABLE``` 开启
+
+![](docs/imgs/user.png)
+
+1. 验证方式为 Bearer token，可在 Header 中填入 ```Authorization: Bearer ***``` 进行验证
+2. token 管理参考 api 文档中的 token 小节  
+相关 api 需通过 admin token 验证，配置为 ```APP_AUTH_ADMIN_TOKEN```，默认为 admin
+3. 创建 token 需填入大模型 base_url 和 api_key，创建的 assistant 将使用相关配置访问大模型
+
 ## 社区与支持
 
 - 加入 [Slack](https://join.slack.com/t/openassistant-qbu7007/shared_invite/zt-29t8j9y12-9og5KZL6GagXTEvbEDf6UQ)

app/api/deps.py

@@ -1,8 +1,103 @@
 from sqlmodel import Session
 
+from fastapi import Depends, Request
+
+from fastapi.security import APIKeyHeader
+from config.config import settings
+from app.exceptions.exception import AuthenticationError, AuthorizationError, ResourceNotFoundError
 from app.providers import database
+from app.models.token_relation import RelationType, TokenRelationQuery
+from app.models.token import Token
+from app.services.token.token_relation import TokenRelationService
+from app.services.token.token import TokenService
 
 
 def get_session():
     with Session(database.engine) as session:
         yield session
+
+
+class OAuth2Bearer(APIKeyHeader):
+    """
+    it use to fetch token from header
+    """
+
+    def __init__(
+        self, *, name: str, scheme_name: str | None = None, description: str | None = None, auto_error: bool = True
+    ):
+        super().__init__(name=name, scheme_name=scheme_name, description=description, auto_error=auto_error)
+
+    async def __call__(self, request: Request) -> str:
+        authorization_header_value = request.headers.get(self.model.name)
+        if authorization_header_value:
+            scheme, _, param = authorization_header_value.partition(" ")
+            if scheme.lower() == "bearer" and param.strip() != "":
+                return param.strip()
+        return None
+
+
+oauth_token = OAuth2Bearer(name="Authorization")
+
+
+async def verify_admin_token(token=Depends(oauth_token)) -> Token:
+    """
+    admin token authentication
+    """
+    if token is None:
+        raise AuthenticationError()
+    if settings.AUTH_ADMIN_TOKEN != token:
+        raise AuthorizationError()
+
+
+async def get_token(session=Depends(get_session), token=Depends(oauth_token)) -> Token:
+    """
+    get token info
+    """
+    if token and token != "":
+        try:
+            return TokenService.get_token(session=session, token=token)
+        except ResourceNotFoundError:
+            pass
+    return None
+
+
+async def verfiy_token(token: Token = Depends(get_token)):
+    if token is None:
+        raise AuthenticationError()
+
+
+async def get_token_id(token: Token = Depends(get_token)):
+    """
+    Return token_id, which can be considered as user information.
+    """
+    return token.id if token is not None else None
+
+
+def get_param(name: str):
+    """
+    extract param from Request
+    """
+
+    async def get_param_from_request(request: Request):
+        if name in request.path_params:
+            return request.path_params[name]
+        if name in request.query_params:
+            return request.query_params[name]
+        body = await request.json()
+        if name in body:
+            return body[name]
+
+    return get_param_from_request
+
+
+def verify_token_relation(relation_type: RelationType, name: str):
+    async def verify_authorization(
+        session=Depends(get_session), token_id=Depends(get_token_id), relation_id=Depends(get_param(name))
+    ):
+        if token_id and relation_id:
+            verify = TokenRelationQuery(token_id=token_id, relation_type=relation_type, relation_id=relation_id)
+            if TokenRelationService.verify_relation(session=session, verify=verify):
+                return
+        raise AuthorizationError()
+
+    return verify_authorization

app/api/routes.py

@@ -1,11 +1,14 @@
 from fastapi import APIRouter
-from app.api.v1 import assistant, assistant_file, thread, message, files, runs
+from app.api.v1 import assistant, assistant_file, thread, message, files, runs, token
 
 api_router = APIRouter(prefix="/v1")
 
-api_router.include_router(assistant.router, prefix="/assistants", tags=["assistants"])
-api_router.include_router(assistant_file.router, prefix="/assistants", tags=["assistants"])
-api_router.include_router(thread.router, prefix="/threads", tags=["threads"])
-api_router.include_router(message.router, prefix="/threads", tags=["messages"])
-api_router.include_router(runs.router, prefix="/threads", tags=["runs"])
-api_router.include_router(files.router, prefix="/files", tags=["files"])
+
+def router_init():
+    api_router.include_router(assistant.router, prefix="/assistants", tags=["assistants"])
+    api_router.include_router(assistant_file.router, prefix="/assistants", tags=["assistants"])
+    api_router.include_router(thread.router, prefix="/threads", tags=["threads"])
+    api_router.include_router(message.router, prefix="/threads", tags=["messages"])
+    api_router.include_router(runs.router, prefix="/threads", tags=["runs"])
+    api_router.include_router(files.router, prefix="/files", tags=["files"])
+    api_router.include_router(token.router, prefix="/tokens", tags=["tokens"])

app/api/v1/assistant.py

@@ -1,29 +1,36 @@
 from fastapi import APIRouter, Depends
 from sqlmodel import Session, select
 
-from app.api.deps import get_session
+from app.api.deps import get_token_id, get_session
 from app.models.assistant import Assistant, AssistantUpdate, AssistantCreate
 from app.libs.paginate import cursor_page, CommonPage
+from app.models.token_relation import RelationType
+from app.providers.auth_provider import auth_policy
 from app.schemas.common import DeleteResponse
 from app.services.assistant.assistant import AssistantService
 
 router = APIRouter()
 
 
 @router.get("", response_model=CommonPage[Assistant])
-def list_assistants(*, session: Session = Depends(get_session)):
+def list_assistants(*, session: Session = Depends(get_session), token_id=Depends(get_token_id)):
     """
     Returns a list of assistants.
     """
-    return cursor_page(select(Assistant), session)
+    statement = auth_policy.token_filter(
+        select(Assistant), field=Assistant.id, relation_type=RelationType.Assistant, token_id=token_id
+    )
+    return cursor_page(statement, session)
 
 
 @router.post("", response_model=Assistant)
-def create_assistant(*, session: Session = Depends(get_session), body: AssistantCreate) -> Assistant:
+def create_assistant(
+    *, session: Session = Depends(get_session), body: AssistantCreate, token_id=Depends(get_token_id)
+) -> Assistant:
     """
     Create an assistant with a model and instructions.
     """
-    return AssistantService.create_assistant(session=session, body=body)
+    return AssistantService.create_assistant(session=session, body=body, token_id=token_id)
 
 
 @router.get("/{assistant_id}", response_model=Assistant)

app/api/v1/thread.py

@@ -1,7 +1,7 @@
 from fastapi import APIRouter, Depends
 from sqlmodel import Session
 
-from app.api.deps import get_session
+from app.api.deps import get_session, get_token_id
 from app.models.thread import Thread, ThreadUpdate, ThreadCreate
 from app.schemas.common import DeleteResponse
 from app.services.thread.thread import ThreadService
@@ -10,11 +10,13 @@
 
 
 @router.post("", response_model=Thread)
-def create_thread(*, session: Session = Depends(get_session), body: ThreadCreate) -> Thread:
+def create_thread(
+    *, session: Session = Depends(get_session), body: ThreadCreate, token_id=Depends(get_token_id)
+) -> Thread:
     """
     Create a thread.
     """
-    return ThreadService.create_thread(session=session, body=body)
+    return ThreadService.create_thread(session=session, body=body, token_id=token_id)
 
 
 @router.get("/{thread_id}", response_model=Thread)

app/api/v1/token.py

@@ -0,0 +1,44 @@
+from sqlmodel import Session, select
+from fastapi import APIRouter, Depends
+
+from app.api.deps import get_session, verify_admin_token
+from app.libs.paginate import CommonPage, cursor_page
+from app.models.token import Token, TokenCreate, TokenUpdate
+from app.services.token.token import TokenService
+
+router = APIRouter()
+
+
+@router.get("/list", response_model=CommonPage[Token], dependencies=[Depends(verify_admin_token)])
+def list_tokens(*, session: Session = Depends(get_session)):
+    """
+    Returns a list of tokens.
+    """
+    statement = select(Token)
+    return cursor_page(statement, session)
+
+
+@router.post("", response_model=Token, dependencies=[Depends(verify_admin_token)])
+def create_token(*, session: Session = Depends(get_session), body: TokenCreate) -> Token:
+    """
+    Create a token with a llm url & token.
+    """
+    return TokenService.create_token(session=session, body=body)
+
+
+@router.get("", response_model=Token, dependencies=[Depends(verify_admin_token)])
+def get_token(*, session: Session = Depends(get_session), token: str) -> Token:
+    """
+    Retrieves a token.
+    """
+    return TokenService.get_token(session=session, token=token)
+
+
+@router.get("/refresh_token", response_model=Token, dependencies=[Depends(verify_admin_token)])
+def refresh_token(*, session: Session = Depends(get_session), token: str) -> Token:
+    return TokenService.refresh_token(session=session, token=token)
+
+
+@router.post("/modify_token", response_model=Token, dependencies=[Depends(verify_admin_token)])
+def modify_token(*, session: Session = Depends(get_session), update: TokenUpdate) -> Token:
+    return TokenService.modify_token(session=session, update=update)

app/core/runner/llm_backend.py

@@ -1,24 +1,21 @@
 import logging
 from typing import List
 
-from openai import OpenAI
-from openai import Stream
+from openai import OpenAI, Stream
 from openai.types.chat import ChatCompletionChunk, ChatCompletion
 
-from config.llm import LLMSettings
-
 
 class LLMBackend:
     """
     openai chat 接口封装
     """
 
-    def __init__(self, llm_settings: LLMSettings) -> None:
-        self.base_url = llm_settings.OPENAI_API_BASE if llm_settings.OPENAI_API_BASE else None
-        self.api_key = llm_settings.OPENAI_API_KEY
+    def __init__(self, base_url: str, api_key) -> None:
+        self.base_url = base_url + "/" if base_url else None
+        self.api_key = api_key
         self.client = OpenAI(
-          api_key=self.api_key,
-          base_url=self.base_url
+          base_url=self.base_url,
+          api_key=self.api_key 
         )
 
     def run(
@@ -32,7 +29,7 @@ def run(
         if tools:
             chat_params['tools'] = tools
             chat_params['tool_choice'] = tool_choice if tool_choice else "auto"
-        logging.info(f"chat_params: {chat_params}")
+        logging.info("chat_params: %s", chat_params)
         response = self.client.chat.completions.create(**chat_params)
-        logging.info(f"chat_response: {response}")
+        logging.info("chat_response: %s", response)
         return response