✨ Add Datamate knowledge base SDK

Author: Zhi-a

sdk/nexent/core/tools/datamate_search_tool.py

@@ -1,15 +1,14 @@
 import json
 import logging
-from typing import List, Optional
+from typing import Optional
 
-import httpx
 from pydantic import Field
 from smolagents.tools import Tool
 
+from ...datamate import DataMateClient
 from ..utils.observer import MessageObserver, ProcessType
 from ..utils.tools_common_message import SearchResultTextMessage, ToolCategory, ToolSign
 
-
 # Get logger instance
 logger = logging.getLogger("datamate_search_tool")
 
@@ -61,10 +60,10 @@ class DataMateSearchTool(Tool):
     tool_sign = ToolSign.DATAMATE_KNOWLEDGE_BASE.value
 
     def __init__(
-        self,
-        server_ip: str = Field(description="DataMate server IP or hostname"),
-        server_port: int = Field(description="DataMate server port"),
-        observer: MessageObserver = Field(description="Message observer", default=None, exclude=True),
+            self,
+            server_ip: str = Field(description="DataMate server IP or hostname"),
+            server_port: int = Field(description="DataMate server port"),
+            observer: MessageObserver = Field(description="Message observer", default=None, exclude=True),
     ):
         """Initialize the DataMateSearchTool.
 
@@ -88,6 +87,9 @@ def __init__(
         # Build base URL: http://host:port
         self.server_base_url = f"http://{self.server_ip}:{self.server_port}".rstrip("/")
 
+        # Initialize DataMate SDK client
+        self.datamate_client = DataMateClient(base_url=self.server_base_url)
+
         self.kb_page = 0
         self.kb_page_size = 20
         self.observer = observer
@@ -97,12 +99,12 @@ def __init__(
         self.running_prompt_en = "Searching the DataMate knowledge base..."
 
     def forward(
-        self,
-        query: str,
-        top_k: int = 10,
-        threshold: float = 0.2,
-        kb_page: int = 0,
-        kb_page_size: int = 20,
+            self,
+            query: str,
+            top_k: int = 10,
+            threshold: float = 0.2,
+            kb_page: int = 0,
+            kb_page_size: int = 20,
     ) -> str:
         """Execute DataMate search.
 
@@ -130,17 +132,37 @@ def forward(
         )
 
         try:
-            # Step 1: Get knowledge base list
-            knowledge_base_ids = self._get_knowledge_base_list()
+            # Step 1: Get knowledge base list using SDK
+            knowledge_bases = self.datamate_client.list_knowledge_bases(
+                page=self.kb_page,
+                size=self.kb_page_size
+            )
+
+            # Extract knowledge base IDs
+            knowledge_base_ids = []
+            for kb in knowledge_bases:
+                kb_id = kb.get("id")
+                chunk_count = kb.get("chunkCount")
+                if kb_id and chunk_count:
+                    knowledge_base_ids.append(str(kb_id))
+
             if not knowledge_base_ids:
                 return json.dumps("No knowledge base found. No relevant information found.", ensure_ascii=False)
 
-            # Step 2: Retrieve knowledge base content
-            kb_search_results = self._retrieve_knowledge_base_content(query, knowledge_base_ids, top_k, threshold
-            )
+            # Step 2: Retrieve knowledge base content using SDK
+            kb_search_results = []
+            for knowledge_base_id in knowledge_base_ids:
 
-            if not kb_search_results:
-                raise Exception("No results found! Try a less restrictive/shorter query.")
+                kb_search = self.datamate_client.retrieve_knowledge_base(
+                    query=query,
+                    knowledge_base_ids=[knowledge_base_id],
+                    top_k=top_k,
+                    threshold=threshold
+                )
+
+                if not kb_search:
+                    raise Exception("No results found! Try a less restrictive/shorter query.")
+                kb_search_results.extend(kb_search)
 
             # Format search results
             search_results_json = []  # Organize search results into a unified format
@@ -150,8 +172,8 @@ def forward(
                 entity_data = single_search_result.get("entity", {})
                 metadata = self._parse_metadata(entity_data.get("metadata"))
                 dataset_id = self._extract_dataset_id(metadata.get("absolute_directory_path", ""))
-                file_id = entity_data.get("id")
-                download_url = self._build_file_download_url(dataset_id, file_id)
+                file_id = metadata.get("original_file_id")
+                download_url = self.datamate_client.build_file_download_url(dataset_id, file_id)
 
                 score_details = entity_data.get("scoreDetails", {}) or {}
                 score_details.update({
@@ -191,100 +213,6 @@ def forward(
             logger.error(error_msg)
             raise Exception(error_msg)
 
-    def _get_knowledge_base_list(self) -> List[str]:
-        """Get knowledge base list from DataMate API.
-
-        Returns:
-            List[str]: List of knowledge base IDs.
-        """
-        try:
-            url = f"{self.server_base_url}/api/knowledge-base/list"
-            payload = {"page": self.kb_page, "size": self.kb_page_size}
-
-            with httpx.Client(timeout=30) as client:
-                response = client.post(url, json=payload)
-
-            if response.status_code != 200:
-                error_detail = (
-                    response.json().get("detail", "unknown error")
-                    if response.headers.get("content-type", "").startswith("application/json")
-                    else response.text
-                )
-                raise Exception(f"Failed to get knowledge base list (status {response.status_code}): {error_detail}")
-
-            result = response.json()
-            # Extract knowledge base IDs from response
-            # Assuming the response structure contains a list of knowledge bases with 'id' field
-            data = result.get("data", {})
-            knowledge_bases = data.get("content", []) if data else []
-
-            knowledge_base_ids = []
-            for kb in knowledge_bases:
-                kb_id = kb.get("id")
-                chunk_count = kb.get("chunkCount")
-                if kb_id and chunk_count:
-                    knowledge_base_ids.append(str(kb_id))
-
-            logger.info(f"Retrieved {len(knowledge_base_ids)} knowledge base(s): {knowledge_base_ids}")
-            return knowledge_base_ids
-
-        except httpx.TimeoutException:
-            raise Exception("Timeout while getting knowledge base list from DataMate API")
-        except httpx.RequestError as e:
-            raise Exception(f"Request error while getting knowledge base list: {str(e)}")
-        except Exception as e:
-            raise Exception(f"Error getting knowledge base list: {str(e)}")
-
-    def _retrieve_knowledge_base_content(
-        self, query: str, knowledge_base_ids: List[str], top_k: int, threshold: float
-    ) -> List[dict]:
-        """Retrieve knowledge base content from DataMate API.
-
-        Args:
-            query (str): Search query.
-            knowledge_base_ids (List[str]): List of knowledge base IDs to search.
-            top_k (int): Maximum number of results to return.
-            threshold (float): Similarity threshold.
-
-        Returns:
-            List[dict]: List of search results.
-        """
-        search_results = []
-        for knowledge_base_id in knowledge_base_ids:
-            try:
-                url = f"{self.server_base_url}/api/knowledge-base/retrieve"
-                payload = {
-                    "query": query,
-                    "topK": top_k,
-                    "threshold": threshold,
-                    "knowledgeBaseIds": [knowledge_base_id],
-                }
-
-                with httpx.Client(timeout=60) as client:
-                    response = client.post(url, json=payload)
-
-                if response.status_code != 200:
-                    error_detail = (
-                        response.json().get("detail", "unknown error")
-                        if response.headers.get("content-type", "").startswith("application/json")
-                        else response.text
-                    )
-                    raise Exception(
-                        f"Failed to retrieve knowledge base content (status {response.status_code}): {error_detail}")
-
-                result = response.json()
-                # Extract search results from response
-                for data in result.get("data", {}):
-                    search_results.append(data)
-            except httpx.TimeoutException:
-                raise Exception("Timeout while retrieving knowledge base content from DataMate API")
-            except httpx.RequestError as e:
-                raise Exception(f"Request error while retrieving knowledge base content: {str(e)}")
-            except Exception as e:
-                raise Exception(f"Error retrieving knowledge base content: {str(e)}")
-        logger.info(f"Retrieved {len(search_results)} search result(s)")
-        return search_results
-
     @staticmethod
     def _parse_metadata(metadata_raw: Optional[str]) -> dict:
         """Parse metadata payload safely."""
@@ -304,10 +232,4 @@ def _extract_dataset_id(absolute_path: str) -> str:
         if not absolute_path:
             return ""
         segments = [segment for segment in absolute_path.strip("/").split("/") if segment]
-        return segments[-1] if segments else ""
-
-    def _build_file_download_url(self, dataset_id: str, file_id: str) -> str:
-        """Build the download URL for a dataset file."""
-        if not (self.server_ip and dataset_id and file_id):
-            return ""
-        return f"{self.server_ip}/api/data-management/datasets/{dataset_id}/files/{file_id}/download"
+        return segments[-1] if segments else ""

sdk/nexent/datamate/__init__.py

@@ -0,0 +1,7 @@
+"""
+DataMate SDK client for interacting with DataMate knowledge base APIs.
+"""
+from .datamate_client import DataMateClient
+
+__all__ = ["DataMateClient"]
+