veadk-python 0.2.8__py3-none-any.whl → 0.2.10__py3-none-any.whl

veadk/integrations/ve_viking_db_memory/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

veadk/integrations/ve_viking_db_memory/ve_viking_db_memory.py

@@ -0,0 +1,293 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import threading
+
+from volcengine.ApiInfo import ApiInfo
+from volcengine.auth.SignerV4 import SignerV4
+from volcengine.base.Service import Service
+from volcengine.Credentials import Credentials
+from volcengine.ServiceInfo import ServiceInfo
+
+
+class VikingDBMemoryException(Exception):
+    def __init__(self, code, request_id, message=None):
+        self.code = code
+        self.request_id = request_id
+        self.message = "{}, code:{}，request_id:{}".format(
+            message, self.code, self.request_id
+        )
+
+    def __str__(self):
+        return self.message
+
+
+class VikingDBMemoryClient(Service):
+    _instance_lock = threading.Lock()
+
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(VikingDBMemoryClient, "_instance"):
+            with VikingDBMemoryClient._instance_lock:
+                if not hasattr(VikingDBMemoryClient, "_instance"):
+                    VikingDBMemoryClient._instance = object.__new__(cls)
+        return VikingDBMemoryClient._instance
+
+    def __init__(
+        self,
+        host="api-knowledgebase.mlp.cn-beijing.volces.com",
+        region="cn-beijing",
+        ak="",
+        sk="",
+        sts_token="",
+        scheme="http",
+        connection_timeout=30,
+        socket_timeout=30,
+    ):
+        self.service_info = VikingDBMemoryClient.get_service_info(
+            host, region, scheme, connection_timeout, socket_timeout
+        )
+        self.api_info = VikingDBMemoryClient.get_api_info()
+        super(VikingDBMemoryClient, self).__init__(self.service_info, self.api_info)
+        if ak:
+            self.set_ak(ak)
+        if sk:
+            self.set_sk(sk)
+        if sts_token:
+            self.set_session_token(session_token=sts_token)
+        try:
+            self.get_body("Ping", {}, json.dumps({}))
+        except Exception as e:
+            raise VikingDBMemoryException(
+                1000028, "missed", "host or region is incorrect: {}".format(str(e))
+            ) from None
+
+    def setHeader(self, header):
+        api_info = VikingDBMemoryClient.get_api_info()
+        for key in api_info:
+            for item in header:
+                api_info[key].header[item] = header[item]
+        self.api_info = api_info
+
+    @staticmethod
+    def get_service_info(host, region, scheme, connection_timeout, socket_timeout):
+        service_info = ServiceInfo(
+            host,
+            {"Host": host},
+            Credentials("", "", "air", region),
+            connection_timeout,
+            socket_timeout,
+            scheme=scheme,
+        )
+        return service_info
+
+    @staticmethod
+    def get_api_info():
+        api_info = {
+            "CreateCollection": ApiInfo(
+                "POST",
+                "/api/memory/collection/create",
+                {},
+                {},
+                {"Accept": "application/json", "Content-Type": "application/json"},
+            ),
+            "GetCollection": ApiInfo(
+                "POST",
+                "/api/memory/collection/info",
+                {},
+                {},
+                {"Accept": "application/json", "Content-Type": "application/json"},
+            ),
+            "DropCollection": ApiInfo(
+                "POST",
+                "/api/memory/collection/delete",
+                {},
+                {},
+                {"Accept": "application/json", "Content-Type": "application/json"},
+            ),
+            "UpdateCollection": ApiInfo(
+                "POST",
+                "/api/memory/collection/update",
+                {},
+                {},
+                {"Accept": "application/json", "Content-Type": "application/json"},
+            ),
+            "SearchMemory": ApiInfo(
+                "POST",
+                "/api/memory/search",
+                {},
+                {},
+                {"Accept": "application/json", "Content-Type": "application/json"},
+            ),
+            "AddMessages": ApiInfo(
+                "POST",
+                "/api/memory/messages/add",
+                {},
+                {},
+                {"Accept": "application/json", "Content-Type": "application/json"},
+            ),
+            "Ping": ApiInfo(
+                "GET",
+                "/api/memory/ping",
+                {},
+                {},
+                {"Accept": "application/json", "Content-Type": "application/json"},
+            ),
+        }
+        return api_info
+
+    def get_body(self, api, params, body):
+        if api not in self.api_info:
+            raise Exception("no such api")
+        api_info = self.api_info[api]
+        r = self.prepare_request(api_info, params)
+        r.headers["Content-Type"] = "application/json"
+        r.headers["Traffic-Source"] = "SDK"
+        r.body = body
+
+        SignerV4.sign(r, self.service_info.credentials)
+
+        url = r.build()
+        resp = self.session.get(
+            url,
+            headers=r.headers,
+            data=r.body,
+            timeout=(
+                self.service_info.connection_timeout,
+                self.service_info.socket_timeout,
+            ),
+        )
+        if resp.status_code == 200:
+            return json.dumps(resp.json())
+        else:
+            raise Exception(resp.text.encode("utf-8"))
+
+    def get_body_exception(self, api, params, body):
+        try:
+            res = self.get_body(api, params, body)
+        except Exception as e:
+            try:
+                res_json = json.loads(e.args[0].decode("utf-8"))
+            except Exception as e:
+                raise VikingDBMemoryException(
+                    1000028, "missed", "json load res error, res:{}".format(str(e))
+                ) from None
+            code = res_json.get("code", 1000028)
+            request_id = res_json.get("request_id", 1000028)
+            message = res_json.get("message", None)
+
+            raise VikingDBMemoryException(code, request_id, message)
+
+        if res == "":
+            raise VikingDBMemoryException(
+                1000028,
+                "missed",
+                "empty response due to unknown error, please contact customer service",
+            ) from None
+        return res
+
+    def get_exception(self, api, params):
+        try:
+            res = self.get(api, params)
+        except Exception as e:
+            try:
+                res_json = json.loads(e.args[0].decode("utf-8"))
+            except Exception as e:
+                raise VikingDBMemoryException(
+                    1000028, "missed", "json load res error, res:{}".format(str(e))
+                ) from None
+            code = res_json.get("code", 1000028)
+            request_id = res_json.get("request_id", 1000028)
+            message = res_json.get("message", None)
+            raise VikingDBMemoryException(code, request_id, message)
+        if res == "":
+            raise VikingDBMemoryException(
+                1000028,
+                "missed",
+                "empty response due to unknown error, please contact customer service",
+            ) from None
+        return res
+
+    def create_collection(
+        self,
+        collection_name,
+        description="",
+        custom_event_type_schemas=[],
+        custom_entity_type_schemas=[],
+        builtin_event_types=[],
+        builtin_entity_types=[],
+    ):
+        params = {
+            "CollectionName": collection_name,
+            "Description": description,
+            "CustomEventTypeSchemas": custom_event_type_schemas,
+            "CustomEntityTypeSchemas": custom_entity_type_schemas,
+            "BuiltinEventTypes": builtin_event_types,
+            "BuiltinEntityTypes": builtin_entity_types,
+        }
+        res = self.json("CreateCollection", {}, json.dumps(params))
+        return json.loads(res)
+
+    def get_collection(self, collection_name):
+        params = {"CollectionName": collection_name}
+        res = self.json("GetCollection", {}, json.dumps(params))
+        return json.loads(res)
+
+    def drop_collection(self, collection_name):
+        params = {"CollectionName": collection_name}
+        res = self.json("DropCollection", {}, json.dumps(params))
+        return json.loads(res)
+
+    def update_collection(
+        self,
+        collection_name,
+        custom_event_type_schemas=[],
+        custom_entity_type_schemas=[],
+        builtin_event_types=[],
+        builtin_entity_types=[],
+    ):
+        params = {
+            "CollectionName": collection_name,
+            "CustomEventTypeSchemas": custom_event_type_schemas,
+            "CustomEntityTypeSchemas": custom_entity_type_schemas,
+            "BuiltinEventTypes": builtin_event_types,
+            "BuiltinEntityTypes": builtin_entity_types,
+        }
+        res = self.json("UpdateCollection", {}, json.dumps(params))
+        return json.loads(res)
+
+    def search_memory(self, collection_name, query, filter, limit=10):
+        params = {
+            "collection_name": collection_name,
+            "limit": limit,
+            "filter": filter,
+        }
+        if query:
+            params["query"] = query
+        res = self.json("SearchMemory", {}, json.dumps(params))
+        return json.loads(res)
+
+    def add_messages(
+        self, collection_name, session_id, messages, metadata, entities=None
+    ):
+        params = {
+            "collection_name": collection_name,
+            "session_id": session_id,
+            "messages": messages,
+            "metadata": metadata,
+        }
+        if entities is not None:
+            params["entities"] = entities
+        res = self.json("AddMessages", {}, json.dumps(params))
+        return json.loads(res)

