Merge branch 'dev' into feat/model-config

backend/apps/ollama/main.py

@@ -43,6 +43,7 @@ from utils.utils import (
 from config import (
     SRC_LOG_LEVELS,
     OLLAMA_BASE_URLS,
+    ENABLE_OLLAMA_API,
     ENABLE_MODEL_FILTER,
     MODEL_FILTER_LIST,
     UPLOAD_DIR,
@@ -68,6 +69,8 @@ app.state.config.ENABLE_MODEL_FILTER = ENABLE_MODEL_FILTER
 app.state.config.MODEL_FILTER_LIST = MODEL_FILTER_LIST
 app.state.MODEL_CONFIG = Models.get_all_models()
 
+
+app.state.config.ENABLE_OLLAMA_API = ENABLE_OLLAMA_API
 app.state.config.OLLAMA_BASE_URLS = OLLAMA_BASE_URLS
 app.state.MODELS = {}
 
@@ -97,6 +100,21 @@ async def get_status():
     return {"status": True}
 
 
+@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}
+
+
 @app.get("/urls")
 async def get_ollama_api_urls(user=Depends(get_admin_user)):
     return {"OLLAMA_BASE_URLS": app.state.config.OLLAMA_BASE_URLS}
@@ -157,17 +175,24 @@ def merge_models_lists(model_lists):
 
 async def get_all_models():
     log.info("get_all_models()")
-    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,
+    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
+                )
             )
-        )
-    }
+        }
+
+    else:
+        models = {"models": []}
+        
     for model in models["models"]:
         add_custom_info_to_model(model)
 

backend/apps/openai/main.py

@@ -316,6 +316,7 @@ async def get_models(url_idx: Optional[int] = None, user=Depends(get_current_use
 @app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
 async def proxy(path: str, request: Request, user=Depends(get_verified_user)):
     idx = 0
+    pipeline = False
 
     body = await request.body()
     # TODO: Remove below after gpt-4-vision fix from Open AI
@@ -324,7 +325,15 @@ async def proxy(path: str, request: Request, user=Depends(get_verified_user)):
         body = body.decode("utf-8")
         body = json.loads(body)
 
-        idx = app.state.MODELS[body.get("model")]["urlIdx"]
+        model = app.state.MODELS[body.get("model")]
+
+        idx = model["urlIdx"]
+
+        if "pipeline" in model:
+            pipeline = model.get("pipeline")
+
+        if pipeline:
+            body["user"] = {"name": user.name, "id": user.id}
 
         # Check if the model is "gpt-4-vision-preview" and set "max_tokens" to 4000
         # This is a workaround until OpenAI fixes the issue with this model

backend/apps/web/internal/db.py

@@ -3,7 +3,7 @@ import json
 from peewee import *
 from peewee_migrate import Router
 from playhouse.db_url import connect
-from config import SRC_LOG_LEVELS, DATA_DIR, DATABASE_URL
+from config import SRC_LOG_LEVELS, DATA_DIR, DATABASE_URL, BACKEND_DIR
 import os
 import logging
 
@@ -30,6 +30,8 @@ else:
 
 DB = connect(DATABASE_URL)
 log.info(f"Connected to a {DB.__class__.__name__} database.")
-router = Router(DB, migrate_dir="apps/web/internal/migrations", logger=log)
+router = Router(
+    DB, migrate_dir=BACKEND_DIR / "apps" / "web" / "internal" / "migrations", logger=log
+)
 router.run()
 DB.connect(reuse_if_open=True)