wip: 对于用户模型进行了优化,增加了权限模型和权限校验

This commit is contained in:
朝夕 2026-04-21 23:00:59 +08:00
parent 446e208193
commit 81da2e9f81
7 changed files with 231 additions and 11 deletions

View File

@ -18,7 +18,7 @@
- 】~~使用fastapi-users完善用户系统~~(2026/4/19 fastapi-users会严重摧毁代码的优雅性) - 】~~使用fastapi-users完善用户系统~~(2026/4/19 fastapi-users会严重摧毁代码的优雅性)
- [ ] 升级auth功能 - [ ] 升级auth功能
- [x] /pretor/api的接口函数进行重构 - [x] /pretor/api的接口函数进行重构
- [ ] /dockerfile待完善 - [x] /dockerfile待完善
- [ ] 完善沙箱功能 - [ ] 完善沙箱功能
- [ ] 完善爬虫功能 - [ ] 完善爬虫功能
- [ ] 对接更多的provider - [ ] 对接更多的provider
@ -30,7 +30,7 @@
- [x] /pretor/core/individual每个template进行优化 - [x] /pretor/core/individual每个template进行优化
- [ ] /pretor/worker_individual待完善复合子个体和基础子个体 - [ ] /pretor/worker_individual待完善复合子个体和基础子个体
- [ ] /pretor/api待完善 - [ ] /pretor/api待完善
- [ ] /dockerfile待完善 - [x] /dockerfile待完善
#### 2026/4/16 #### 2026/4/16
- [ ] 发布v0.1.0正式版 - [ ] 发布v0.1.0正式版
@ -42,7 +42,7 @@
- [ ] 完善爬虫功能 - [ ] 完善爬虫功能
- [ ] 对接更多的provider - [ ] 对接更多的provider
- [ ] 优化import - [ ] 优化import
- [ ] 升级auth功能 - [x] 升级auth功能
#### 2026/4/20 #### 2026/4/20
- [ ] 优化安全架构防止模型注入 - [ ] 优化安全架构防止模型注入

View File

