main.py 12.5 KB
Newer Older
Michael Poluektov's avatar
Michael Poluektov committed
1
from fastapi import FastAPI
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
2
from fastapi.responses import StreamingResponse
3
from fastapi.middleware.cors import CORSMiddleware
4
from apps.webui.routers import (
Timothy J. Baek's avatar
Timothy J. Baek committed
5
6
7
8
    auths,
    users,
    chats,
    documents,
Timothy J. Baek's avatar
Timothy J. Baek committed
9
    tools,
Timothy J. Baek's avatar
Timothy J. Baek committed
10
    models,
Timothy J. Baek's avatar
Timothy J. Baek committed
11
12
    prompts,
    configs,
Timothy J. Baek's avatar
Timothy J. Baek committed
13
    memories,
Timothy J. Baek's avatar
Timothy J. Baek committed
14
    utils,
Timothy J. Baek's avatar
Timothy J. Baek committed
15
    files,
Timothy J. Baek's avatar
Timothy J. Baek committed
16
    functions,
Timothy J. Baek's avatar
Timothy J. Baek committed
17
)
Timothy J. Baek's avatar
Timothy J. Baek committed
18
from apps.webui.models.functions import Functions
Timothy J. Baek's avatar
Timothy J. Baek committed
19
from apps.webui.models.models import Models
Timothy J. Baek's avatar
Timothy J. Baek committed
20
from apps.webui.utils import load_function_module_by_id
Timothy J. Baek's avatar
Timothy J. Baek committed
21

22
23
24
25
26
from utils.misc import (
    stream_message_template,
    whole_message_template,
    add_or_update_system_message,
)
Timothy J. Baek's avatar
Timothy J. Baek committed
27
28
from utils.task import prompt_template

Timothy J. Baek's avatar
Timothy J. Baek committed
29

Timothy J. Baek's avatar
Timothy J. Baek committed
30
from config import (
31
32
    SHOW_ADMIN_DETAILS,
    ADMIN_EMAIL,
Timothy J. Baek's avatar
Timothy J. Baek committed
33
34
35
36
37
    WEBUI_AUTH,
    DEFAULT_MODELS,
    DEFAULT_PROMPT_SUGGESTIONS,
    DEFAULT_USER_ROLE,
    ENABLE_SIGNUP,
38
    ENABLE_LOGIN_FORM,
Timothy J. Baek's avatar
Timothy J. Baek committed
39
    USER_PERMISSIONS,
Timothy J. Baek's avatar
Timothy J. Baek committed
40
    WEBHOOK_URL,
41
    WEBUI_AUTH_TRUSTED_EMAIL_HEADER,
42
    WEBUI_AUTH_TRUSTED_NAME_HEADER,
43
    JWT_EXPIRES_IN,
Timothy J. Baek's avatar
Timothy J. Baek committed
44
    WEBUI_BANNERS,
45
    ENABLE_COMMUNITY_SHARING,
Timothy J. Baek's avatar
Timothy J. Baek committed
46
    AppConfig,
47
    OAUTH_USERNAME_CLAIM,
Sergey Mihaylin's avatar
Sergey Mihaylin committed
48
    OAUTH_PICTURE_CLAIM,
Timothy J. Baek's avatar
Timothy J. Baek committed
49
)
50

Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
51
52
from apps.socket.main import get_event_call, get_event_emitter

Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
53
54
55
import inspect
import json

Michael Poluektov's avatar
Michael Poluektov committed
56
from typing import Iterator, Generator, AsyncGenerator
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
57
58
from pydantic import BaseModel

59
60
61
62
app = FastAPI()

origins = ["*"]

63
app.state.config = AppConfig()
Timothy J. Baek's avatar
Timothy J. Baek committed
64

65
app.state.config.ENABLE_SIGNUP = ENABLE_SIGNUP
66
app.state.config.ENABLE_LOGIN_FORM = ENABLE_LOGIN_FORM
67
app.state.config.JWT_EXPIRES_IN = JWT_EXPIRES_IN
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
68
app.state.AUTH_TRUSTED_EMAIL_HEADER = WEBUI_AUTH_TRUSTED_EMAIL_HEADER
69
app.state.AUTH_TRUSTED_NAME_HEADER = WEBUI_AUTH_TRUSTED_NAME_HEADER
70

71
72
73
74
75

app.state.config.SHOW_ADMIN_DETAILS = SHOW_ADMIN_DETAILS
app.state.config.ADMIN_EMAIL = ADMIN_EMAIL


