✨ Develop datamate core final version part2

backend/agents/create_agent_info.py

@@ -150,6 +150,8 @@ async def create_agent_config(
                     tenant_id=tenant_id, user_id=user_id)
                 if knowledge_info_list:
                     for knowledge_info in knowledge_info_list:
+                        if knowledge_info.get('knowledge_sources') != 'elasticsearch':
+                            continue
                         knowledge_name = knowledge_info.get("index_name")
                         try:
                             message = ElasticSearchService().get_summary(index_name=knowledge_name)
@@ -239,13 +241,22 @@ async def create_tool_config_list(agent_id, tenant_id, user_id):
             knowledge_info_list = get_selected_knowledge_list(
                 tenant_id=tenant_id, user_id=user_id)
             index_names = [knowledge_info.get(
-                "index_name") for knowledge_info in knowledge_info_list]
+                "index_name") for knowledge_info in knowledge_info_list if knowledge_info.get('knowledge_sources') == 'elasticsearch']
             tool_config.metadata = {
                 "index_names": index_names,
                 "vdb_core": get_vector_db_core(),
                 "embedding_model": get_embedding_model(tenant_id=tenant_id),
                 "name_resolver": build_knowledge_name_mapping(tenant_id=tenant_id, user_id=user_id),
             }
+        elif tool_config.class_name == "DataMateSearchTool":
+            knowledge_info_list = get_selected_knowledge_list(
+                tenant_id=tenant_id, user_id=user_id)
+            index_names = [knowledge_info.get(
+                "index_name") for knowledge_info in knowledge_info_list if
+                knowledge_info.get('knowledge_sources') == 'datamate']
+            tool_config.metadata = {
+                "index_names": index_names,
+            }
         elif tool_config.class_name == "AnalyzeTextFileTool":
             tool_config.metadata = {
                 "llm_model": get_llm_model(tenant_id=tenant_id),

backend/apps/config_app.py

@@ -6,6 +6,7 @@
 
 from apps.agent_app import agent_config_router as agent_router
 from apps.config_sync_app import router as config_sync_router
+from apps.datamate_app import router as datamate_router
 from apps.vectordatabase_app import router as vectordatabase_router
 from apps.file_management_app import file_management_config_router as file_manager_router
 from apps.image_app import router as proxy_router
@@ -43,6 +44,7 @@
 app.include_router(config_sync_router)
 app.include_router(agent_router)
 app.include_router(vectordatabase_router)
+app.include_router(datamate_router)
 app.include_router(voice_router)
 app.include_router(file_manager_router)
 app.include_router(proxy_router)

backend/apps/config_sync_app.py

@@ -5,9 +5,11 @@
 from fastapi import APIRouter, Header, Request, HTTPException
 from fastapi.responses import JSONResponse
 
+from consts.const import DATAMATE_URL
 from consts.model import GlobalConfig
 from services.config_sync_service import save_config_impl, load_config_impl
 from utils.auth_utils import get_current_user_id, get_current_user_info
+from utils.config_utils import tenant_config_manager
 
 router = APIRouter(prefix="/config")
 logger = logging.getLogger("config_sync_app")
@@ -27,7 +29,43 @@ async def save_config(config: GlobalConfig, authorization: Optional[str] = Heade
         )
     except Exception as e:
         logger.error(f"Failed to save configuration: {str(e)}")
-        raise HTTPException(status_code=HTTPStatus.BAD_REQUEST, detail="Failed to save configuration.")
+        raise HTTPException(status_code=HTTPStatus.BAD_REQUEST,
+                            detail="Failed to save configuration.")
+
+
+@router.post("/save_datamate_url")
+async def save_datamate_url(data: dict, authorization: Optional[str] = Header(None)):
+    """
+    Save DataMate URL configuration
+
+    Args:
+        data: Dictionary containing datamate_url
+
+    Returns:
+        JSONResponse: Success message
+    """
+    try:
+        user_id, tenant_id = get_current_user_id(authorization)
+        datamate_url = data.get("datamate_url", "").strip()
+
+        if datamate_url:
+            tenant_config_manager.set_single_config(
+                user_id, tenant_id, DATAMATE_URL, datamate_url)
+            logger.info(f"DataMate URL saved successfully")
+        else:
+            # If empty, delete the configuration
+            tenant_config_manager.delete_single_config(tenant_id, DATAMATE_URL)
+            logger.info("DataMate URL deleted (empty value)")
+
+        return JSONResponse(
+            status_code=HTTPStatus.OK,
+            content={"message": "DataMate URL saved successfully",
+                     "status": "saved"}
+        )
+    except Exception as e:
+        logger.error(f"Failed to save DataMate URL: {str(e)}")
+        raise HTTPException(status_code=HTTPStatus.BAD_REQUEST,
+                            detail="Failed to save DataMate URL.")
 
 
 @router.get("/load_config")
@@ -49,4 +87,5 @@ async def load_config(authorization: Optional[str] = Header(None), request: Requ
         )
     except Exception as e:
         logger.error(f"Failed to load configuration: {str(e)}")
-        raise HTTPException(status_code=HTTPStatus.BAD_REQUEST, detail="Failed to load configuration.")
+        raise HTTPException(status_code=HTTPStatus.BAD_REQUEST,
+                            detail="Failed to load configuration.")

backend/apps/datamate_app.py

@@ -0,0 +1,48 @@
+import logging
+from typing import Optional
+
+from fastapi import APIRouter, Header, HTTPException, Path
+from fastapi.responses import JSONResponse
+from http import HTTPStatus
+
+from services.datamate_service import (
+    sync_datamate_knowledge_bases_and_create_records,
+    fetch_datamate_knowledge_base_file_list
+)
+from utils.auth_utils import get_current_user_id
+
+router = APIRouter(prefix="/datamate")
+logger = logging.getLogger("datamate_app")
+
+
+@router.post("/sync_datamate_knowledges")
+async def sync_datamate_knowledges(
+    authorization: Optional[str] = Header(None)
+):
+    """Sync DataMate knowledge bases and create knowledge records in local database."""
+    try:
+        user_id, tenant_id = get_current_user_id(authorization)
+
+        return await sync_datamate_knowledge_bases_and_create_records(
+            tenant_id=tenant_id,
+            user_id=user_id
+        )
+    except Exception as e:
+        raise HTTPException(
+            status_code=HTTPStatus.INTERNAL_SERVER_ERROR, detail=f"Error syncing DataMate knowledge bases and creating records: {str(e)}")
+
+
+@router.get("/{knowledge_base_id}/files")
+async def get_datamate_knowledge_base_files_endpoint(
+    knowledge_base_id: str = Path(...,
+                                  description="ID of the DataMate knowledge base"),
+    authorization: Optional[str] = Header(None)
+):
+    """Get all files from a DataMate knowledge base."""
+    try:
+        user_id, tenant_id = get_current_user_id(authorization)
+        result = await fetch_datamate_knowledge_base_file_list(knowledge_base_id, tenant_id)
+        return JSONResponse(status_code=HTTPStatus.OK, content=result)
+    except Exception as e:
+        raise HTTPException(
+            status_code=HTTPStatus.INTERNAL_SERVER_ERROR, detail=f"Error fetching DataMate knowledge base files: {str(e)}")

backend/apps/tenant_config_app.py

@@ -6,6 +6,7 @@
 from fastapi.responses import JSONResponse
 
 from consts.const import DEPLOYMENT_VERSION, APP_VERSION
+from consts.model import UpdateKnowledgeListRequest
 from services.tenant_config_service import get_selected_knowledge_list, update_selected_knowledge
 from utils.auth_utils import get_current_user_id
 
@@ -61,16 +62,37 @@ def load_knowledge_list(
 @router.post("/update_knowledge_list")
 def update_knowledge_list(
     authorization: Optional[str] = Header(None),
-    knowledge_list: List[str] = Body(None)
+    request: UpdateKnowledgeListRequest = Body(...)
 ):
     try:
         user_id, tenant_id = get_current_user_id(authorization)
+
+        # Convert grouped request to flat lists
+        knowledge_list = []
+        knowledge_sources = []
+
+        if request.nexent:
+            knowledge_list.extend(request.nexent)
+            knowledge_sources.extend(["nexent"] * len(request.nexent))
+
+        if request.datamate:
+            knowledge_list.extend(request.datamate)
+            knowledge_sources.extend(["datamate"] * len(request.datamate))
+
         result = update_selected_knowledge(
-            tenant_id=tenant_id, user_id=user_id, index_name_list=knowledge_list)
+            tenant_id=tenant_id, user_id=user_id, index_name_list=knowledge_list, knowledge_sources=knowledge_sources)
         if result:
+            # 获取更新后的知识库信息
+            selected_knowledge_info = get_selected_knowledge_list(
+                tenant_id=tenant_id, user_id=user_id)
+
+            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=HTTPStatus.OK,
-                content={"message": "update success", "status": "success"}
+                content={"content": content, "message": "update success", "status": "success"}
             )
         else:
             raise HTTPException(

backend/database/client.py

@@ -73,7 +73,7 @@ def clean_string_values(data: Dict[str, Any]) -> Dict[str, Any]:
 class MinioClient:
     """
     MinIO client wrapper using storage SDK
-    
+
     This class maintains backward compatibility with the existing MinioClient interface
     while using the new storage SDK under the hood.
     """
@@ -86,7 +86,8 @@ def __new__(cls):
 
     def __init__(self):
         # Determine if endpoint uses HTTPS
-        secure = MINIO_ENDPOINT.startswith('https://') if MINIO_ENDPOINT else True
+        secure = MINIO_ENDPOINT.startswith(
+            'https://') if MINIO_ENDPOINT else True
         # Initialize storage client using SDK factory
         self.storage_config = MinIOStorageConfig(
             endpoint=MINIO_ENDPOINT,
@@ -96,7 +97,8 @@ def __init__(self):
             default_bucket=MINIO_DEFAULT_BUCKET,
             secure=secure
         )
-        self._storage_client = create_storage_client_from_config(self.storage_config)
+        self._storage_client = create_storage_client_from_config(
+            self.storage_config)
 
     def upload_file(
             self,

backend/database/db_models.py

@@ -1,4 +1,4 @@
-from sqlalchemy import Boolean, Column, Integer, JSON, Numeric, Sequence, String, Text, TIMESTAMP
+from sqlalchemy import BigInteger, Boolean, Column, Integer, JSON, Numeric, Sequence, String, Text, TIMESTAMP
 from sqlalchemy.orm import DeclarativeBase
 from sqlalchemy.sql import func
 
@@ -248,7 +248,7 @@ class KnowledgeRecord(TableBase):
     __tablename__ = "knowledge_record_t"
     __table_args__ = {"schema": "nexent"}
 
-    knowledge_id = Column(Integer, Sequence("knowledge_record_t_knowledge_id_seq", schema="nexent"),
+    knowledge_id = Column(BigInteger, Sequence("knowledge_record_t_knowledge_id_seq", schema="nexent"),
                           primary_key=True, nullable=False, doc="Knowledge base ID, unique primary key")
     index_name = Column(String(100), doc="Internal Elasticsearch index name")
     knowledge_name = Column(String(100), doc="User-facing knowledge base name")

backend/database/knowledge_db.py

@@ -79,7 +79,58 @@ def create_knowledge_record(query: Dict[str, Any]) -> Dict[str, Any]:
                 "knowledge_name": new_record.knowledge_name,
             }
     except SQLAlchemyError as e:
-        session.rollback()
+        raise e
+
+
+def upsert_knowledge_record(query: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Create or update a knowledge base record (upsert operation).
+    If a record with the same index_name and tenant_id exists, update it.
+    Otherwise, create a new record.
+
+    Args:
+        query: Dictionary containing knowledge base data, must include:
+            - index_name: Knowledge base name (used as unique identifier)
+            - tenant_id: Tenant ID
+            - knowledge_name: User-facing knowledge base name
+            - knowledge_describe: Knowledge base description
+            - knowledge_sources: Knowledge base sources (optional, default 'elasticsearch')
+            - embedding_model_name: Embedding model name
+            - user_id: User ID for created_by and updated_by fields
+
+    Returns:
+        Dict[str, Any]: Dictionary with 'knowledge_id' and 'index_name'
+    """
+    try:
+        with get_db_session() as session:
+            # Check if record exists
+            existing_record = session.query(KnowledgeRecord).filter(
+                KnowledgeRecord.index_name == query['index_name'],
+                KnowledgeRecord.tenant_id == query['tenant_id'],
+                KnowledgeRecord.delete_flag != 'Y'
+            ).first()
+
+            if existing_record:
+                # Update existing record
+                existing_record.knowledge_name = query.get('knowledge_name') or query.get('index_name')
+                existing_record.knowledge_describe = query.get('knowledge_describe', '')
+                existing_record.knowledge_sources = query.get('knowledge_sources', 'elasticsearch')
+                existing_record.embedding_model_name = query.get('embedding_model_name')
+                existing_record.updated_by = query.get('user_id')
+                existing_record.update_time = func.current_timestamp()
+
+                session.flush()
+                session.commit()
+                return {
+                    "knowledge_id": existing_record.knowledge_id,
+                    "index_name": existing_record.index_name,
+                    "knowledge_name": existing_record.knowledge_name,
+                }
+            else:
+                # Create new record
+                return create_knowledge_record(query)
+
+    except SQLAlchemyError as e:
         raise e
 
 
@@ -115,7 +166,6 @@ def update_knowledge_record(query: Dict[str, Any]) -> bool:
             session.commit()
             return True
     except SQLAlchemyError as e:
-        session.rollback()
         raise e
 
 
@@ -152,7 +202,6 @@ def delete_knowledge_record(query: Dict[str, Any]) -> bool:
             session.commit()
             return True
     except SQLAlchemyError as e:
-        session.rollback()
         raise e
 
 
@@ -239,6 +288,29 @@ def get_knowledge_info_by_tenant_id(tenant_id: str) -> List[Dict[str, Any]]:
         raise e
 
 
+def get_knowledge_info_by_tenant_and_source(tenant_id: str, knowledge_sources: str) -> List[Dict[str, Any]]:
+    """
+    Get knowledge base records by tenant ID and knowledge sources.
+
+    Args:
+        tenant_id: Tenant ID to filter by
+        knowledge_sources: Knowledge sources to filter by (e.g., 'datamate')
+
+    Returns:
+        List[Dict[str, Any]]: List of knowledge base record dictionaries
+    """
+    try:
+        with get_db_session() as session:
+            result = session.query(KnowledgeRecord).filter(
+                KnowledgeRecord.tenant_id == tenant_id,
+                KnowledgeRecord.knowledge_sources == knowledge_sources,
+                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:

backend/services/agent_service.py

@@ -1022,7 +1022,7 @@ async def export_agent_by_agent_id(agent_id: int, tenant_id: str, user_id: str)
 
     # Check if any tool is KnowledgeBaseSearchTool and set its metadata to empty dict
     for tool in tool_list:
-        if tool.class_name in ["KnowledgeBaseSearchTool", "AnalyzeTextFileTool", "AnalyzeImageTool"]:
+        if tool.class_name in ["KnowledgeBaseSearchTool", "AnalyzeTextFileTool", "AnalyzeImageTool", "DataMateSearchTool"]:
             tool.metadata = {}
 
     # Get model_id and model display name from agent_info

backend/services/config_sync_service.py

@@ -1,11 +1,12 @@
 import logging
-from typing import Optional
+from typing import Optional, Any
 
 from consts.const import (
     APP_DESCRIPTION,
     APP_NAME,
     AVATAR_URI,
     CUSTOM_ICON_URL,
+    DATAMATE_URL,
     DEFAULT_APP_DESCRIPTION_EN,
     DEFAULT_APP_DESCRIPTION_ZH,
     DEFAULT_APP_NAME_EN,
@@ -142,8 +143,9 @@ def build_app_config(language: str, tenant_id: str) -> dict:
             "avatarUri": tenant_config_manager.get_app_config(AVATAR_URI, tenant_id=tenant_id) or "",
             "customUrl": tenant_config_manager.get_app_config(CUSTOM_ICON_URL, tenant_id=tenant_id) or ""
         },
+        "datamateUrl": tenant_config_manager.get_app_config(DATAMATE_URL, tenant_id=tenant_id) or "",
         "modelEngineEnabled": str(MODEL_ENGINE_ENABLED).lower() == "true"
-    }
+        }
 
 
 def build_models_config(tenant_id: str) -> dict: