wip: 对于用户模型进行了优化,增加了权限模型和权限校验
This commit is contained in:
parent
446e208193
commit
81da2e9f81
|
|
@ -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
|
||||||
- [ ] 优化安全架构防止模型注入
|
- [ ] 优化安全架构防止模型注入
|
||||||
|
|
|
||||||
|
|
@ -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"])
|
||||||
|
|
||||||
|
|
@ -50,4 +53,87 @@ async def load_agent(agent_register: Union[AgentRegister, AgentLocalRegister],
|
||||||
node.create_agent.remote(global_state_machine,agent_register.provider_title,agent_register.model_id)
|
node.create_agent.remote(global_state_machine,agent_register.provider_title,agent_register.model_id)
|
||||||
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"}
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
@ -65,4 +65,10 @@ class AuthDatabase:
|
||||||
user = results.scalar_one_or_none()
|
user = results.scalar_one_or_none()
|
||||||
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
|
||||||
|
|
@ -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:
|
||||||
|
|
@ -59,4 +61,23 @@ class PostgresDatabase:
|
||||||
return await self.auth_database.delete_user(**kwargs)
|
return await self.auth_database.delete_user(**kwargs)
|
||||||
|
|
||||||
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)
|
||||||
|
|
@ -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)
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
||||||
Loading…
Reference in New Issue