enh: tools user info

This commit is contained in:
Timothy Jaeryang Baek
2024-11-18 06:19:34 -08:00
parent 43ffd61aeb
commit c50b678dce
8 changed files with 113 additions and 49 deletions

View File

@@ -2,7 +2,7 @@ import time
from typing import Optional
from open_webui.apps.webui.internal.db import Base, get_db
from open_webui.apps.webui.models.groups import Groups
from open_webui.apps.webui.models.users import Users, UserResponse
from pydantic import BaseModel, ConfigDict
from sqlalchemy import BigInteger, Column, String, Text, JSON
@@ -57,6 +57,10 @@ class PromptModel(BaseModel):
####################
class PromptUserResponse(PromptModel):
user: Optional[UserResponse] = None
class PromptForm(BaseModel):
command: str
title: str
@@ -97,15 +101,21 @@ class PromptsTable:
except Exception:
return None
def get_prompts(self) -> list[PromptModel]:
def get_prompts(self) -> list[PromptUserResponse]:
with get_db() as db:
return [
PromptModel.model_validate(prompt) for prompt in db.query(Prompt).all()
PromptUserResponse.model_validate(
{
**PromptModel.model_validate(prompt).model_dump(),
"user": Users.get_user_by_id(prompt.user_id).model_dump(),
}
)
for prompt in db.query(Prompt).all()
]
def get_prompts_by_user_id(
self, user_id: str, permission: str = "write"
) -> list[PromptModel]:
) -> list[PromptUserResponse]:
prompts = self.get_prompts()
return [

View File

@@ -3,7 +3,7 @@ import time
from typing import Optional
from open_webui.apps.webui.internal.db import Base, JSONField, get_db
from open_webui.apps.webui.models.users import Users
from open_webui.apps.webui.models.users import Users, UserResponse
from open_webui.env import SRC_LOG_LEVELS
from pydantic import BaseModel, ConfigDict
from sqlalchemy import BigInteger, Column, String, Text, JSON
@@ -86,6 +86,10 @@ class ToolResponse(BaseModel):
created_at: int # timestamp in epoch
class ToolUserResponse(ToolResponse):
user: Optional[UserResponse] = None
class ToolForm(BaseModel):
id: str
name: str
@@ -134,13 +138,21 @@ class ToolsTable:
except Exception:
return None
def get_tools(self) -> list[ToolModel]:
def get_tools(self) -> list[ToolUserResponse]:
with get_db() as db:
return [ToolModel.model_validate(tool) for tool in db.query(Tool).all()]
return [
ToolUserResponse.model_validate(
{
**ToolModel.model_validate(tool).model_dump(),
"user": Users.get_user_by_id(tool.user_id).model_dump(),
}
)
for tool in db.query(Tool).order_by(Tool.updated_at.desc()).all()
]
def get_tools_by_user_id(
self, user_id: str, permission: str = "write"
) -> list[ToolModel]:
) -> list[ToolUserResponse]:
tools = self.get_tools()
return [

View File

@@ -1,6 +1,11 @@
from typing import Optional
from open_webui.apps.webui.models.prompts import PromptForm, PromptModel, Prompts
from open_webui.apps.webui.models.prompts import (
PromptForm,
PromptUserResponse,
PromptModel,
Prompts,
)
from open_webui.constants import ERROR_MESSAGES
from fastapi import APIRouter, Depends, HTTPException, status, Request
from open_webui.utils.utils import get_admin_user, get_verified_user
@@ -23,7 +28,7 @@ async def get_prompts(user=Depends(get_verified_user)):
return prompts
@router.get("/list", response_model=list[PromptModel])
@router.get("/list", response_model=list[PromptUserResponse])
async def get_prompt_list(user=Depends(get_verified_user)):
if user.role == "admin":
prompts = Prompts.get_prompts()

View File

@@ -2,7 +2,13 @@ import os
from pathlib import Path
from typing import Optional
from open_webui.apps.webui.models.tools import ToolForm, ToolModel, ToolResponse, Tools
from open_webui.apps.webui.models.tools import (
ToolForm,
ToolModel,
ToolResponse,
ToolUserResponse,
Tools,
)
from open_webui.apps.webui.utils import load_tools_module_by_id, replace_imports
from open_webui.config import CACHE_DIR, DATA_DIR
from open_webui.constants import ERROR_MESSAGES
@@ -19,7 +25,7 @@ router = APIRouter()
############################
@router.get("/", response_model=list[ToolResponse])
@router.get("/", response_model=list[ToolUserResponse])
async def get_tools(user=Depends(get_verified_user)):
if user.role == "admin":
tools = Tools.get_tools()
@@ -33,7 +39,7 @@ async def get_tools(user=Depends(get_verified_user)):
############################
@router.get("/list", response_model=list[ToolResponse])
@router.get("/list", response_model=list[ToolUserResponse])
async def get_tool_list(user=Depends(get_verified_user)):
if user.role == "admin":
tools = Tools.get_tools()