veadk/knowledgebase/backends/base_backend.py

@@ -29,19 +29,19 @@ class BaseKnowledgebaseBackend(ABC, BaseModel):
         """
 
     @abstractmethod
-    def add_from_directory(self, directory: str, **kwargs) -> bool:
+    def add_from_directory(self, directory: str, *args, **kwargs) -> bool:
         """Add knowledge from file path to knowledgebase"""
 
     @abstractmethod
-    def add_from_files(self, files: list[str], **kwargs) -> bool:
+    def add_from_files(self, files: list[str], *args, **kwargs) -> bool:
         """Add knowledge (e.g, documents, strings, ...) to knowledgebase"""
 
     @abstractmethod
-    def add_from_text(self, text: str | list[str], **kwargs) -> bool:
+    def add_from_text(self, text: str | list[str], *args, **kwargs) -> bool:
         """Add knowledge from text to knowledgebase"""
 
     @abstractmethod
-    def search(self, **kwargs) -> list:
+    def search(self, *args, **kwargs) -> list:
         """Search knowledge from knowledgebase"""
 
     # Optional methods for future use:

veadk/knowledgebase/backends/vikingdb_knowledge_backend.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import asyncio
+import json
 import re
 from pathlib import Path
 from typing import Any, Literal
@@ -24,9 +25,9 @@ from typing_extensions import override
 import veadk.config  # noqa E401
 from veadk.config import getenv
 from veadk.configs.database_configs import NormalTOSConfig, TOSConfig
-from veadk.consts import DEFAULT_TOS_BUCKET_NAME
 from veadk.knowledgebase.backends.base_backend import BaseKnowledgebaseBackend
 from veadk.knowledgebase.backends.utils import build_vikingdb_knowledgebase_request
+from veadk.knowledgebase.entry import KnowledgebaseEntry
 from veadk.utils.logger import get_logger
 from veadk.utils.misc import formatted_timestamp
 
@@ -48,13 +49,6 @@ def _read_file_to_bytes(file_path: str) -> tuple[bytes, str]:
     return file_content, file_name
 
 
-def _extract_tos_attributes(**kwargs) -> tuple[str, str]:
-    """Extract TOS attributes from kwargs"""
-    tos_bucket_name = kwargs.get("tos_bucket_name", DEFAULT_TOS_BUCKET_NAME)
-    tos_bucket_path = kwargs.get("tos_bucket_path", "knowledgebase")
-    return tos_bucket_name, tos_bucket_path
-
-
 def get_files_in_directory(directory: str):
     dir_path = Path(directory)
     if not dir_path.is_dir():
@@ -109,15 +103,27 @@ class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
         )
 
     @override
-    def add_from_directory(self, directory: str, **kwargs) -> bool:
-        """
+    def add_from_directory(
+        self,
+        directory: str,
+        tos_bucket_name: str | None = None,
+        tos_bucket_path: str = "knowledgebase",
+        metadata: dict | None = None,
+        **kwargs,
+    ) -> bool:
+        """Add knowledge from a directory to the knowledgebase.
+
         Args:
-            directory: str, the directory to add to knowledgebase
-            **kwargs:
-                - tos_bucket_name: str, the bucket name of TOS
-                - tos_bucket_path: str, the path of TOS bucket
+            directory (str): The directory to add to knowledgebase.
+            tos_bucket_name (str | None, optional): The bucket name of TOS. Defaults to None.
+            tos_bucket_path (str, optional): The path of TOS bucket. Defaults to "knowledgebase".
+            metadata (dict | None, optional): The metadata of the files. Defaults to None.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            bool: True if successful, False otherwise.
         """
