main.py 29.9 KB
Newer Older
Timothy J. Baek's avatar
Timothy J. Baek committed
1
2
3
4
5
6
7
8
from fastapi import (
    FastAPI,
    Request,
    HTTPException,
    Depends,
    UploadFile,
    File,
)
9
10
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
Timothy J. Baek's avatar
Timothy J. Baek committed
11

Timothy J. Baek's avatar
Timothy J. Baek committed
12
from pydantic import BaseModel, ConfigDict
13

Timothy J. Baek's avatar
Timothy J. Baek committed
14
import os
15
import re
16
import random
Timothy J. Baek's avatar
Timothy J. Baek committed
17
18
import requests
import json
19
20
import aiohttp
import asyncio
21
import logging
22
import time
Timothy J. Baek's avatar
Timothy J. Baek committed
23
24
25
from urllib.parse import urlparse
from typing import Optional, List, Union

26
27
from starlette.background import BackgroundTask

28
from apps.webui.models.models import Models
29
from constants import ERROR_MESSAGES
30
31
32
33
from utils.utils import (
    get_verified_user,
    get_admin_user,
)
Timothy J. Baek's avatar
Timothy J. Baek committed
34

Timothy J. Baek's avatar
Timothy J. Baek committed
35
36
37
from config import (
    SRC_LOG_LEVELS,
    OLLAMA_BASE_URLS,
Timothy J. Baek's avatar
Timothy J. Baek committed
38
    ENABLE_OLLAMA_API,
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
39
    AIOHTTP_CLIENT_TIMEOUT,
Timothy J. Baek's avatar
Timothy J. Baek committed
40
    ENABLE_MODEL_FILTER,
Timothy J. Baek's avatar
Timothy J. Baek committed
41
42
    MODEL_FILTER_LIST,
    UPLOAD_DIR,
43
    AppConfig,
Timothy J. Baek's avatar
Timothy J. Baek committed
44
)
Michael Poluektov's avatar
Michael Poluektov committed
45
46
47
48
49
50
from utils.misc import (
    apply_model_params_to_body_ollama,
    calculate_sha256,
    apply_model_params_to_body_openai,
    apply_model_system_prompt_to_body,
)
51

52
53
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["OLLAMA"])
54

55
56
57
58
59
60
61
62
app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)
Timothy J. Baek's avatar
Timothy J. Baek committed
63

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

Timothy J. Baek's avatar
Timothy J. Baek committed
66
67
app.state.config.ENABLE_MODEL_FILTER = ENABLE_MODEL_FILTER
app.state.config.MODEL_FILTER_LIST = MODEL_FILTER_LIST
Timothy J. Baek's avatar
Timothy J. Baek committed
68
69

app.state.config.ENABLE_OLLAMA_API = ENABLE_OLLAMA_API
70
app.state.config.OLLAMA_BASE_URLS = OLLAMA_BASE_URLS
71
app.state.MODELS = {}
Timothy J. Baek's avatar
Timothy J. Baek committed
72
73


Timothy J. Baek's avatar
Timothy J. Baek committed
74
75
76
77
78
# TODO: Implement a more intelligent load balancing mechanism for distributing requests among multiple backend instances.
# Current implementation uses a simple round-robin approach (random.choice). Consider incorporating algorithms like weighted round-robin,
# least connections, or least response time for better resource utilization and performance optimization.


79
80
81
82
83
84
85
86
87
88
89
@app.middleware("http")
async def check_url(request: Request, call_next):
    if len(app.state.MODELS) == 0:
        await get_all_models()
    else:
        pass

    response = await call_next(request)
    return response


Timothy J. Baek's avatar
Timothy J. Baek committed
90
91
92
93
94
95
@app.head("/")
@app.get("/")
async def get_status():
    return {"status": True}


Timothy J. Baek's avatar
Timothy J. Baek committed
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
@app.get("/config")
async def get_config(user=Depends(get_admin_user)):
    return {"ENABLE_OLLAMA_API": app.state.config.ENABLE_OLLAMA_API}


class OllamaConfigForm(BaseModel):
    enable_ollama_api: Optional[bool] = None


@app.post("/config/update")
async def update_config(form_data: OllamaConfigForm, user=Depends(get_admin_user)):
    app.state.config.ENABLE_OLLAMA_API = form_data.enable_ollama_api
    return {"ENABLE_OLLAMA_API": app.state.config.ENABLE_OLLAMA_API}


111
112
@app.get("/urls")
async def get_ollama_api_urls(user=Depends(get_admin_user)):
113
    return {"OLLAMA_BASE_URLS": app.state.config.OLLAMA_BASE_URLS}
114

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

116
class UrlUpdateForm(BaseModel):
117
    urls: List[str]
118
119


120
@app.post("/urls/update")
Timothy J. Baek's avatar
Timothy J. Baek committed
121
async def update_ollama_api_url(form_data: UrlUpdateForm, user=Depends(get_admin_user)):
122
    app.state.config.OLLAMA_BASE_URLS = form_data.urls
