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 [