feat: add update api for embedding/reranker models (pingcap#677)

part of pingcap#667

Author: Mini256

backend/Dockerfile

@@ -27,6 +27,6 @@ ENV PYTHONPATH=/app
 COPY . /app/
 
 # Default number of workers
-ENV FASTAPI_WORKERS=1
+ENV WEB_CONCURRENCY=4
 
-CMD ["sh", "-c", "fastapi run app/api_server.py --host 0.0.0.0 --port 80 --workers ${FASTAPI_WORKERS}"]
+CMD ["sh", "-c", "fastapi run app/api_server.py --host 0.0.0.0 --port 80 --workers ${WEB_CONCURRENCY}"]

backend/Makefile

@@ -25,8 +25,12 @@ test:
 	@uv run pytest -v tests/
 
 dev_backend:
-	@echo "Running development backend server..."
-	@uv run fastapi dev app/api_server.py --host 0.0.0.0 --port 5001
+	@echo "Running backend server in development mode..."
+	@uv run fastapi dev app/api_server.py --host 127.0.0.1 --port 5001
+
+run_backend:
+	@echo "Running backend server..."
+	@uv run fastapi run app/api_server.py --host 0.0.0.0 --port 5001 --workers 4
 
 dev_celery_flower:
 	@echo "Running Celery Flower..."

backend/app/api/admin_routes/embedding_model/models.py

@@ -27,8 +27,7 @@ def vector_dimension_must_gt_1(cls, v: int) -> int:
 class EmbeddingModelUpdate(BaseModel):
     name: Optional[str] = None
     config: Optional[dict | list] = None
-    credentials: Optional[Any] = None
-    is_default: Optional[bool] = False
+    credentials: Optional[str | dict] = None
 
 
 class EmbeddingModelItem(BaseModel):

backend/app/api/admin_routes/embedding_model/routes.py

@@ -1,4 +1,3 @@
-import logging
 from typing import List
 
 from fastapi import APIRouter, Depends
@@ -12,43 +11,29 @@
     EmbeddingModelCreate,
 )
 from app.api.deps import CurrentSuperuserDep, SessionDep
-from app.exceptions import EmbeddingModelNotFound, InternalServerError
-from app.repositories.embedding_model import embed_model_repo
+from app.repositories.embedding_model import embedding_model_repo
 from app.rag.embeddings.provider import (
     EmbeddingProviderOption,
     embedding_provider_options,
 )
 from app.rag.embeddings.resolver import resolve_embed_model
+from app.logger import logger
 
 router = APIRouter()
-logger = logging.getLogger(__name__)
 
 
-@router.get("/admin/embedding-models/provider/options")
+@router.get("/admin/embedding-models/providers/options")
 def list_embedding_model_provider_options(
     user: CurrentSuperuserDep,
 ) -> List[EmbeddingProviderOption]:
     return embedding_provider_options
 
 
-@router.get("/admin/embedding-models/options", deprecated=True)
-def get_embedding_model_options(
-    user: CurrentSuperuserDep,
-) -> List[EmbeddingProviderOption]:
-    return embedding_provider_options
-
-
-@router.post("/admin/embedding-models")
-def create_embedding_model(
-    session: SessionDep,
-    user: CurrentSuperuserDep,
-    create: EmbeddingModelCreate,
-) -> EmbeddingModelDetail:
-    try:
-        return embed_model_repo.create(session, create)
-    except Exception as e:
-        logger.exception(e)
-        raise InternalServerError()
+@router.get("/admin/embedding-models")
+def list_embedding_models(
+    db_session: SessionDep, user: CurrentSuperuserDep, params: Params = Depends()
+) -> Page[EmbeddingModelItem]:
+    return embedding_model_repo.paginate(db_session, params)
 
 
 @router.post("/admin/embedding-models/test")
@@ -72,60 +57,50 @@ def test_embedding_model(
         success = True
         error = ""
     except Exception as e:
+        logger.info(f"Failed to test embedding model: {e}")
         success = False
         error = str(e)
     return EmbeddingModelTestResult(success=success, error=error)
 
 
-@router.get("/admin/embedding-models")
-def list_embedding_models(
-    session: SessionDep, user: CurrentSuperuserDep, params: Params = Depends()
-) -> Page[EmbeddingModelItem]:
-    return embed_model_repo.paginate(session, params)
+@router.post("/admin/embedding-models")
+def create_embedding_model(
+    db_session: SessionDep,
+    user: CurrentSuperuserDep,
+    create: EmbeddingModelCreate,
+) -> EmbeddingModelDetail:
+    return embedding_model_repo.create(db_session, create)
 
 
 @router.get("/admin/embedding-models/{model_id}")
 def get_embedding_model_detail(
-    session: SessionDep, user: CurrentSuperuserDep, model_id: int
+    db_session: SessionDep, user: CurrentSuperuserDep, model_id: int
 ) -> EmbeddingModelDetail:
-    try:
-        return embed_model_repo.must_get(session, model_id)
-    except EmbeddingModelNotFound as e:
-        raise e
-    except Exception as e:
-        logger.exception(e)
-        raise InternalServerError()
+    return embedding_model_repo.must_get(db_session, model_id)
 
 
 @router.put("/admin/embedding-models/{model_id}")
 def update_embedding_model(
-    session: SessionDep,
+    db_session: SessionDep,
     user: CurrentSuperuserDep,
     model_id: int,
     update: EmbeddingModelUpdate,
 ) -> EmbeddingModelDetail:
-    try:
-        embed_model = embed_model_repo.must_get(session, model_id)
-        embed_model_repo.update(session, embed_model, update)
-        return embed_model
-    except EmbeddingModelNotFound as e:
-        raise e
-    except Exception as e:
-        logger.exception(e)
-        raise InternalServerError()
+    embed_model = embedding_model_repo.must_get(db_session, model_id)
+    return embedding_model_repo.update(db_session, embed_model, update)
+
+
+@router.delete("/admin/embedding-models/{model_id}")
+def delete_embedding_model(
+    db_session: SessionDep, user: CurrentSuperuserDep, model_id: int
+) -> None:
+    embedding_model = embedding_model_repo.must_get(db_session, model_id)
+    embedding_model_repo.delete(db_session, embedding_model)
 
 
 @router.put("/admin/embedding-models/{model_id}/set_default")
 def set_default_embedding_model(
-    session: SessionDep, user: CurrentSuperuserDep, model_id: int
+    db_session: SessionDep, user: CurrentSuperuserDep, model_id: int
 ) -> EmbeddingModelDetail:
-    try:
-        embed_model = embed_model_repo.must_get(session, model_id)
-        embed_model_repo.set_default_model(session, model_id)
-        session.refresh(embed_model)
-        return embed_model
-    except EmbeddingModelNotFound as e:
-        raise e
-    except Exception as e:
-        logger.exception(e)
-        raise InternalServerError()
+    embed_model = embedding_model_repo.must_get(db_session, model_id)
+    return embedding_model_repo.set_default(db_session, embed_model)

backend/app/api/admin_routes/knowledge_base/routes.py

@@ -29,7 +29,7 @@
     KnowledgeBase,
 )
 from app.repositories import (
-    embed_model_repo,
+    embedding_model_repo,
     llm_repo,
     data_source_repo,
     knowledge_base_repo,
@@ -77,7 +77,9 @@ def create_knowledge_base(
             create.llm_id = llm_repo.must_get_default(session).id
 
         if not create.embedding_model_id:
-            create.embedding_model_id = embed_model_repo.must_get_default(session).id
+            create.embedding_model_id = embedding_model_repo.must_get_default(
+                session
+            ).id
 
         knowledge_base = KnowledgeBase(
             name=create.name,

backend/app/api/admin_routes/llm/routes.py

@@ -4,11 +4,10 @@
 from fastapi_pagination import Page, Params
 from llama_index.core.base.llms.types import ChatMessage
 from pydantic import BaseModel
-from sqlalchemy import update
 
 from app.api.deps import CurrentSuperuserDep, SessionDep
 from app.logger import logger
-from app.models import AdminLLM, ChatEngine, KnowledgeBase, LLM, LLMUpdate
+from app.models import AdminLLM, LLM, LLMUpdate
 from app.rag.llms.provider import LLMProviderOption, llm_provider_options
 from app.rag.llms.resolver import resolve_llm
 from app.repositories.llm import llm_repo
@@ -17,7 +16,7 @@
 router = APIRouter()
 
 
-@router.get("/admin/llms/provider/options")
+@router.get("/admin/llms/providers/options")
 def list_llm_provider_options(user: CurrentSuperuserDep) -> List[LLMProviderOption]:
     return llm_provider_options
 
@@ -31,15 +30,6 @@ def list_llms(
     return llm_repo.paginate(db_session, params)
 
 
-@router.post("/admin/llms")
-def create_llm(
-    db_session: SessionDep,
-    user: CurrentSuperuserDep,
-    llm: LLM,
-) -> AdminLLM:
-    return llm_repo.create(db_session, llm)
-
-
 class LLMTestResult(BaseModel):
     success: bool
     error: str = ""
@@ -72,12 +62,21 @@ def test_llm(
         success = True
         error = ""
     except Exception as e:
-        logger.error(f"Failed to test LLM: {e}")
+        logger.info(f"Failed to test LLM: {e}")
         success = False
         error = str(e)
     return LLMTestResult(success=success, error=error)
 
 
+@router.post("/admin/llms")
+def create_llm(
+    db_session: SessionDep,
+    user: CurrentSuperuserDep,
+    llm: LLM,
+) -> AdminLLM:
+    return llm_repo.create(db_session, llm)
+
+
 @router.get("/admin/llms/{llm_id}")
 def get_llm(
     db_session: SessionDep,
@@ -103,26 +102,9 @@ def delete_llm(
     db_session: SessionDep,
     user: CurrentSuperuserDep,
     llm_id: int,
-) -> AdminLLM:
+) -> None:
     llm = llm_repo.must_get(db_session, llm_id)
-
-    # TODO: Support to specify a new LLM to replace the current LLM.
-    db_session.exec(
-        update(ChatEngine).where(ChatEngine.llm_id == llm_id).values(llm_id=None)
-    )
-    db_session.exec(
-        update(ChatEngine)
-        .where(ChatEngine.fast_llm_id == llm_id)
-        .values(fast_llm_id=None)
-    )
-    db_session.exec(
-        update(KnowledgeBase).where(KnowledgeBase.llm_id == llm_id).values(llm_id=None)
-    )
-
-    # TODO: Should using soft deletion.
-    db_session.delete(llm)
-    db_session.commit()
-    return llm
+    llm_repo.delete(db_session, llm)
 
 
 @router.put("/admin/llms/{llm_id}/set_default")