123

124
125
    log.info(f"app.state.config.OLLAMA_BASE_URLS: {app.state.config.OLLAMA_BASE_URLS}")
    return {"OLLAMA_BASE_URLS": app.state.config.OLLAMA_BASE_URLS}
Timothy J. Baek's avatar
Timothy J. Baek committed
126
127


128
async def fetch_url(url):
Timothy J. Baek's avatar
Timothy J. Baek committed
129
    timeout = aiohttp.ClientTimeout(total=5)
130
    try:
131
        async with aiohttp.ClientSession(timeout=timeout, trust_env=True) as session:
132
133
134
135
            async with session.get(url) as response:
                return await response.json()
    except Exception as e:
        # Handle connection error here
136
        log.error(f"Connection error: {e}")
137
138
139
        return None


140
141
142
143
144
145
146
147
148
149
async def cleanup_response(
    response: Optional[aiohttp.ClientResponse],
    session: Optional[aiohttp.ClientSession],
):
    if response:
        response.close()
    if session:
        await session.close()


150
async def post_streaming_url(url: str, payload: str, stream: bool = True):
151
152
    r = None
    try:
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
153
154
155
        session = aiohttp.ClientSession(
            trust_env=True, timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT)
        )
156
157
158
        r = await session.post(url, data=payload)
        r.raise_for_status()

159
160
161
162
163
164
165
166
167
168
169
170
171
172
        if stream:
            return StreamingResponse(
                r.content,
                status_code=r.status,
                headers=dict(r.headers),
                background=BackgroundTask(
                    cleanup_response, response=r, session=session
                ),
            )
        else:
            res = await r.json()
            await cleanup_response(r, session)
            return res

173
174
175
176
177
178
179
    except Exception as e:
        error_detail = "Open WebUI: Server Connection Error"
        if r is not None:
            try:
                res = await r.json()
                if "error" in res:
                    error_detail = f"Ollama: {res['error']}"
Michael Poluektov's avatar
cleanup  
Michael Poluektov committed
180
            except Exception:
181
182
183
184
185
186
187
188
                error_detail = f"Ollama: {e}"

        raise HTTPException(
            status_code=r.status if r else 500,
            detail=error_detail,
        )


189
190
191
192
def merge_models_lists(model_lists):
    merged_models = {}

    for idx, model_list in enumerate(model_lists):
Timothy J. Baek's avatar
Timothy J. Baek committed
193
194
195
196
197
198
199
200
        if model_list is not None:
            for model in model_list:
                digest = model["digest"]
                if digest not in merged_models:
                    model["urls"] = [idx]
                    merged_models[digest] = model
                else:
                    merged_models[digest]["urls"].append(idx)
201
202
203
204
205

    return list(merged_models.values())


async def get_all_models():
206
    log.info("get_all_models()")
Timothy J. Baek's avatar
Timothy J. Baek committed
207
208
209
210
211
212
213
214
215
216
217
218

    if app.state.config.ENABLE_OLLAMA_API:
        tasks = [
            fetch_url(f"{url}/api/tags") for url in app.state.config.OLLAMA_BASE_URLS
        ]
        responses = await asyncio.gather(*tasks)

        models = {
            "models": merge_models_lists(
                map(
                    lambda response: response["models"] if response else None, responses
                )
219
            )
Timothy J. Baek's avatar
Timothy J. Baek committed
220
221
222
223
        }

    else:
        models = {"models": []}
Timothy J. Baek's avatar
Timothy J. Baek committed
224

225
226
227
228
229
230
231
232
    app.state.MODELS = {model["model"]: model for model in models["models"]}

    return models


@app.get("/api/tags")
@app.get("/api/tags/{url_idx}")
async def get_ollama_tags(
233
    url_idx: Optional[int] = None, user=Depends(get_verified_user)
234
):
Michael Poluektov's avatar
cleanup  
Michael Poluektov committed
235
    if url_idx is None:
Timothy J. Baek's avatar
Timothy J. Baek committed
236
        models = await get_all_models()
237

Timothy J. Baek's avatar
Timothy J. Baek committed
238
        if app.state.config.ENABLE_MODEL_FILTER:
Timothy J. Baek's avatar
Timothy J. Baek committed
239
            if user.role == "user":
240
241
                models["models"] = list(
                    filter(
Timothy J. Baek's avatar
Timothy J. Baek committed
242
243
                        lambda model: model["name"]
                        in app.state.config.MODEL_FILTER_LIST,
244
245
                        models["models"],
                    )
Timothy J. Baek's avatar
Timothy J. Baek committed
246
247
248
                )
                return models
        return models
249
    else:
250
        url = app.state.config.OLLAMA_BASE_URLS[url_idx]
251
252

        r = None
253
254
255
256
257
258
        try:
            r = requests.request(method="GET", url=f"{url}/api/tags")
            r.raise_for_status()

            return r.json()
        except Exception as e:
259
            log.exception(e)
