main.py 12.7 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
from utils.misc import stream_message_template, whole_message_template
Timothy J. Baek's avatar
Timothy J. Baek committed
23
24
from utils.task import prompt_template

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

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

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

Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
49
50
51
52
53
import inspect
import uuid
import time
import json

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

57
58
59
60
app = FastAPI()

origins = ["*"]

61
app.state.config = AppConfig()
Timothy J. Baek's avatar
Timothy J. Baek committed
62

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

69
70
71
72
73

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


74
75
76
77
78
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
79
app.state.config.BANNERS = WEBUI_BANNERS
Timothy J. Baek's avatar
Timothy J. Baek committed
80

81
app.state.config.ENABLE_COMMUNITY_SHARING = ENABLE_COMMUNITY_SHARING
Timothy J. Baek's avatar
Timothy J. Baek committed
82

83
84
85
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
86
app.state.MODELS = {}
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
87
app.state.TOOLS = {}
Timothy J. Baek's avatar
Timothy J. Baek committed
88
app.state.FUNCTIONS = {}
Timothy J. Baek's avatar
Timothy J. Baek committed
89

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

Timothy J. Baek's avatar
Timothy J. Baek committed
98
99

app.include_router(configs.router, prefix="/configs", tags=["configs"])
100
app.include_router(auths.router, prefix="/auths", tags=["auths"])
101
102
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
103

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

Timothy J. Baek's avatar
Timothy J. Baek committed
108
app.include_router(memories.router, prefix="/memories", tags=["memories"])
Timothy J. Baek's avatar
Timothy J. Baek committed
109
110
111
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
112

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


@app.get("/")
async def get_status():
Timothy J. Baek's avatar
Timothy J. Baek committed
118
119
120
    return {
        "status": True,
        "auth": WEBUI_AUTH,
121
122
        "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
123
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
124
125


Michael Poluektov's avatar
Michael Poluektov committed
126
127
128
129
130
131
132
133
134
135
136
137
138
139
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


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

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

Timothy J. Baek's avatar
Timothy J. Baek committed
147
148
        # Check if function is a manifold
        if hasattr(function_module, "type"):
Michael Poluektov's avatar
Michael Poluektov committed
149
150
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
            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
180
        else:
Timothy J. Baek's avatar
Timothy J. Baek committed
181
182
183
184
            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
185
186
187
188
189
190
191
            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
192
                    "pipe": pipe_flag,
Timothy J. Baek's avatar
Timothy J. Baek committed
193
194
195
196
                }
            )

    return pipe_models
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
197
198


Michael Poluektov's avatar
Michael Poluektov committed
199
200
201
202
203
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
204

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

206
async def get_message_content(res: str | Generator | AsyncGenerator) -> str:
Michael Poluektov's avatar
Michael Poluektov committed
207
208
209
210
211
212
    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
213
214


Michael Poluektov's avatar
Michael Poluektov committed
215
216
217
218
219
220
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
221

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

Michael Poluektov's avatar
Michael Poluektov committed
227
228
    if line.startswith("data:"):
        return f"{line}\n\n"
Timothy J. Baek's avatar
Timothy J. Baek committed
229
    else:
Michael Poluektov's avatar
Michael Poluektov committed
230
231
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
        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
267

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


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

Michael Poluektov's avatar
Michael Poluektov committed
277
278
279
280
281
282
    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)

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


290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
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


310
311
312
313
async def generate_function_chat_completion(form_data, user):
    print("entry point")
    model_id = form_data.get("model")
    model_info = Models.get_model_by_id(model_id)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
314

315
316
317
318
319
320
321
322
323
    metadata = form_data.pop("metadata", None)
    extra_params = get_extra_params(metadata)

    if model_info:
        if model_info.base_model_id:
            form_data["model"] = model_info.base_model_id

        params = model_info.params.model_dump()
        system = params.get("system", None)
324
325
        form_data = add_model_params(params, form_data)

326
327
328
329
330
331
332
333
        if system:
            if user:
                template_params = {
                    "user_name": user.name,
                    "user_location": user.info.get("location") if user.info else None,
                }
            else:
                template_params = {}
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
334

335
            system = prompt_template(system, **template_params)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
336

337
338
339
340
341
342
343
344
345
346
347
            # Check if the payload already has a system message
            # If not, add a system message to the payload
            for message in form_data.get("messages", []):
                if message.get("role") == "system":
                    message["content"] = system + message["content"]
                    break
            else:
                if form_data.get("messages"):
                    form_data["messages"].insert(
                        0, {"role": "system", "content": system}
                    )
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
348

Michael Poluektov's avatar
Michael Poluektov committed
349
350
351
    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
352

Michael Poluektov's avatar
Michael Poluektov committed
353
354
        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
355

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

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

                    # 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
371
372
373
374
375
376
377
378
379
380
381
                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
382
                        yield process_line(form_data, line)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
383

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

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

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

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

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

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

    return await job()