✨ record embedding model used by each knowledge base in PG

✨ record embedding model used by each knowledge base in PG

Author: liutao12138

backend/apps/elasticsearch_app.py

@@ -135,7 +135,7 @@ def get_list_indices(
     """List all user indices with optional stats"""
     try:
         user_id, tenant_id = get_current_user_id(authorization)
-        return ElasticSearchService.list_indices(pattern, include_stats, tenant_id, es_core)
+        return ElasticSearchService.list_indices(pattern, include_stats, tenant_id, user_id, es_core)
     except Exception as e:
         raise HTTPException(
             status_code=500, detail=f"Error get index: {str(e)}")

backend/apps/knowledge_app.py

@@ -1,11 +1,11 @@
 import logging
+from http import HTTPStatus
 from typing import List, Optional
 
-import requests
-from fastapi import APIRouter, Body, Header
+from fastapi import APIRouter, Body, Header, HTTPException
 from fastapi.responses import JSONResponse
 
-from consts.const import DEPLOYMENT_VERSION, ELASTICSEARCH_SERVICE
+from consts.const import DEPLOYMENT_VERSION
 from services.tenant_config_service import get_selected_knowledge_list, update_selected_knowledge
 from utils.auth_utils import get_current_user_id
 
@@ -20,16 +20,15 @@ def get_deployment_version():
     """
     try:
         return JSONResponse(
-            status_code=200,
+            status_code=HTTPStatus.OK,
             content={"deployment_version": DEPLOYMENT_VERSION,
                      "status": "success"}
         )
     except Exception as e:
         logger.error(f"Failed to get deployment version, error: {e}")
-        return JSONResponse(
-            status_code=500,
-            content={"message": "Failed to get deployment version",
-                     "status": "error"}
+        raise HTTPException(
+            status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
+            detail="Failed to get deployment version"
         )
 
 
@@ -42,44 +41,19 @@ def load_knowledge_list(
         selected_knowledge_info = get_selected_knowledge_list(
             tenant_id=tenant_id, user_id=user_id)
 
-        kb_name_list = [item["index_name"] for item in selected_knowledge_info]
-
-        # Query embedding models from Elasticsearch service
-        embedding_models = []
-        if kb_name_list:
-            try:
-                api_url = f"{ELASTICSEARCH_SERVICE}/indices?include_stats=true"
-                response = requests.get(api_url, timeout=10)
-                response.raise_for_status()
-
-                index_to_model = {}
-                for index_info in response.json().get("indices_info", []):
-                    index_name = index_info.get("name")
-                    embedding_model = index_info.get("stats", {}).get(
-                        "base_info", {}).get("embedding_model", "")
-                    if index_name:
-                        index_to_model[index_name] = embedding_model
-
-                embedding_models = [index_to_model.get(
-                    kb_name, "") for kb_name in kb_name_list]
-
-            except Exception as e:
-                logger.error(f"API call failed: {e}")
-                embedding_models = []
-
-        content = {"selectedKbNames": kb_name_list,
-                   "selectedKbModels": embedding_models,
+        content = {"selectedKbNames": [item["index_name"] for item in selected_knowledge_info],
+                   "selectedKbModels": [item["embedding_model_name"] for item in selected_knowledge_info],
                    "selectedKbSources": [item["knowledge_sources"] for item in selected_knowledge_info]}
 
         return JSONResponse(
-            status_code=200,
+            status_code=HTTPStatus.OK,
             content={"content": content, "status": "success"}
         )
     except Exception as e:
         logger.error(f"load knowledge list failed, error: {e}")
-        return JSONResponse(
-            status_code=400,
-            content={"message": "Failed to load configuration", "status": "error"}
+        raise HTTPException(
+            status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
+            detail="Failed to load configuration"
         )
 
 
@@ -94,18 +68,17 @@ def update_knowledge_list(
             tenant_id=tenant_id, user_id=user_id, index_name_list=knowledge_list)
         if result:
             return JSONResponse(
-                status_code=200,
+                status_code=HTTPStatus.OK,
                 content={"message": "update success", "status": "success"}
             )
         else:
-            return JSONResponse(
-                status_code=400,
-                content={"message": "update failed", "status": "error"}
+            raise HTTPException(
+                status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
+                detail="Failed to update configuration"
             )
     except Exception as e:
         logger.error(f"update knowledge list failed, error: {e}")
-        return JSONResponse(
-            status_code=400,
-            content={"message": "Failed to update configuration",
-                     "status": "error"}
+        raise HTTPException(
+            status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
+            detail="Failed to update configuration"
         )

backend/apps/tenant_config_app.py

@@ -282,6 +282,7 @@ class KnowledgeRecord(TableBase):
     index_name = Column(String(100), doc="Knowledge base name")
     knowledge_describe = Column(String(3000), doc="Knowledge base description")
     knowledge_sources = Column(String(300), doc="Knowledge base sources")
+    embedding_model_name = Column(String(200), doc="Embedding model name, used to record the embedding model used by the knowledge base")
     tenant_id = Column(String(100), doc="Tenant ID")
     delete_flag = Column(String(
         1), default="N", doc="Knowledge base status. Currently defaults to 1, if knowledge base status is 0, then this knowledge base is unavailable")

backend/database/db_models.py

@@ -18,6 +18,7 @@ def create_knowledge_record(query: Dict[str, Any]) -> int:
             - knowledge_status: Knowledge base status
             - user_id: Optional user ID for created_by and updated_by fields
             - tenant_id: Optional tenant ID for created_by and updated_by fields
+            - embedding_model_name: embedding model name for the knowledge base
 
     Returns:
         int: Newly created knowledge base ID
@@ -31,7 +32,8 @@ def create_knowledge_record(query: Dict[str, Any]) -> int:
                 "created_by": query.get("user_id"),
                 "updated_by": query.get("user_id"),
                 "knowledge_sources": query.get("knowledge_sources", "elasticsearch"),
-                "tenant_id": query.get("tenant_id")
+                "tenant_id": query.get("tenant_id"),
+                "embedding_model_name": query.get("embedding_model_name")
             }
 
             # Create new record
@@ -163,7 +165,8 @@ def get_knowledge_info_by_knowledge_ids(knowledge_ids: List[str]) -> List[Dict[s
                 knowledge_info.append({
                     "knowledge_id": item.knowledge_id,
                     "index_name": item.index_name,
-                    "knowledge_sources": item.knowledge_sources
+                    "knowledge_sources": item.knowledge_sources,
+                    "embedding_model_name": item.embedding_model_name
                 })
             return knowledge_info
     except SQLAlchemyError as e:
@@ -180,3 +183,29 @@ def get_knowledge_ids_by_index_names(index_names: List[str]) -> List[str]:
             return [item.knowledge_id for item in result]
     except SQLAlchemyError as e:
         raise e
+
+
+def get_knowledge_info_by_tenant_id(tenant_id: str) -> List[Dict[str, Any]]:
+    try:
+        with get_db_session() as session:
+            result = session.query(KnowledgeRecord).filter(
+                KnowledgeRecord.tenant_id == tenant_id,
+                KnowledgeRecord.delete_flag != 'Y'
+            ).all()
+            return [as_dict(item) for item in result]
+    except SQLAlchemyError as e:
+        raise e
+
+
+def update_model_name_by_index_name(index_name: str, embedding_model_name: str, tenant_id: str, user_id: str) -> bool:
+    try:
+        with get_db_session() as session:
+            session.query(KnowledgeRecord).filter(
+                KnowledgeRecord.index_name == index_name,
+                KnowledgeRecord.delete_flag != 'Y',
+                KnowledgeRecord.tenant_id == tenant_id
+            ).update({"embedding_model_name": embedding_model_name, "updated_by": user_id})
+            session.commit()
+            return True
+    except SQLAlchemyError as e:
+        raise e

backend/database/knowledge_db.py

@@ -1,8 +1,13 @@
+import logging
 from typing import Any, Dict
 
-from .client import get_db_session
-from .db_models import TenantConfig
+from sqlalchemy.exc import SQLAlchemyError
 
+from database.client import get_db_session
+from database.db_models import TenantConfig
+
+
+logger = logging.getLogger("tenant_config_db")
 
 def get_all_configs_by_tenant_id(tenant_id: str):
     with get_db_session() as session:
@@ -65,8 +70,9 @@ def insert_config(insert_data: Dict[str, Any]):
             session.add(TenantConfig(**insert_data))
             session.commit()
             return True
-        except Exception as e:
+        except SQLAlchemyError as e:
             session.rollback()
+            logger.error(f"insert config failed, error: {e}")
             return False
 
 
@@ -79,8 +85,9 @@ def delete_config_by_tenant_config_id(tenant_config_id: int):
             ).update({"delete_flag": "Y"})
             session.commit()
             return True
-        except Exception as e:
+        except SQLAlchemyError as e:
             session.rollback()
+            logger.error(f"delete config by tenant config id failed, error: {e}")
             return False
 
 
@@ -96,8 +103,9 @@ def delete_config(tenant_id: str, user_id: str, select_key: str, config_value: s
             ).update({"delete_flag": "Y"})
             session.commit()
             return True
-        except Exception as e:
+        except SQLAlchemyError as e:
             session.rollback()
+            logger.error(f"delete config failed, error: {e}")
             return False
 
 
@@ -110,8 +118,9 @@ def update_config_by_tenant_config_id(tenant_config_id: int, update_value: str):
             ).update({"config_value": update_value})
             session.commit()
             return True
-        except Exception as e:
+        except SQLAlchemyError as e:
             session.rollback()
+            logger.error(f"update config by tenant config id failed, error: {e}")
             return False
 
 
@@ -124,6 +133,7 @@ def update_config_by_tenant_config_id_and_data(tenant_config_id: int, insert_dat
             ).update(insert_data)
             session.commit()
             return True
-        except Exception as e:
+        except SQLAlchemyError as e:
             session.rollback()
+            logger.error(f"update config by tenant config id and data failed, error: {e}")
             return False