260
261
262
263
264
265
            error_detail = "Open WebUI: Server Connection Error"
            if r is not None:
                try:
                    res = r.json()
                    if "error" in res:
                        error_detail = f"Ollama: {res['error']}"
Michael Poluektov's avatar
cleanup  
Michael Poluektov committed
266
                except Exception:
267
268
269
270
271
272
273
274
275
276
277
                    error_detail = f"Ollama: {e}"

            raise HTTPException(
                status_code=r.status_code if r else 500,
                detail=error_detail,
            )


@app.get("/api/version")
@app.get("/api/version/{url_idx}")
async def get_ollama_versions(url_idx: Optional[int] = None):
278
    if app.state.config.ENABLE_OLLAMA_API:
Michael Poluektov's avatar
cleanup  
Michael Poluektov committed
279
        if url_idx is None:
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
            # returns lowest version
            tasks = [
                fetch_url(f"{url}/api/version")
                for url in app.state.config.OLLAMA_BASE_URLS
            ]
            responses = await asyncio.gather(*tasks)
            responses = list(filter(lambda x: x is not None, responses))

            if len(responses) > 0:
                lowest_version = min(
                    responses,
                    key=lambda x: tuple(
                        map(int, re.sub(r"^v|-.*", "", x["version"]).split("."))
                    ),
                )
295

296
297
298
299
300
301
                return {"version": lowest_version["version"]}
            else:
                raise HTTPException(
                    status_code=500,
                    detail=ERROR_MESSAGES.OLLAMA_NOT_FOUND,
                )
Timothy J. Baek's avatar
Timothy J. Baek committed
302
        else:
303
            url = app.state.config.OLLAMA_BASE_URLS[url_idx]
304

305
306
307
308
309
310
311
312
313
314
315
316
317
318
            r = None
            try:
                r = requests.request(method="GET", url=f"{url}/api/version")
                r.raise_for_status()

                return r.json()
            except Exception as e:
                log.exception(e)
                error_detail = "Open WebUI: Server Connection Error"
                if r is not None:
                    try:
                        res = r.json()
                        if "error" in res:
                            error_detail = f"Ollama: {res['error']}"
Michael Poluektov's avatar
cleanup  
Michael Poluektov committed
319
                    except Exception:
320
321
322
323
324
325
326
327
                        error_detail = f"Ollama: {e}"

                raise HTTPException(
                    status_code=r.status_code if r else 500,
                    detail=error_detail,
                )
    else:
        return {"version": False}
328
329
330
331
332
333
334
335
336
337
338


class ModelNameForm(BaseModel):
    name: str


@app.post("/api/pull")
@app.post("/api/pull/{url_idx}")
async def pull_model(
    form_data: ModelNameForm, url_idx: int = 0, user=Depends(get_admin_user)
):
339
    url = app.state.config.OLLAMA_BASE_URLS[url_idx]
340
    log.info(f"url: {url}")
341

Timothy J. Baek's avatar
Timothy J. Baek committed
342
343
344
    # Admin should be able to pull models from any source
    payload = {**form_data.model_dump(exclude_none=True), "insecure": True}

345
    return await post_streaming_url(f"{url}/api/pull", json.dumps(payload))
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360


class PushModelForm(BaseModel):
    name: str
    insecure: Optional[bool] = None
    stream: Optional[bool] = None


@app.delete("/api/push")
@app.delete("/api/push/{url_idx}")
async def push_model(
    form_data: PushModelForm,
    url_idx: Optional[int] = None,
    user=Depends(get_admin_user),
):
Michael Poluektov's avatar
cleanup  
Michael Poluektov committed
361
    if url_idx is None:
362
363
364
365
366
        if form_data.name in app.state.MODELS:
            url_idx = app.state.MODELS[form_data.name]["urls"][0]
        else:
            raise HTTPException(
                status_code=400,
Timothy J. Baek's avatar
Timothy J. Baek committed
367
                detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.name),
368
369
            )

370
    url = app.state.config.OLLAMA_BASE_URLS[url_idx]
371
    log.debug(f"url: {url}")
372

373
374
375
    return await post_streaming_url(
        f"{url}/api/push", form_data.model_dump_json(exclude_none=True).encode()
    )
376
377
378
379
380
381
382
383
384
385
386
387
388
389


class CreateModelForm(BaseModel):
    name: str
    modelfile: Optional[str] = None
    stream: Optional[bool] = None
    path: Optional[str] = None


@app.post("/api/create")
@app.post("/api/create/{url_idx}")
async def create_model(
    form_data: CreateModelForm, url_idx: int = 0, user=Depends(get_admin_user)
):
390
    log.debug(f"form_data: {form_data}")
391
    url = app.state.config.OLLAMA_BASE_URLS[url_idx]
392
    log.info(f"url: {url}")
393

394
395
396
    return await post_streaming_url(
        f"{url}/api/create", form_data.model_dump_json(exclude_none=True).encode()
    )
397
398
399
400
401
402
403
404
405
406
407
408
409
410


class CopyModelForm(BaseModel):
    source: str
    destination: str