-        tos_bucket_name, tos_bucket_path = _extract_tos_attributes(**kwargs)
+        tos_bucket_name = tos_bucket_name or self.tos_config.bucket
         files = get_files_in_directory(directory=directory)
         for _file in files:
             content, file_name = _read_file_to_bytes(_file)
@@ -125,40 +131,64 @@ class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
                 content,
                 tos_bucket_name=tos_bucket_name,
                 object_key=f"{tos_bucket_path}/{file_name}",
+                metadata=metadata,
             )
             self._add_doc(tos_url=tos_url)
         return True
 
     @override
-    def add_from_files(self, files: list[str], **kwargs) -> bool:
-        """
+    def add_from_files(
+        self,
+        files: list[str],
+        tos_bucket_name: str | None = None,
+        tos_bucket_path: str = "knowledgebase",
+        metadata: dict | None = None,
+        **kwargs,
+    ) -> bool:
+        """Add knowledge from a directory to the knowledgebase.
+
         Args:
-            files:  list[str], the files to add to knowledgebase
-            **kwargs:
-                - tos_bucket_name: str, the bucket name of TOS
-                - tos_bucket_path: str, the path of TOS bucket
+            files (list[str]): The files to add to knowledgebase.
+            tos_bucket_name (str | None, optional): The bucket name of TOS. Defaults to None.
+            tos_bucket_path (str, optional): The path of TOS bucket. Defaults to "knowledgebase".
+            metadata (dict | None, optional): The metadata of the files. Defaults to None.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            bool: True if successful, False otherwise.
         """
-        tos_bucket_name, tos_bucket_path = _extract_tos_attributes(**kwargs)
+        tos_bucket_name = tos_bucket_name or self.tos_config.bucket
         for _file in files:
             content, file_name = _read_file_to_bytes(_file)
             tos_url = self._upload_bytes_to_tos(
                 content,
                 tos_bucket_name=tos_bucket_name,
                 object_key=f"{tos_bucket_path}/{file_name}",
+                metadata=metadata,
             )
             self._add_doc(tos_url=tos_url)
         return True
 
     @override
-    def add_from_text(self, text: str | list[str], **kwargs) -> bool:
-        """
+    def add_from_text(
+        self,
+        text: str | list[str],
+        tos_bucket_name: str | None = None,
+        tos_bucket_path: str = "knowledgebase",
+        metadata: dict | None = None,
+        **kwargs,
+    ) -> bool:
+        """Add knowledge from text to the knowledgebase.
+
         Args:
-            text:   str or list[str], the text to add to knowledgebase
-            **kwargs:
-                - tos_bucket_name: str, the bucket name of TOS
-                - tos_bucket_path: str, the path of TOS bucket
+            text (str | list[str]): The text to add to knowledgebase.
+            tos_bucket_name (str | None, optional): The bucket name of TOS. Defaults to None.
+            tos_bucket_path (str, optional): The path of TOS bucket. Defaults to "knowledgebase".
+
+        Returns:
+            bool: True if successful, False otherwise.
         """
