✨Added Datamate vector knowledge base core

Author: Zhi-a

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/datamate_app.py

@@ -6,8 +6,7 @@
 from http import HTTPStatus
 
 from services.datamate_service import (
-    sync_datamate_knowledge_bases,
-    fetch_datamate_knowledge_bases,
+    sync_datamate_knowledge_bases_and_create_records,
     fetch_datamate_knowledge_base_file_list
 )
 from utils.auth_utils import get_current_user_id
@@ -16,62 +15,23 @@
 logger = logging.getLogger("datamate_app")
 
 
-@router.get("/knowledge_bases")
-async def get_datamate_knowledge_bases_endpoint(
-    authorization: Optional[str] = Header(None)
-):
-    """Get list of DataMate knowledge bases."""
-    try:
-        knowledge_bases = await fetch_datamate_knowledge_bases()
-
-        # Transform to the same format as list_indices method
-        indices = []
-        indices_info = []
-
-        for kb in knowledge_bases:
-            kb_id = kb.get("id")
-            kb_name = kb.get("name") or kb_id
-
-            # Get stats from the knowledge base data
-            stats = kb.get("stats", {})
-            chunk_count = kb.get("chunkCount", 0)
-            doc_count = kb.get("docCount", 0)
 
-            indices.append(kb_name)
-            indices_info.append({
-                "name": kb_id,  # Internal index name (used as ID)
-                "display_name": kb_name,  # User-facing knowledge base name
-                "stats": {
-                    "base_info": {
-                        "doc_count": doc_count,
-                        "chunk_count": chunk_count,
-                        "creation_date": kb.get("createdAt"),
-                        "update_date": kb.get("updatedAt"),
-                        "embedding_model": kb.get("embeddingModel", "unknown"),
-                    }
-                }
-            })
 
-        return {
-            "indices": indices,
-            "count": len(indices),
-            "indices_info": indices_info,
-        }
-    except Exception as e:
-        raise HTTPException(
-            status_code=HTTPStatus.INTERNAL_SERVER_ERROR, detail=f"Error fetching DataMate knowledge bases: {str(e)}")
-
-
-@router.post("/sync")
-async def sync_datamate_knowledge_bases_endpoint(
+@router.post("/sync_and_create_records")
+async def sync_datamate_and_create_records_endpoint(
     authorization: Optional[str] = Header(None)
 ):
-    """Sync DataMate knowledge bases and their files."""
+    """Sync DataMate knowledge bases and create knowledge records in local database."""
     try:
-        return await sync_datamate_knowledge_bases()
+        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: {str(e)}")
+            status_code=HTTPStatus.INTERNAL_SERVER_ERROR, detail=f"Error syncing DataMate knowledge bases and creating records: {str(e)}")
 
 
 @router.get("/{knowledge_base_id}/files")

backend/database/knowledge_db.py

@@ -79,6 +79,59 @@ def create_knowledge_record(query: Dict[str, Any]) -> Dict[str, Any]:
         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:
+        session.rollback()
+        raise e
+
+
 def update_knowledge_record(query: Dict[str, Any]) -> bool:
     """
     Update a knowledge base record
@@ -230,6 +283,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

@@ -97,45 +97,45 @@ def _resolve_model_with_fallback(
 ) -> str | None:
     """
     Resolve model_id from model_display_name with fallback to quick config LLM model.
-    
+
     Args:
         model_display_name: Display name of the model to lookup
         exported_model_id: Original model_id from export (for logging only)
         model_label: Label for logging (e.g., "Model", "Business logic model")
         tenant_id: Tenant ID for model lookup
-    
+
     Returns:
         Resolved model_id or None if not found and no fallback available
     """
     if not model_display_name:
         return None
-    
+
     # Try to find model by display name in current tenant
     resolved_id = get_model_id_by_display_name(model_display_name, tenant_id)
-    
+
     if resolved_id:
         logger.info(
             f"{model_label} '{model_display_name}' found in tenant {tenant_id}, "
             f"mapped to model_id: {resolved_id} (exported model_id was: {exported_model_id})")
         return resolved_id
-    
+
     # Model not found, try fallback to quick config LLM model
     logger.warning(
         f"{model_label} '{model_display_name}' (exported model_id: {exported_model_id}) "
         f"not found in tenant {tenant_id}, falling back to quick config LLM model.")
-    
+
     quick_config_model = tenant_config_manager.get_model_config(
         key=MODEL_CONFIG_MAPPING["llm"],
         tenant_id=tenant_id
     )
-    
+
     if quick_config_model:
         fallback_id = quick_config_model.get("model_id")
         logger.info(
             f"Using quick config LLM model for {model_label.lower()}: "
             f"{quick_config_model.get('display_name')} (model_id: {fallback_id})")
         return fallback_id
-    
+
     logger.warning(f"No quick config LLM model found for tenant {tenant_id}")
     return None
 
@@ -998,7 +998,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
@@ -1132,7 +1132,7 @@ async def import_agent_by_agent_id(
     if not import_agent_info.name.isidentifier():
         raise ValueError(
             f"Invalid agent name: {import_agent_info.name}. agent name must be a valid python variable name.")
-    
+
     # Resolve model IDs with fallback
     # Note: We use model_display_name for cross-tenant compatibility
     # The exported model_id is kept for reference/debugging only
@@ -1142,7 +1142,7 @@ async def import_agent_by_agent_id(
         model_label="Model",
         tenant_id=tenant_id
     )
-    
+
     business_logic_model_id = _resolve_model_with_fallback(
         model_display_name=import_agent_info.business_logic_model_name,
         exported_model_id=import_agent_info.business_logic_model_id,
@@ -1344,44 +1344,44 @@ def check_agent_availability(
 ) -> tuple[bool, list[str]]:
     """
     Check if an agent is available based on its tools and model configuration.
-    
+
     Args:
         agent_id: The agent ID to check
         tenant_id: The tenant ID
         agent_info: Optional pre-fetched agent info (to avoid duplicate DB queries)
         model_cache: Optional model cache for performance optimization
-        
+
     Returns:
         tuple: (is_available: bool, unavailable_reasons: list[str])
     """
     unavailable_reasons: list[str] = []
-    
+
     if model_cache is None:
         model_cache = {}
-    
+
     # Fetch agent info if not provided
     if agent_info is None:
         agent_info = search_agent_info_by_agent_id(agent_id, tenant_id)
-    
+
     if not agent_info:
         return False, ["agent_not_found"]
-    
+
     # Check tool availability
     tool_info = search_tools_for_sub_agent(agent_id=agent_id, tenant_id=tenant_id)
     tool_id_list = [tool["tool_id"] for tool in tool_info if tool.get("tool_id") is not None]
     if tool_id_list:
         tool_statuses = check_tool_is_available(tool_id_list)
         if not all(tool_statuses):
             unavailable_reasons.append("tool_unavailable")
-    
+
     # Check model availability
     model_reasons = _collect_model_availability_reasons(
         agent=agent_info,
         tenant_id=tenant_id,
         model_cache=model_cache
     )
     unavailable_reasons.extend(model_reasons)
-    
+
     is_available = len(unavailable_reasons) == 0
     return is_available, unavailable_reasons
 
@@ -1935,4 +1935,4 @@ def get_sub_agents_recursive(parent_agent_id: int, depth: int = 0, max_depth: in
     except Exception as e:
         logger.exception(
             f"Failed to get agent call relationship for agent {agent_id}: {str(e)}")
-        raise ValueError(f"Failed to get agent call relationship: {str(e)}")
+        raise ValueError(f"Failed to get agent call relationship: {str(e)}")