@app.post("/api/copy")
@app.post("/api/copy/{url_idx}")
async def copy_model(
    form_data: CopyModelForm,
    url_idx: Optional[int] = None,
    user=Depends(get_admin_user),
):
Michael Poluektov's avatar
cleanup  
Michael Poluektov committed
411
    if url_idx is None:
412
413
414
415
416
        if form_data.source in app.state.MODELS:
            url_idx = app.state.MODELS[form_data.source]["urls"][0]
        else:
            raise HTTPException(
                status_code=400,
Timothy J. Baek's avatar
Timothy J. Baek committed
417
                detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.source),
418
419
            )

420
    url = app.state.config.OLLAMA_BASE_URLS[url_idx]
421
    log.info(f"url: {url}")
Michael Poluektov's avatar
cleanup  
Michael Poluektov committed
422
423
424
425
426
    r = requests.request(
        method="POST",
        url=f"{url}/api/copy",
        data=form_data.model_dump_json(exclude_none=True).encode(),
    )
427
428
429
430

    try:
        r.raise_for_status()

431
        log.debug(f"r.text: {r.text}")
432
433
434

        return True
    except Exception as e:
435
        log.exception(e)
436
437
438
439
440
441
        error_detail = "Open WebUI: Server Connection Error"
        if r is not None:
            try:
                res = r.json()
                if "error" in res:
                    error_detail = f"Ollama: {res['error']}"
Michael Poluektov's avatar
cleanup  
Michael Poluektov committed
442
            except Exception:
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
                error_detail = f"Ollama: {e}"

        raise HTTPException(
            status_code=r.status_code if r else 500,
            detail=error_detail,
        )


@app.delete("/api/delete")
@app.delete("/api/delete/{url_idx}")
async def delete_model(
    form_data: ModelNameForm,
    url_idx: Optional[int] = None,
    user=Depends(get_admin_user),
):
Michael Poluektov's avatar
cleanup  
Michael Poluektov committed
458
    if url_idx is None:
459
460
461
462
463
        if form_data.name in app.state.MODELS:
            url_idx = app.state.MODELS[form_data.name]["urls"][0]
        else:
            raise HTTPException(
                status_code=400,
Timothy J. Baek's avatar
Timothy J. Baek committed
464
                detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.name),
465
466
            )

467
    url = app.state.config.OLLAMA_BASE_URLS[url_idx]
468
    log.info(f"url: {url}")
469

Michael Poluektov's avatar
cleanup  
Michael Poluektov committed
470
471
472
473
474
    r = requests.request(
        method="DELETE",
        url=f"{url}/api/delete",
        data=form_data.model_dump_json(exclude_none=True).encode(),
    )
475
476
477
    try:
        r.raise_for_status()

478
        log.debug(f"r.text: {r.text}")
479
480
481

        return True
    except Exception as e:
482
        log.exception(e)
483
484
485
486
487
488
        error_detail = "Open WebUI: Server Connection Error"
        if r is not None:
            try:
                res = r.json()
                if "error" in res:
                    error_detail = f"Ollama: {res['error']}"
Michael Poluektov's avatar
cleanup  
Michael Poluektov committed
489
            except Exception:
490
491
492
493
494
495
496
497
498
                error_detail = f"Ollama: {e}"

        raise HTTPException(
            status_code=r.status_code if r else 500,
            detail=error_detail,
        )


@app.post("/api/show")
499
async def show_model_info(form_data: ModelNameForm, user=Depends(get_verified_user)):
500
501
502
    if form_data.name not in app.state.MODELS:
        raise HTTPException(
            status_code=400,
Timothy J. Baek's avatar
Timothy J. Baek committed
503
            detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.name),
504
505
506
        )

    url_idx = random.choice(app.state.MODELS[form_data.name]["urls"])
507
    url = app.state.config.OLLAMA_BASE_URLS[url_idx]
508
    log.info(f"url: {url}")
509

Michael Poluektov's avatar
cleanup  
Michael Poluektov committed
510
511
512
513
514
    r = requests.request(
        method="POST",
        url=f"{url}/api/show",
        data=form_data.model_dump_json(exclude_none=True).encode(),
    )
515
516
517
518
519
    try:
        r.raise_for_status()

        return r.json()
    except Exception as e:
520
        log.exception(e)
521
522
523
524
525
526
        error_detail = "Open WebUI: Server Connection Error"
        if r is not None:
            try:
                res = r.json()
                if "error" in res:
                    error_detail = f"Ollama: {res['error']}"
Michael Poluektov's avatar
cleanup  
Michael Poluektov committed
527
            except Exception:
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
                error_detail = f"Ollama: {e}"

        raise HTTPException(
            status_code=r.status_code if r else 500,
            detail=error_detail,
        )


class GenerateEmbeddingsForm(BaseModel):
    model: str
    prompt: str
    options: Optional[dict] = None
    keep_alive: Optional[Union[int, str]] = None