-        tos_bucket_name, tos_bucket_path = _extract_tos_attributes(**kwargs)
+        tos_bucket_name = tos_bucket_name or self.tos_config.bucket
         if isinstance(text, list):
             object_keys = kwargs.get(
                 "tos_object_keys",
@@ -170,7 +200,7 @@ class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
             for _text, _object_key in zip(text, object_keys):
                 _content = _text.encode("utf-8")
                 tos_url = self._upload_bytes_to_tos(
-                    _content, tos_bucket_name, _object_key
+                    _content, tos_bucket_name, _object_key, metadata=metadata
                 )
                 self._add_doc(tos_url=tos_url)
             return True
@@ -179,26 +209,42 @@ class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
             object_key = kwargs.get(
                 "object_key", f"veadk/knowledgebase/{formatted_timestamp()}.txt"
             )
-            tos_url = self._upload_bytes_to_tos(content, tos_bucket_name, object_key)
+            tos_url = self._upload_bytes_to_tos(
+                content, tos_bucket_name, object_key, metadata=metadata
+            )
             self._add_doc(tos_url=tos_url)
         else:
             raise ValueError("text must be str or list[str]")
         return True
 
-    def add_from_bytes(self, content: bytes, file_name: str, **kwargs) -> bool:
-        """
+    def add_from_bytes(
+        self,
+        content: bytes,
+        file_name: str,
+        tos_bucket_name: str | None = None,
+        tos_bucket_path: str = "knowledgebase",
+        metadata: dict | None = None,
+        **kwargs,
+    ) -> bool:
+        """Add knowledge from bytes to the knowledgebase.
+
         Args:
-            content: bytes, the content to add to knowledgebase, bytes
-            file_name: str, the file name of the content
-            **kwargs:
-                - tos_bucket_name: str, the bucket name of TOS
-                - tos_bucket_path: str, the path of TOS bucket
+            content (bytes): The content to add to knowledgebase.
+            file_name (str): The file name of the content.
+            tos_bucket_name (str | None, optional): The bucket name of TOS. Defaults to None.
+            tos_bucket_path (str, optional): The path of TOS bucket. Defaults to "knowledgebase".
+            metadata (dict | None, optional): The metadata of the files. Defaults to None.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            bool: True if successful, False otherwise.
         """
-        tos_bucket_name, tos_bucket_path = _extract_tos_attributes(**kwargs)
+        tos_bucket_name = tos_bucket_name or self.tos_config.bucket
         tos_url = self._upload_bytes_to_tos(
             content,
             tos_bucket_name=tos_bucket_name,
             object_key=f"{tos_bucket_path}/{file_name}",
+            metadata=metadata,
         )
         response = self._add_doc(tos_url=tos_url)
         if response["code"] == 0:
@@ -206,8 +252,16 @@ class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
         return False
 
     @override
-    def search(self, query: str, top_k: int = 5) -> list:
-        return self._search_knowledge(query=query, top_k=top_k)
+    def search(
+        self,
+        query: str,
+        top_k: int = 5,
+        metadata: dict | None = None,
+        rerank: bool = True,
+    ) -> list:
+        return self._search_knowledge(
+            query=query, top_k=top_k, metadata=metadata, rerank=rerank
+        )
 
     def delete_collection(self) -> bool:
         DELETE_COLLECTION_PATH = "/api/knowledge/collection/delete"
@@ -330,7 +384,7 @@ class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
         response = self._do_request(
             body={
                 "name": self.index,
-                "project": "default",
+                "project": self.volcengine_project,
                 "description": "Created by Volcengine Agent Development Kit (VeADK).",
             },
             path=CREATE_COLLECTION_PATH,
@@ -343,10 +397,27 @@ class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
             )
 
     def _upload_bytes_to_tos(
-        self, content: bytes, tos_bucket_name: str, object_key: str
+        self,
+        content: bytes,
+        tos_bucket_name: str,
+        object_key: str,
+        metadata: dict | None = None,
     ) -> str:
+        # Here, we set the metadata via the TOS object, ref: https://www.volcengine.com/docs/84313/1254624
         self._tos_client.bucket_name = tos_bucket_name
-        asyncio.run(self._tos_client.upload(object_key=object_key, data=content))
+        coro = self._tos_client.upload(
+            object_key=object_key,
+            bucket_name=tos_bucket_name,
+            data=content,
+            metadata=metadata,
+        )
+        try:
+            loop = asyncio.get_running_loop()
+            loop.run_until_complete(
+                coro
+            ) if not loop.is_running() else asyncio.ensure_future(coro)
+        except RuntimeError:
+            asyncio.run(coro)
         return f"{self._tos_client.bucket_name}/{object_key}"
 
     def _add_doc(self, tos_url: str) -> Any:
@@ -355,7 +426,7 @@ class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
         response = self._do_request(
             body={
                 "collection_name": self.index,
-                "project": "default",
+                "project": self.volcengine_project,
                 "add_type": "tos",
                 "tos_path": tos_url,
             },
@@ -364,14 +435,43 @@ class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
         )
         return response
 
-    def _search_knowledge(self, query: str, top_k: int = 5) -> list[str]:
+    def _search_knowledge(
+        self,
+        query: str,
+        top_k: int = 5,
+        metadata: dict | None = None,
+        rerank: bool = True,
+        chunk_diffusion_count: int | None = 3,
+    ) -> list[KnowledgebaseEntry]:
         SEARCH_KNOWLEDGE_PATH = "/api/knowledge/collection/search_knowledge"
 
+        query_param = (
+            {
+                "doc_filter": {
+                    "op": "and",
+                    "conds": [
+                        {"op": "must", "field": str(k), "conds": [str(v)]}
+                        for k, v in metadata.items()
+                    ],
+                }
+            }
+            if metadata
+            else None
+        )
+
+        post_precessing = {
+            "rerank_swich": rerank,
+            "chunk_diffusion_count": chunk_diffusion_count,
+        }
+
         response = self._do_request(
             body={
                 "name": self.index,
+                "project": self.volcengine_project,
                 "query": query,
                 "limit": top_k,
+                "query_param": query_param,
+                "post_processing": post_precessing,
             },
             path=SEARCH_KNOWLEDGE_PATH,
             method="POST",
@@ -382,11 +482,19 @@ class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
                 f"Error during knowledge search: {response.get('code')}, message: {response.get('message')}"
             )
 
-        search_result_list = response.get("data", {}).get("result_list", [])
+        entries = []
+        for result in response.get("data", {}).get("result_list", []):
+            doc_meta_raw_str = result.get("doc_info", {}).get("doc_meta")
+            doc_meta_list = json.loads(doc_meta_raw_str) if doc_meta_raw_str else []
+            metadata = {}
+            for meta in doc_meta_list:
+                metadata[meta["field_name"]] = meta["field_value"]
 
-        return [
-            search_result.get("content", "") for search_result in search_result_list
-        ]
+            entries.append(
+                KnowledgebaseEntry(content=result.get("content", ""), metadata=metadata)
+            )
+
+        return entries
 
     def _do_request(
         self,
@@ -409,4 +517,8 @@ class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
             headers=request.headers,
             data=request.body,
         )
+        if not response.ok:
+            logger.error(
+                f"VikingDBKnowledgeBackend error during request: {response.json()}"
+            )
         return response.json()