tools.py 5.25 KB
Newer Older
1
2
from pydantic import BaseModel, ConfigDict
from typing import List, Optional
Timothy J. Baek's avatar
Timothy J. Baek committed
3
4
import time
import logging
5
6
from sqlalchemy import String, Column, BigInteger

7
from apps.webui.internal.db import Base, JSONField, Session
Timothy J. Baek's avatar
Timothy J. Baek committed
8
from apps.webui.models.users import Users
Timothy J. Baek's avatar
Timothy J. Baek committed
9
10

import json
11
12
import copy

Timothy J. Baek's avatar
Timothy J. Baek committed
13
14
15
16
17
18
19
20
21
22
23

from config import SRC_LOG_LEVELS

log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["MODELS"])

####################
# Tools DB Schema
####################


24
25
class Tool(Base):
    __tablename__ = "tool"
Timothy J. Baek's avatar
Timothy J. Baek committed
26

27
28
29
30
31
32
33
34
35
    id = Column(String, primary_key=True)
    user_id = Column(String)
    name = Column(String)
    content = Column(String)
    specs = Column(JSONField)
    meta = Column(JSONField)
    valves = Column(JSONField)
    updated_at = Column(BigInteger)
    created_at = Column(BigInteger)
Timothy J. Baek's avatar
Timothy J. Baek committed
36
37
38
39


class ToolMeta(BaseModel):
    description: Optional[str] = None
40
    manifest: Optional[dict] = {}
Timothy J. Baek's avatar
Timothy J. Baek committed
41
42
43
44
45
46
47


class ToolModel(BaseModel):
    id: str
    user_id: str
    name: str
    content: str
Timothy J. Baek's avatar
Timothy J. Baek committed
48
    specs: List[dict]
Timothy J. Baek's avatar
Timothy J. Baek committed
49
50
51
52
    meta: ToolMeta
    updated_at: int  # timestamp in epoch
    created_at: int  # timestamp in epoch

53
54
    model_config = ConfigDict(from_attributes=True)

Timothy J. Baek's avatar
Timothy J. Baek committed
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76

####################
# Forms
####################


class ToolResponse(BaseModel):
    id: str
    user_id: str
    name: str
    meta: ToolMeta
    updated_at: int  # timestamp in epoch
    created_at: int  # timestamp in epoch


class ToolForm(BaseModel):
    id: str
    name: str
    content: str
    meta: ToolMeta


Timothy J. Baek's avatar
Timothy J. Baek committed
77
78
79
80
class ToolValves(BaseModel):
    valves: Optional[dict] = None


Timothy J. Baek's avatar
Timothy J. Baek committed
81
82
83
class ToolsTable:

    def insert_new_tool(
84
        self, user_id: str, form_data: ToolForm, specs: List[dict]
Timothy J. Baek's avatar
Timothy J. Baek committed
85
86
87
88
89
90
91
92
93
94
95
96
    ) -> Optional[ToolModel]:
        tool = ToolModel(
            **{
                **form_data.model_dump(),
                "specs": specs,
                "user_id": user_id,
                "updated_at": int(time.time()),
                "created_at": int(time.time()),
            }
        )

        try:
97
98
99
100
101
102
103
104
            result = Tool(**tool.model_dump())
            Session.add(result)
            Session.commit()
            Session.refresh(result)
            if result:
                return ToolModel.model_validate(result)
            else:
                return None
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
105
106
        except Exception as e:
            print(f"Error creating tool: {e}")
Timothy J. Baek's avatar
Timothy J. Baek committed
107
108
            return None

109
    def get_tool_by_id(self, id: str) -> Optional[ToolModel]:
Timothy J. Baek's avatar
Timothy J. Baek committed
110
        try:
111
112
            tool = Session.get(Tool, id)
            return ToolModel.model_validate(tool)
Timothy J. Baek's avatar
Timothy J. Baek committed
113
114
115
        except:
            return None

116
    def get_tools(self) -> List[ToolModel]:
117
        return [ToolModel.model_validate(tool) for tool in Session.query(Tool).all()]
Timothy J. Baek's avatar
Timothy J. Baek committed
118

Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
119
    def get_tool_valves_by_id(self, id: str) -> Optional[dict]:
Timothy J. Baek's avatar
Timothy J. Baek committed
120
        try:
121
122
            tool = Session.get(Tool, id)
            return tool.valves if tool.valves else {}
Timothy J. Baek's avatar
Timothy J. Baek committed
123
124
125
126
127
128
        except Exception as e:
            print(f"An error occurred: {e}")
            return None

    def update_tool_valves_by_id(self, id: str, valves: dict) -> Optional[ToolValves]:
        try:
129
130
131
132
133
            Session.query(Tool).filter_by(id=id).update(
                {"valves": valves, "updated_at": int(time.time())}
            )
            Session.commit()
            return self.get_tool_by_id(id)
Timothy J. Baek's avatar
Timothy J. Baek committed
134
135
136
        except:
            return None

Timothy J. Baek's avatar
Timothy J. Baek committed
137
138
139
140
141
    def get_user_valves_by_id_and_user_id(
        self, id: str, user_id: str
    ) -> Optional[dict]:
        try:
            user = Users.get_user_by_id(user_id)
142
            user_settings = user.settings.model_dump()
Timothy J. Baek's avatar
Timothy J. Baek committed
143
144

            # Check if user has "tools" and "valves" settings
145
146
147
148
            if "tools" not in user_settings:
                user_settings["tools"] = {}
            if "valves" not in user_settings["tools"]:
                user_settings["tools"]["valves"] = {}
Timothy J. Baek's avatar
Timothy J. Baek committed
149

150
            return user_settings["tools"]["valves"].get(id, {})
Timothy J. Baek's avatar
Timothy J. Baek committed
151
152
153
154
155
156
157
158
159
        except Exception as e:
            print(f"An error occurred: {e}")
            return None

    def update_user_valves_by_id_and_user_id(
        self, id: str, user_id: str, valves: dict
    ) -> Optional[dict]:
        try:
            user = Users.get_user_by_id(user_id)
160
            user_settings = user.settings.model_dump()
Timothy J. Baek's avatar
Timothy J. Baek committed
161
162

            # Check if user has "tools" and "valves" settings
163
164
165
166
            if "tools" not in user_settings:
                user_settings["tools"] = {}
            if "valves" not in user_settings["tools"]:
                user_settings["tools"]["valves"] = {}
Timothy J. Baek's avatar
Timothy J. Baek committed
167

168
            user_settings["tools"]["valves"][id] = valves
Timothy J. Baek's avatar
Timothy J. Baek committed
169
170

            # Update the user settings in the database
171
            Users.update_user_by_id(user_id, {"settings": user_settings})
Timothy J. Baek's avatar
Timothy J. Baek committed
172

173
            return user_settings["tools"]["valves"][id]
Timothy J. Baek's avatar
Timothy J. Baek committed
174
175
176
177
        except Exception as e:
            print(f"An error occurred: {e}")
            return None

Timothy J. Baek's avatar
Timothy J. Baek committed
178
179
    def update_tool_by_id(self, id: str, updated: dict) -> Optional[ToolModel]:
        try:
180
181
182
183
184
185
            tool = Session.get(Tool, id)
            tool.update(**updated)
            tool.updated_at = int(time.time())
            Session.commit()
            Session.refresh(tool)
            return ToolModel.model_validate(tool)
Timothy J. Baek's avatar
Timothy J. Baek committed
186
187
188
        except:
            return None

189
    def delete_tool_by_id(self, id: str) -> bool:
Timothy J. Baek's avatar
Timothy J. Baek committed
190
        try:
191
            Session.query(Tool).filter_by(id=id).delete()
Timothy J. Baek's avatar
Timothy J. Baek committed
192
193
194
195
196
            return True
        except:
            return False


197
Tools = ToolsTable()