@app.post("/api/embeddings")
@app.post("/api/embeddings/{url_idx}")
async def generate_embeddings(
    form_data: GenerateEmbeddingsForm,
    url_idx: Optional[int] = None,
548
    user=Depends(get_verified_user),
549
):
Michael Poluektov's avatar
cleanup  
Michael Poluektov committed
550
    if url_idx is None:
Timothy J. Baek's avatar
Timothy J. Baek committed
551
552
553
554
555
556
        model = form_data.model

        if ":" not in model:
            model = f"{model}:latest"

        if model in app.state.MODELS:
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
557
            url_idx = random.choice(app.state.MODELS[model]["urls"])
558
559
560
        else:
            raise HTTPException(
                status_code=400,
Timothy J. Baek's avatar
Timothy J. Baek committed
561
                detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.model),
562
563
            )

564
    url = app.state.config.OLLAMA_BASE_URLS[url_idx]
565
    log.info(f"url: {url}")
566

Michael Poluektov's avatar
cleanup  
Michael Poluektov committed
567
568
569
570
571
    r = requests.request(
        method="POST",
        url=f"{url}/api/embeddings",
        data=form_data.model_dump_json(exclude_none=True).encode(),
    )
572
573
574
575
576
    try:
        r.raise_for_status()

        return r.json()
    except Exception as e:
577
        log.exception(e)
578
579
580
581
582
583
        error_detail = "Open WebUI: Server Connection Error"
        if r is not None:
            try:
                res = r.json()
                if "error" in res:
                    error_detail = f"Ollama: {res['error']}"
Michael Poluektov's avatar
cleanup  
Michael Poluektov committed
584
            except Exception:
585
586
587
588
589
590
591
592
                error_detail = f"Ollama: {e}"

        raise HTTPException(
            status_code=r.status_code if r else 500,
            detail=error_detail,
        )


593
594
595
596
def generate_ollama_embeddings(
    form_data: GenerateEmbeddingsForm,
    url_idx: Optional[int] = None,
):
597
    log.info(f"generate_ollama_embeddings {form_data}")
Timothy J. Baek's avatar
Timothy J. Baek committed
598

Michael Poluektov's avatar
cleanup  
Michael Poluektov committed
599
    if url_idx is None:
600
601
602
603
604
605
606
607
608
609
610
611
612
        model = form_data.model

        if ":" not in model:
            model = f"{model}:latest"

        if model in app.state.MODELS:
            url_idx = random.choice(app.state.MODELS[model]["urls"])
        else:
            raise HTTPException(
                status_code=400,
                detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.model),
            )

613
    url = app.state.config.OLLAMA_BASE_URLS[url_idx]
614
615
    log.info(f"url: {url}")

Michael Poluektov's avatar
cleanup  
Michael Poluektov committed
616
617
618
619
620
    r = requests.request(
        method="POST",
        url=f"{url}/api/embeddings",
        data=form_data.model_dump_json(exclude_none=True).encode(),
    )
621
622
623
624
625
    try:
        r.raise_for_status()

        data = r.json()

626
        log.info(f"generate_ollama_embeddings {data}")
Timothy J. Baek's avatar
Timothy J. Baek committed
627

628
629
630
        if "embedding" in data:
            return data["embedding"]
        else:
Michael Poluektov's avatar
cleanup  
Michael Poluektov committed
631
            raise Exception("Something went wrong :/")
632
633
634
635
636
637
638
639
    except Exception as e:
        log.exception(e)
        error_detail = "Open WebUI: Server Connection Error"
        if r is not None:
            try:
                res = r.json()
                if "error" in res:
                    error_detail = f"Ollama: {res['error']}"
Michael Poluektov's avatar
cleanup  
Michael Poluektov committed
640
            except Exception:
641
642
                error_detail = f"Ollama: {e}"

Michael Poluektov's avatar
cleanup  
Michael Poluektov committed
643
        raise Exception(error_detail)
644
645


646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
class GenerateCompletionForm(BaseModel):
    model: str
    prompt: str
    images: Optional[List[str]] = None
    format: Optional[str] = None
    options: Optional[dict] = None
    system: Optional[str] = None
    template: Optional[str] = None
    context: Optional[str] = None
    stream: Optional[bool] = True
    raw: Optional[bool] = None
    keep_alive: Optional[Union[int, str]] = None


@app.post("/api/generate")
@app.post("/api/generate/{url_idx}")
async def generate_completion(
    form_data: GenerateCompletionForm,
    url_idx: Optional[int] = None,
665
    user=Depends(get_verified_user),
666
):
Michael Poluektov's avatar
cleanup  
Michael Poluektov committed
667
    if url_idx is None:
Timothy J. Baek's avatar
Timothy J. Baek committed
668
669
670
671
672
673
        model = form_data.model

        if ":" not in model:
            model = f"{model}:latest"

        if model in app.state.MODELS:
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
674
            url_idx = random.choice(app.state.MODELS[model]["urls"])
675
676
677
        else:
            raise HTTPException(
                status_code=400,
Timothy J. Baek's avatar
Timothy J. Baek committed
678
                detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.model),
679
680
            )

681
    url = app.state.config.OLLAMA_BASE_URLS[url_idx]
682
    log.info(f"url: {url}")
683

684
685
686
    return await post_streaming_url(
        f"{url}/api/generate", form_data.model_dump_json(exclude_none=True).encode()
    )
687
688
689
690
691
692
693
694
695
696
697
698
699
700


class ChatMessage(BaseModel):
    role: str
    content: str
    images: Optional[List[str]] = None


class GenerateChatCompletionForm(BaseModel):
    model: str
    messages: List[ChatMessage]
    format: Optional[str] = None
    options: Optional[dict] = None
    template: Optional[str] = None
Timothy J. Baek's avatar
Timothy J. Baek committed
701
    stream: Optional[bool] = None
702
703
704
    keep_alive: Optional[Union[int, str]] = None


Michael Poluektov's avatar
Michael Poluektov committed
705
706
707
708
709
710
711
712
713
714
715
716
def get_ollama_url(url_idx: Optional[int], model: str):
    if url_idx is None:
        if model not in app.state.MODELS:
            raise HTTPException(
                status_code=400,
                detail=ERROR_MESSAGES.MODEL_NOT_FOUND(model),
            )
        url_idx = random.choice(app.state.MODELS[model]["urls"])
    url = app.state.config.OLLAMA_BASE_URLS[url_idx]
    return url


717
718
@app.post("/api/chat")
@app.post("/api/chat/{url_idx}")
Timothy J. Baek's avatar
Timothy J. Baek committed
719
async def generate_chat_completion(
720
721
    form_data: GenerateChatCompletionForm,
    url_idx: Optional[int] = None,
722
    user=Depends(get_verified_user),
723
):
Michael Poluektov's avatar
Michael Poluektov committed
724
    log.debug(f"{form_data.model_dump_json(exclude_none=True).encode()}=")
Timothy J. Baek's avatar
Timothy J. Baek committed
725
726

    payload = {
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
727
        **form_data.model_dump(exclude_none=True, exclude=["metadata"]),
Timothy J. Baek's avatar
Timothy J. Baek committed
728
    }
Michael Poluektov's avatar
Michael Poluektov committed
729
    payload.pop("metadata", None)
Timothy J. Baek's avatar
Timothy J. Baek committed
730
731
732
733
734
735
736
737

    model_id = form_data.model
    model_info = Models.get_model_by_id(model_id)

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

Michael Poluektov's avatar
cleanup  
Michael Poluektov committed
738
        params = model_info.params.model_dump()
Timothy J. Baek's avatar
Timothy J. Baek committed
739

Michael Poluektov's avatar
cleanup  
Michael Poluektov committed
740
        if params:
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
741
742
            if payload.get("options") is None:
                payload["options"] = {}
Timothy J. Baek's avatar
Timothy J. Baek committed
743

Michael Poluektov's avatar
Michael Poluektov committed
744
745
            payload["options"] = apply_model_params_to_body_ollama(
                params, payload["options"]
746
            )
Michael Poluektov's avatar
Michael Poluektov committed
747
            payload = apply_model_system_prompt_to_body(params, payload, user)
748

Michael Poluektov's avatar
Michael Poluektov committed
749
750
    if ":" not in payload["model"]:
        payload["model"] = f"{payload['model']}:latest"
751

Michael Poluektov's avatar
Michael Poluektov committed
752
    url = get_ollama_url(url_idx, payload["model"])
753
    log.info(f"url: {url}")
754
    log.debug(payload)
Timothy J. Baek's avatar
Timothy J. Baek committed
755

756
    return await post_streaming_url(f"{url}/api/chat", json.dumps(payload))
Timothy J. Baek's avatar
Timothy J. Baek committed
757
758
759


# TODO: we should update this part once Ollama supports other types
760
761
762
763
764
class OpenAIChatMessageContent(BaseModel):
    type: str
    model_config = ConfigDict(extra="allow")


Timothy J. Baek's avatar
Timothy J. Baek committed
765
766
class OpenAIChatMessage(BaseModel):
    role: str
767
    content: Union[str, OpenAIChatMessageContent]
Timothy J. Baek's avatar
Timothy J. Baek committed
768
769
770
771
772
773
774
775
776
777
778
779
780
781

    model_config = ConfigDict(extra="allow")


class OpenAIChatCompletionForm(BaseModel):
    model: str
    messages: List[OpenAIChatMessage]

    model_config = ConfigDict(extra="allow")


@app.post("/v1/chat/completions")
@app.post("/v1/chat/completions/{url_idx}")
async def generate_openai_chat_completion(
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
782
    form_data: dict,
Timothy J. Baek's avatar
Timothy J. Baek committed
783
    url_idx: Optional[int] = None,
784
    user=Depends(get_verified_user),
Timothy J. Baek's avatar
Timothy J. Baek committed
785
):
Michael Poluektov's avatar
Michael Poluektov committed
786
787
    completion_form = OpenAIChatCompletionForm(**form_data)
    payload = {**completion_form.model_dump(exclude_none=True, exclude=["metadata"])}
Michael Poluektov's avatar
Michael Poluektov committed
788
    payload.pop("metadata", None)
Timothy J. Baek's avatar
Timothy J. Baek committed
789

Michael Poluektov's avatar
Michael Poluektov committed
790
    model_id = completion_form.model
Timothy J. Baek's avatar
Timothy J. Baek committed
791
    model_info = Models.get_model_by_id(model_id)
Timothy J. Baek's avatar
Timothy J. Baek committed
792

Timothy J. Baek's avatar
Timothy J. Baek committed
793
794
795
796
    if model_info:
        if model_info.base_model_id:
            payload["model"] = model_info.base_model_id

Michael Poluektov's avatar
Michael Poluektov committed
797
        params = model_info.params.model_dump()
Timothy J. Baek's avatar
Timothy J. Baek committed
798

Michael Poluektov's avatar
Michael Poluektov committed
799
800
801
        if params:
            payload = apply_model_params_to_body_openai(params, payload)
            payload = apply_model_system_prompt_to_body(params, payload, user)
Timothy J. Baek's avatar
Timothy J. Baek committed
802

Michael Poluektov's avatar
Michael Poluektov committed
803
804
    if ":" not in payload["model"]:
        payload["model"] = f"{payload['model']}:latest"
Timothy J. Baek's avatar
Timothy J. Baek committed
805

Michael Poluektov's avatar
Michael Poluektov committed
806
    url = get_ollama_url(url_idx, payload["model"])
807
    log.info(f"url: {url}")
Timothy J. Baek's avatar
Timothy J. Baek committed
808

809
810
811
812
813
    return await post_streaming_url(
        f"{url}/v1/chat/completions",
        json.dumps(payload),
        stream=payload.get("stream", False),
    )
814
815


816
817
818
819
820
821
@app.get("/v1/models")
@app.get("/v1/models/{url_idx}")
async def get_openai_models(
    url_idx: Optional[int] = None,
    user=Depends(get_verified_user),
):
Michael Poluektov's avatar
cleanup  
Michael Poluektov committed
822
    if url_idx is None:
823
824
        models = await get_all_models()

Timothy J. Baek's avatar
Timothy J. Baek committed
825
        if app.state.config.ENABLE_MODEL_FILTER:
826
827
828
            if user.role == "user":
                models["models"] = list(
                    filter(
Timothy J. Baek's avatar
Timothy J. Baek committed
829
830
                        lambda model: model["name"]
                        in app.state.config.MODEL_FILTER_LIST,
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
                        models["models"],
                    )
                )

        return {
            "data": [
                {
                    "id": model["model"],
                    "object": "model",
                    "created": int(time.time()),
                    "owned_by": "openai",
                }
                for model in models["models"]
            ],
            "object": "list",
        }

    else:
849
        url = app.state.config.OLLAMA_BASE_URLS[url_idx]
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
        try:
            r = requests.request(method="GET", url=f"{url}/api/tags")
            r.raise_for_status()

            models = r.json()

            return {
                "data": [
                    {
                        "id": model["model"],
                        "object": "model",
                        "created": int(time.time()),
                        "owned_by": "openai",
                    }
                    for model in models["models"]
                ],
                "object": "list",
            }

        except Exception as e:
            log.exception(e)
            error_detail = "Open WebUI: Server Connection Error"
            if r is not None:
                try:
                    res = r.json()
                    if "error" in res:
                        error_detail = f"Ollama: {res['error']}"
Michael Poluektov's avatar
cleanup  
Michael Poluektov committed
877
                except Exception:
878
879
880
881
882
883
884
885
                    error_detail = f"Ollama: {e}"

            raise HTTPException(
                status_code=r.status_code if r else 500,
                detail=error_detail,
            )


Timothy J. Baek's avatar
Timothy J. Baek committed
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
class UrlForm(BaseModel):
    url: str


class UploadBlobForm(BaseModel):
    filename: str


def parse_huggingface_url(hf_url):
    try:
        # Parse the URL
        parsed_url = urlparse(hf_url)

        # Get the path and split it into components
        path_components = parsed_url.path.split("/")

        # Extract the desired output
        model_file = path_components[-1]

        return model_file
    except ValueError:
        return None


async def download_file_stream(
    ollama_url, file_url, file_path, file_name, chunk_size=1024 * 1024
):
    done = False

    if os.path.exists(file_path):
        current_size = os.path.getsize(file_path)
    else:
        current_size = 0

    headers = {"Range": f"bytes={current_size}-"} if current_size > 0 else {}

    timeout = aiohttp.ClientTimeout(total=600)  # Set the timeout

924
    async with aiohttp.ClientSession(timeout=timeout, trust_env=True) as session:
Timothy J. Baek's avatar
Timothy J. Baek committed
925
926
927
928
929
930
931
932
933
934
        async with session.get(file_url, headers=headers) as response:
            total_size = int(response.headers.get("content-length", 0)) + current_size

            with open(file_path, "ab+") as file:
                async for data in response.content.iter_chunked(chunk_size):
                    current_size += len(data)
                    file.write(data)

                    done = current_size == total_size
                    progress = round((current_size / total_size) * 100, 2)