76
77
78
79
80
app.state.config.DEFAULT_MODELS = DEFAULT_MODELS
app.state.config.DEFAULT_PROMPT_SUGGESTIONS = DEFAULT_PROMPT_SUGGESTIONS
app.state.config.DEFAULT_USER_ROLE = DEFAULT_USER_ROLE
app.state.config.USER_PERMISSIONS = USER_PERMISSIONS
app.state.config.WEBHOOK_URL = WEBHOOK_URL
Timothy J. Baek's avatar
Timothy J. Baek committed
81
app.state.config.BANNERS = WEBUI_BANNERS
Timothy J. Baek's avatar
Timothy J. Baek committed
82

83
app.state.config.ENABLE_COMMUNITY_SHARING = ENABLE_COMMUNITY_SHARING
Timothy J. Baek's avatar
Timothy J. Baek committed
84

85
86
87
app.state.config.OAUTH_USERNAME_CLAIM = OAUTH_USERNAME_CLAIM
app.state.config.OAUTH_PICTURE_CLAIM = OAUTH_PICTURE_CLAIM

Timothy J. Baek's avatar
Timothy J. Baek committed
88
app.state.MODELS = {}
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
89
app.state.TOOLS = {}
Timothy J. Baek's avatar
Timothy J. Baek committed
90
app.state.FUNCTIONS = {}
Timothy J. Baek's avatar
Timothy J. Baek committed
91

92
93
94
95
96
97
98
99
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

Timothy J. Baek's avatar
Timothy J. Baek committed
100
101

app.include_router(configs.router, prefix="/configs", tags=["configs"])
102
app.include_router(auths.router, prefix="/auths", tags=["auths"])
103
104
app.include_router(users.router, prefix="/users", tags=["users"])
app.include_router(chats.router, prefix="/chats", tags=["chats"])
Timothy J. Baek's avatar
Timothy J. Baek committed
105

Timothy J. Baek's avatar
Timothy J. Baek committed
106
app.include_router(documents.router, prefix="/documents", tags=["documents"])
Timothy J. Baek's avatar
Timothy J. Baek committed
107
app.include_router(models.router, prefix="/models", tags=["models"])
108
app.include_router(prompts.router, prefix="/prompts", tags=["prompts"])
Timothy J. Baek's avatar
Timothy J. Baek committed
109

Timothy J. Baek's avatar
Timothy J. Baek committed
110
app.include_router(memories.router, prefix="/memories", tags=["memories"])
Timothy J. Baek's avatar
Timothy J. Baek committed
111
112
113
app.include_router(files.router, prefix="/files", tags=["files"])
app.include_router(tools.router, prefix="/tools", tags=["tools"])
app.include_router(functions.router, prefix="/functions", tags=["functions"])
Timothy J. Baek's avatar
Timothy J. Baek committed
114

Timothy J. Baek's avatar
Timothy J. Baek committed
115
app.include_router(utils.router, prefix="/utils", tags=["utils"])
116
117
118
119


@app.get("/")
async def get_status():
Timothy J. Baek's avatar
Timothy J. Baek committed
120
121
122
    return {
        "status": True,
        "auth": WEBUI_AUTH,
123
124
        "default_models": app.state.config.DEFAULT_MODELS,
        "default_prompt_suggestions": app.state.config.DEFAULT_PROMPT_SUGGESTIONS,
Timothy J. Baek's avatar
Timothy J. Baek committed
125
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
126
127


Michael Poluektov's avatar
Michael Poluektov committed
128
129
130
131
132
133
134
135
136
137
138
139
140
141
def get_function_module(pipe_id: str):
    # Check if function is already loaded
    if pipe_id not in app.state.FUNCTIONS:
        function_module, _, _ = load_function_module_by_id(pipe_id)
        app.state.FUNCTIONS[pipe_id] = function_module
    else:
        function_module = app.state.FUNCTIONS[pipe_id]

    if hasattr(function_module, "valves") and hasattr(function_module, "Valves"):
        valves = Functions.get_function_valves_by_id(pipe_id)
        function_module.valves = function_module.Valves(**(valves if valves else {}))
    return function_module


142
143
async def get_pipe_models():
    pipes = Functions.get_functions_by_type("pipe", active_only=True)
Timothy J. Baek's avatar
Timothy J. Baek committed
144
145
146
    pipe_models = []

    for pipe in pipes:
Michael Poluektov's avatar
Michael Poluektov committed
147
        function_module = get_function_module(pipe.id)
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
148

Timothy J. Baek's avatar
Timothy J. Baek committed
149
150
        # Check if function is a manifold
        if hasattr(function_module, "type"):
Michael Poluektov's avatar
Michael Poluektov committed
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
            if not function_module.type == "manifold":
                continue
            manifold_pipes = []

            # Check if pipes is a function or a list
            if callable(function_module.pipes):
                manifold_pipes = function_module.pipes()
            else:
                manifold_pipes = function_module.pipes

            for p in manifold_pipes:
                manifold_pipe_id = f'{pipe.id}.{p["id"]}'
                manifold_pipe_name = p["name"]

                if hasattr(function_module, "name"):
                    manifold_pipe_name = f"{function_module.name}{manifold_pipe_name}"

                pipe_flag = {"type": pipe.type}
                if hasattr(function_module, "ChatValves"):
                    pipe_flag["valves_spec"] = function_module.ChatValves.schema()

                pipe_models.append(
                    {
                        "id": manifold_pipe_id,
                        "name": manifold_pipe_name,
                        "object": "model",
                        "created": pipe.created_at,
                        "owned_by": "openai",
                        "pipe": pipe_flag,
                    }
                )
Timothy J. Baek's avatar
Timothy J. Baek committed
182
        else:
Timothy J. Baek's avatar
Timothy J. Baek committed
183
184
185
186
            pipe_flag = {"type": "pipe"}
            if hasattr(function_module, "ChatValves"):
                pipe_flag["valves_spec"] = function_module.ChatValves.schema()

Timothy J. Baek's avatar
Timothy J. Baek committed
187
188
189
190
191
192
193
            pipe_models.append(
                {
                    "id": pipe.id,
                    "name": pipe.name,
                    "object": "model",
                    "created": pipe.created_at,
                    "owned_by": "openai",
Timothy J. Baek's avatar
Timothy J. Baek committed
194
                    "pipe": pipe_flag,
Timothy J. Baek's avatar
Timothy J. Baek committed
195
196
197
198
                }
            )

    return pipe_models
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
199
200


Michael Poluektov's avatar
Michael Poluektov committed
201
202
203
204
205
async def execute_pipe(pipe, params):
    if inspect.iscoroutinefunction(pipe):
        return await pipe(**params)
    else:
        return pipe(**params)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
206

Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
207

208
async def get_message_content(res: str | Generator | AsyncGenerator) -> str:
Michael Poluektov's avatar
Michael Poluektov committed
209
210
211
212
213
214
    if isinstance(res, str):
        return res
    if isinstance(res, Generator):
        return "".join(map(str, res))
    if isinstance(res, AsyncGenerator):
        return "".join([str(stream) async for stream in res])
Timothy J. Baek's avatar
Timothy J. Baek committed
215
216


Michael Poluektov's avatar
Michael Poluektov committed
217
218
219
220
221
222
def process_line(form_data: dict, line):
    if isinstance(line, BaseModel):
        line = line.model_dump_json()
        line = f"data: {line}"
    if isinstance(line, dict):
        line = f"data: {json.dumps(line)}"
Timothy J. Baek's avatar
Timothy J. Baek committed
223

Michael Poluektov's avatar
Michael Poluektov committed
224
225
226
227
    try:
        line = line.decode("utf-8")
    except Exception:
        pass
Timothy J. Baek's avatar
Timothy J. Baek committed
228

Michael Poluektov's avatar
Michael Poluektov committed
229
230
    if line.startswith("data:"):
        return f"{line}\n\n"
Timothy J. Baek's avatar
Timothy J. Baek committed
231
    else:
Michael Poluektov's avatar
Michael Poluektov committed
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
        line = stream_message_template(form_data["model"], line)
        return f"data: {json.dumps(line)}\n\n"


def get_pipe_id(form_data: dict) -> str:
    pipe_id = form_data["model"]
    if "." in pipe_id:
        pipe_id, _ = pipe_id.split(".", 1)
    print(pipe_id)
    return pipe_id


def get_params_dict(pipe, form_data, user, extra_params, function_module):
    pipe_id = get_pipe_id(form_data)
    # Get the signature of the function
    sig = inspect.signature(pipe)
    params = {"body": form_data}

    for key, value in extra_params.items():
        if key in sig.parameters:
            params[key] = value

    if "__user__" in sig.parameters:
        __user__ = {
            "id": user.id,
            "email": user.email,
            "name": user.name,
            "role": user.role,
        }

        try:
            if hasattr(function_module, "UserValves"):
                __user__["valves"] = function_module.UserValves(
                    **Functions.get_user_valves_by_id_and_user_id(pipe_id, user.id)
                )
        except Exception as e:
            print(e)
Timothy J. Baek's avatar
Timothy J. Baek committed
269

Michael Poluektov's avatar
Michael Poluektov committed
270
271
        params["__user__"] = __user__
    return params
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
272
273


274
def get_extra_params(metadata: dict):
275
276
277
    __event_emitter__ = None
    __event_call__ = None
    __task__ = None
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
278

Michael Poluektov's avatar
Michael Poluektov committed
279
280
281
282
283
284
    if metadata:
        if all(k in metadata for k in ("session_id", "chat_id", "message_id")):
            __event_emitter__ = get_event_emitter(metadata)
            __event_call__ = get_event_call(metadata)
        __task__ = metadata.get("task", None)

285
286
287
288
289
    return {
        "__event_emitter__": __event_emitter__,
        "__event_call__": __event_call__,
        "__task__": __task__,
    }
Michael Poluektov's avatar
Michael Poluektov committed
290
291


292
# inplace function: form_data is modified
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
def add_model_params(params: dict, form_data: dict) -> dict:
    if not params:
        return form_data

    mappings = {
        "temperature": float,
        "top_p": int,
        "max_tokens": int,
        "frequency_penalty": int,
        "seed": lambda x: x,
        "stop": lambda x: [bytes(s, "utf-8").decode("unicode_escape") for s in x],
    }

    for key, cast_func in mappings.items():
        if (value := params.get(key)) is not None:
            form_data[key] = cast_func(value)

    return form_data


313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
# inplace function: form_data is modified
def populate_system_message(params: dict, form_data: dict, user) -> dict:
    system = params.get("system", None)
    if not system:
        return form_data

    if user:
        template_params = {
            "user_name": user.name,
            "user_location": user.info.get("location") if user.info else None,
        }
    else:
        template_params = {}
    system = prompt_template(system, **template_params)
    form_data["messages"] = add_or_update_system_message(
        system, form_data.get("messages", [])
    )
    return form_data


333
334
335
336
337
async def generate_function_chat_completion(form_data, user):
    model_id = form_data.get("model")
    model_info = Models.get_model_by_id(model_id)
    metadata = form_data.pop("metadata", None)

338
339
    # Add extra params such as __event_emitter__
    extra_params = get_extra_params(metadata)
340
341
342
343
344
    if model_info:
        if model_info.base_model_id:
            form_data["model"] = model_info.base_model_id

        params = model_info.params.model_dump()
345
        form_data = add_model_params(params, form_data)
346
        form_data = populate_system_message(params, form_data, user)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
347

Michael Poluektov's avatar
Michael Poluektov committed
348
349
350
    async def job():
        pipe_id = get_pipe_id(form_data)
        function_module = get_function_module(pipe_id)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
351

Michael Poluektov's avatar
Michael Poluektov committed
352
353
        pipe = function_module.pipe
        params = get_params_dict(pipe, form_data, user, extra_params, function_module)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
354

Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
355
356
357
358
        if form_data["stream"]:

            async def stream_content():
                try:
Michael Poluektov's avatar
Michael Poluektov committed
359
                    res = await execute_pipe(pipe, params)
Timothy J. Baek's avatar
Timothy J. Baek committed
360
361
362
363
364
365
366
367
368
369

                    # Directly return if the response is a StreamingResponse
                    if isinstance(res, StreamingResponse):
                        async for data in res.body_iterator:
                            yield data
                        return
                    if isinstance(res, dict):
                        yield f"data: {json.dumps(res)}\n\n"
                        return

Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
370
371
372
373
374
375
376
377
378
379
380
                except Exception as e:
                    print(f"Error: {e}")
                    yield f"data: {json.dumps({'error': {'detail':str(e)}})}\n\n"
                    return

                if isinstance(res, str):
                    message = stream_message_template(form_data["model"], res)
                    yield f"data: {json.dumps(message)}\n\n"

                if isinstance(res, Iterator):
                    for line in res:
Michael Poluektov's avatar
Michael Poluektov committed
381
                        yield process_line(form_data, line)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
382

Timothy J. Baek's avatar
Timothy J. Baek committed
383
384
                if isinstance(res, AsyncGenerator):
                    async for line in res:
Michael Poluektov's avatar
Michael Poluektov committed
385
386
387
                        yield process_line(form_data, line)

                if isinstance(res, str) or isinstance(res, Generator):
388
                    finish_message = stream_message_template(form_data["model"], "")
389
                    finish_message["choices"][0]["finish_reason"] = "stop"
Michael Poluektov's avatar
Michael Poluektov committed
390
391
                    yield f"data: {json.dumps(finish_message)}\n\n"
                    yield "data: [DONE]"
Timothy J. Baek's avatar
Timothy J. Baek committed
392

Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
393
394
395
            return StreamingResponse(stream_content(), media_type="text/event-stream")
        else:
            try:
Michael Poluektov's avatar
Michael Poluektov committed
396
                res = await execute_pipe(pipe, params)
Timothy J. Baek's avatar
Timothy J. Baek committed
397

Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
398
399
400
401
            except Exception as e:
                print(f"Error: {e}")
                return {"error": {"detail": str(e)}}

Michael Poluektov's avatar
Michael Poluektov committed
402
            if isinstance(res, StreamingResponse) or isinstance(res, dict):
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
403
                return res
Michael Poluektov's avatar
Michael Poluektov committed
404
            if isinstance(res, BaseModel):
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
405
                return res.model_dump()
Michael Poluektov's avatar
Michael Poluektov committed
406

407
408
            message = await get_message_content(res)
            return whole_message_template(form_data["model"], message)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
409
410

    return await job()