@ -15,9 +15,12 @@
from typing import Union from typing import Union
from pretor.utils.ray_hook import ray_actor_hook from pretor.utils.ray_hook import ray_actor_hook
from fastapi import APIRouter, Request, Depends from fastapi import APIRouter, Depends
from pydantic import BaseModel from pydantic import BaseModel
from pretor.utils.access import Accessor, TokenData from pretor.utils.access import Accessor, TokenData
from pretor.core.database.table.individual import AgentType
from fastapi import HTTPException
from typing import Optional, List, Dict
agent_router = APIRouter(prefix="/api/v1/agent", tags=["agent"]) agent_router = APIRouter(prefix="/api/v1/agent", tags=["agent"])
@ -51,3 +54,86 @@ async def load_agent(agent_register: Union[AgentRegister, AgentLocalRegister],
case _: case _:
pass pass
return {"message": "创建成功"} return {"message": "创建成功"}
class WorkerIndividualCreate(BaseModel):
agent_name: str
agent_type: AgentType
description: str
provider_title: str
model_id: str
system_prompt: str
output_template: dict
bound_skill: Dict[str, List[str]]
workspace: List[str]
class WorkerIndividualUpdate(BaseModel):
agent_name: Optional[str] = None
agent_type: Optional[AgentType] = None
description: Optional[str] = None
provider_title: Optional[str] = None
model_id: Optional[str] = None
system_prompt: Optional[str] = None
output_template: Optional[dict] = None
bound_skill: Optional[Dict[str, List[str]]] = None
workspace: Optional[List[str]] = None
@agent_router.post("/worker")
async def create_worker_individual(worker_data: WorkerIndividualCreate,
token_data: TokenData = Depends(Accessor.get_current_user)):
postgres_database = ray_actor_hook("postgres_database")
data_dict = worker_data.model_dump()
data_dict["owner_id"] = token_data.user_id
worker = await postgres_database.add_worker_individual.remote(**data_dict)
return {"message": "success", "agent_id": worker.agent_id}
@agent_router.get("/worker")
async def get_worker_individual_list(token_data: TokenData = Depends(Accessor.get_current_user)):
postgres_database = ray_actor_hook("postgres_database")
workers = await postgres_database.get_worker_individual_list.remote(owner_id=token_data.user_id)
return {"workers": workers}
@agent_router.get("/worker/{agent_id}")
async def get_worker_individual(agent_id: str,
token_data: TokenData = Depends(Accessor.get_current_user)):
postgres_database = ray_actor_hook("postgres_database")
worker = await postgres_database.get_worker_individual.remote(agent_id=agent_id)
if not worker:
raise HTTPException(status_code=404, detail="Agent not found")
if worker.owner_id != token_data.user_id:
raise HTTPException(status_code=403, detail="Forbidden: You do not own this agent")
return worker
@agent_router.put("/worker/{agent_id}")
async def update_worker_individual(agent_id: str,
worker_data: WorkerIndividualUpdate,
token_data: TokenData = Depends(Accessor.get_current_user)):
postgres_database = ray_actor_hook("postgres_database")
worker = await postgres_database.get_worker_individual.remote(agent_id=agent_id)
if not worker:
raise HTTPException(status_code=404, detail="Agent not found")
if worker.owner_id != token_data.user_id:
raise HTTPException(status_code=403, detail="Forbidden: You do not own this agent")
update_data = worker_data.model_dump(exclude_unset=True)
updated_worker = await postgres_database.update_worker_individual.remote(agent_id=agent_id, **update_data)
return {"message": "success", "worker": updated_worker}
@agent_router.delete("/worker/{agent_id}")
async def delete_worker_individual(agent_id: str,
token_data: TokenData = Depends(Accessor.get_current_user)):
postgres_database = ray_actor_hook("postgres_database")
worker = await postgres_database.get_worker_individual.remote(agent_id=agent_id)
if not worker:
raise HTTPException(status_code=404, detail="Agent not found")
if worker.owner_id != token_data.user_id:
raise HTTPException(status_code=403, detail="Forbidden: You do not own this agent")
await postgres_database.delete_worker_individual.remote(agent_id=agent_id)
return {"message": "success"}

View File

@ -14,5 +14,63 @@
from pretor.core.database.table import WorkerIndividual from pretor.core.database.table import WorkerIndividual
from sqlmodel import select from sqlmodel import select
from pretor.utils.error import UserNotExistError, UserPasswordError from typing import List, Optional
from pretor.core.database.database_exception import database_exception from pretor.core.database.database_exception import database_exception
from ulid import ULID
class IndividualDatabase:
def __init__(self, async_session_maker):
self.async_session_maker = async_session_maker
@database_exception
async def add_worker_individual(self, **kwargs) -> WorkerIndividual:
async with self.async_session_maker() as session:
agent_id = str(ULID())
individual = WorkerIndividual(agent_id=agent_id, **kwargs)
session.add(individual)
await session.commit()
await session.refresh(individual)
return individual
@database_exception
async def get_worker_individual(self, agent_id: str) -> Optional[WorkerIndividual]:
async with self.async_session_maker() as session:
statement = select(WorkerIndividual).where(WorkerIndividual.agent_id == agent_id)
results = await session.execute(statement)
return results.scalar_one_or_none()
@database_exception
async def get_worker_individual_list(self, owner_id: str) -> List[WorkerIndividual]:
async with self.async_session_maker() as session:
statement = select(WorkerIndividual).where(WorkerIndividual.owner_id == owner_id)
results = await session.execute(statement)
return list(results.scalars().all())
@database_exception
async def update_worker_individual(self, agent_id: str, **kwargs) -> Optional[WorkerIndividual]:
async with self.async_session_maker() as session:
statement = select(WorkerIndividual).where(WorkerIndividual.agent_id == agent_id)
results = await session.execute(statement)
individual = results.scalar_one_or_none()
if not individual:
return None
for key, value in kwargs.items():
if value is not None:
setattr(individual, key, value)
session.add(individual)
await session.commit()
await session.refresh(individual)
return individual
@database_exception
async def delete_worker_individual(self, agent_id: str) -> bool:
async with self.async_session_maker() as session:
statement = select(WorkerIndividual).where(WorkerIndividual.agent_id == agent_id)
results = await session.execute(statement)
individual = results.scalar_one_or_none()
if not individual:
return False
session.delete(individual)
await session.commit()
return True

View File

@ -16,7 +16,7 @@ from pretor.core.database.table import User
from sqlmodel import select from sqlmodel import select
from pretor.utils.error import UserNotExistError, UserPasswordError from pretor.utils.error import UserNotExistError, UserPasswordError
from pretor.core.database.database_exception import database_exception from pretor.core.database.database_exception import database_exception
from pretor.core.database.table.user import UserAuthority
class AuthDatabase: class AuthDatabase:
def __init__(self, async_session_maker): def __init__(self, async_session_maker):
self.async_session_maker = async_session_maker self.async_session_maker = async_session_maker
@ -66,3 +66,9 @@ class AuthDatabase:
if user is None: if user is None:
raise UserNotExistError() raise UserNotExistError()
return user return user
@database_exception
async def get_user_authority(self, user_id: str) -> UserAuthority:
async with self.async_session_maker() as session:
user = await session.get(User, user_id)
return user.user_authority

View File

@ -19,6 +19,7 @@ from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
from sqlalchemy.orm import sessionmaker from sqlalchemy.orm import sessionmaker
from sqlmodel import SQLModel from sqlmodel import SQLModel
from pretor.core.database.module.individual import IndividualDatabase
from pretor.core.database.module.user import AuthDatabase from pretor.core.database.module.user import AuthDatabase
from pretor.core.database.module.provider import ProviderDatabase from pretor.core.database.module.provider import ProviderDatabase
@ -36,6 +37,7 @@ class PostgresDatabase:
self.auth_database = AuthDatabase(self.async_session_maker) self.auth_database = AuthDatabase(self.async_session_maker)
self.provider_database = ProviderDatabase(self.async_session_maker) self.provider_database = ProviderDatabase(self.async_session_maker)
self.individual_database = IndividualDatabase(self.async_session_maker)
async def init_db(self) -> None: async def init_db(self) -> None:
async with self.async_engine.begin() as conn: async with self.async_engine.begin() as conn:
@ -60,3 +62,22 @@ class PostgresDatabase:
async def login_user(self, **kwargs): async def login_user(self, **kwargs):
return await self.auth_database.login_user(**kwargs) return await self.auth_database.login_user(**kwargs)
async def get_user_authority(self, **kwargs):
return await self.auth_database.get_user_authority(**kwargs)
##individual_database 操作
async def add_worker_individual(self, **kwargs):
return await self.individual_database.add_worker_individual(**kwargs)
async def get_worker_individual(self, agent_id: str):
return await self.individual_database.get_worker_individual(agent_id)
async def get_worker_individual_list(self, owner_id: str):
return await self.individual_database.get_worker_individual_list(owner_id)
async def update_worker_individual(self, agent_id: str, **kwargs):
return await self.individual_database.update_worker_individual(agent_id, **kwargs)
async def delete_worker_individual(self, agent_id: str):
return await self.individual_database.delete_worker_individual(agent_id)

View File

@ -13,10 +13,19 @@
# limitations under the License. # limitations under the License.
from sqlmodel import SQLModel, Field from sqlmodel import SQLModel, Field
from enum import IntEnum
class UserAuthority(IntEnum):
SUPER_ADMINISTRATOR = 100
ADMINISTRATOR = 50
USER = 20
UNAUTHORIZED_USER = 10
GUEST = 0
class User(SQLModel): class User(SQLModel):
__tablename__ = 'user' __tablename__ = 'user'
user_id: int = Field(default=None, primary_key=True) user_id: str = Field(primary_key=True)
user_name: str = Field(index=True) user_name: str = Field(index=True)
hashed_password: str hashed_password: str
user_authority: UserAuthority = Field(default=UserAuthority.USER)

View File

@ -0,0 +1,40 @@
# Copyright 2026 zhaoxi826
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from functools import lru_cache
from typing import Annotated
from fastapi import Depends, HTTPException
from pretor.utils.access import Accessor, TokenData
from pretor.core.database.table.user import UserAuthority
from pretor.utils.ray_hook import ray_actor_hook
@lru_cache
async def get_authority(user_id: str) -> UserAuthority:
postgres_database = ray_actor_hook("postgres_database")
user_authority = await postgres_database.get_user_authority.remote(user_id=user_id)
return user_authority
class RoleChecker:
def __init__(self, **kwargs):
self.allowed_roles = kwargs.get("allowed_roles", )
async def __call__(self,
token_data: Annotated[TokenData, Depends(Accessor.get_current_user)]):
user_authority = await get_authority(token_data.user_id)
if user_authority < self.allowed_roles:
raise HTTPException(
status_code=403,
detail={"message": f"User {token_data.user_id} does not have allowed roles"},
)
return token_data