Timothy J. Baek's avatar
Timothy J. Baek committed
935

Timothy J. Baek's avatar
Timothy J. Baek committed
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
                    yield f'data: {{"progress": {progress}, "completed": {current_size}, "total": {total_size}}}\n\n'

                if done:
                    file.seek(0)
                    hashed = calculate_sha256(file)
                    file.seek(0)

                    url = f"{ollama_url}/api/blobs/sha256:{hashed}"
                    response = requests.post(url, data=file)

                    if response.ok:
                        res = {
                            "done": done,
                            "blob": f"sha256:{hashed}",
                            "name": file_name,
                        }
                        os.remove(file_path)

                        yield f"data: {json.dumps(res)}\n\n"
                    else:
                        raise "Ollama: Could not create blob, Please try again."


# url = "https://huggingface.co/TheBloke/stablelm-zephyr-3b-GGUF/resolve/main/stablelm-zephyr-3b.Q2_K.gguf"
@app.post("/models/download")
@app.post("/models/download/{url_idx}")
async def download_model(
    form_data: UrlForm,
    url_idx: Optional[int] = None,
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
965
    user=Depends(get_admin_user),
Timothy J. Baek's avatar
Timothy J. Baek committed
966
):
Timothy J. Baek's avatar
Timothy J. Baek committed
967
968
969
970
971
972
973
974
    allowed_hosts = ["https://huggingface.co/", "https://github.com/"]

    if not any(form_data.url.startswith(host) for host in allowed_hosts):
        raise HTTPException(
            status_code=400,
            detail="Invalid file_url. Only URLs from allowed hosts are permitted.",
        )

Michael Poluektov's avatar
cleanup  
Michael Poluektov committed
975
    if url_idx is None:
Timothy J. Baek's avatar
Timothy J. Baek committed
976
        url_idx = 0
977
    url = app.state.config.OLLAMA_BASE_URLS[url_idx]
Timothy J. Baek's avatar
Timothy J. Baek committed
978
979
980
981
982

    file_name = parse_huggingface_url(form_data.url)

    if file_name:
        file_path = f"{UPLOAD_DIR}/{file_name}"
Timothy J. Baek's avatar
Timothy J. Baek committed
983

Timothy J. Baek's avatar
Timothy J. Baek committed
984
        return StreamingResponse(
Timothy J. Baek's avatar
Timothy J. Baek committed
985
            download_file_stream(url, form_data.url, file_path, file_name),
Timothy J. Baek's avatar
Timothy J. Baek committed
986
987
988
989
990
991
992
        )
    else:
        return None


@app.post("/models/upload")
@app.post("/models/upload/{url_idx}")
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
993
994
995
996
997
def upload_model(
    file: UploadFile = File(...),
    url_idx: Optional[int] = None,
    user=Depends(get_admin_user),
):
Michael Poluektov's avatar
cleanup  
Michael Poluektov committed
998
    if url_idx is None:
Timothy J. Baek's avatar
Timothy J. Baek committed
999
        url_idx = 0
1000
    ollama_url = app.state.config.OLLAMA_BASE_URLS[url_idx]
Timothy J. Baek's avatar
Timothy J. Baek committed
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038

    file_path = f"{UPLOAD_DIR}/{file.filename}"

    # Save file in chunks
    with open(file_path, "wb+") as f:
        for chunk in file.file:
            f.write(chunk)

    def file_process_stream():
        nonlocal ollama_url
        total_size = os.path.getsize(file_path)
        chunk_size = 1024 * 1024
        try:
            with open(file_path, "rb") as f:
                total = 0
                done = False

                while not done:
                    chunk = f.read(chunk_size)
                    if not chunk:
                        done = True
                        continue

                    total += len(chunk)
                    progress = round((total / total_size) * 100, 2)

                    res = {
                        "progress": progress,
                        "total": total_size,
                        "completed": total,
                    }
                    yield f"data: {json.dumps(res)}\n\n"

                if done:
                    f.seek(0)
                    hashed = calculate_sha256(f)
                    f.seek(0)

Timothy J. Baek's avatar
Timothy J. Baek committed
1039
                    url = f"{ollama_url}/api/blobs/sha256:{hashed}"
Timothy J. Baek's avatar
Timothy J. Baek committed
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
                    response = requests.post(url, data=f)

                    if response.ok:
                        res = {
                            "done": done,
                            "blob": f"sha256:{hashed}",
                            "name": file.filename,
                        }
                        os.remove(file_path)
                        yield f"data: {json.dumps(res)}\n\n"
                    else:
                        raise Exception(
                            "Ollama: Could not create blob, Please try again."
                        )

        except Exception as e:
            res = {"error": str(e)}
            yield f"data: {json.dumps(res)}\n\n"

    return StreamingResponse(file_process_stream(), media_type="text/event-stream")