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

veadk/integrations/ve_tls/ve_tls.py

@@ -0,0 +1,208 @@
+# 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 os
+
+from veadk.consts import DEFAULT_TLS_LOG_PROJECT_NAME, DEFAULT_TLS_TRACING_INSTANCE_NAME
+from veadk.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class VeTLS:
+    def __init__(
+        self,
+        access_key: str | None = None,
+        secret_key: str | None = None,
+        region: str = "cn-beijing",
+    ):
+        try:
+            from volcengine.tls.TLSService import TLSService
+
+            from veadk.integrations.ve_tls.utils import ve_tls_request
+        except ImportError:
+            raise ImportError(
+                "Please install volcengine SDK before init VeTLS: pip install volcengine"
+            )
+
+        self._ve_tls_request = ve_tls_request
+
+        self.access_key = (
+            access_key if access_key else os.getenv("VOLCENGINE_ACCESS_KEY", "")
+        )
+        self.secret_key = (
+            secret_key if secret_key else os.getenv("VOLCENGINE_SECRET_KEY", "")
+        )
+        self.region = region
+
+        self._client = TLSService(
+            endpoint=f"https://tls-{self.region}.volces.com",
+            access_key_id=self.access_key,
+            access_key_secret=self.secret_key,
+            region=self.region,
+        )
+
+    def get_project_id_by_name(self, project_name: str) -> str:
+        """Get the ID of a log project by its name.
+
+        Args:
+            project_name (str): The name of the log project.
+
+        Returns:
+            str: The ID of the log project, or None if not found.
+        """
+        logger.info(f"Getting ID for log project '{project_name}' in TLS...")
+
+        request_body = {
+            "ProjectName": project_name,
+            "IsFullName": True,
+        }
+
+        try:
+            res = None
+            res = self._ve_tls_request(
+                client=self._client,
+                api="DescribeProjects",
+                body=request_body,
+                method="GET",
+            )
+            projects = res["Projects"]
+            for project in projects:
+                if project["ProjectName"] == project_name:
+                    return project["ProjectId"]
+            return "<no_project_id_found>"
+        except KeyError:
+            raise ValueError(f"Failed to get log project ID: {res}")
+
+    def create_log_project(self, project_name: str) -> str:
+        """Create a log project in TLS.
+
+        Args:
+            project_name (str): The name of the log project to create.
+
+        Returns:
+            str: The ID of the created log project.
+        """
+        logger.info(f"Creating log project '{project_name}' in TLS...")
+
+        request_body = {
+            "ProjectName": project_name,
+            "Region": self.region,
+            "Description": "Created by Volcengine Agent Development Kit (VeADK)",
+            "Tags": [{"Key": "provider", "Value": "VeADK"}],
+        }
+        try:
+            res = self._ve_tls_request(
+                client=self._client, api="CreateProject", body=request_body
+            )
+
+            if res["ErrorCode"] == "ProjectAlreadyExist":
+                logger.debug(
+                    f"Log project '{project_name}' already exists. Check its ID."
+                )
+                return self.get_project_id_by_name(project_name)
+
+            return res["ProjectId"]
+        except KeyError:
+            raise ValueError(f"Failed to create log project: {res}")
+
+    def get_trace_instance_by_name(self, log_project_id: str, trace_instance_name: str):
+        logger.info(f"Getting trace instance '{trace_instance_name}' in TLS...")
+
+        request_body = {
+            "PageSize": 100,
+            "ProjectId": log_project_id,
+            "TraceInstanceName": trace_instance_name,
+        }
+        try:
+            res = self._ve_tls_request(
+                client=self._client,
+                api="DescribeTraceInstances",
+                body=request_body,
+                method="GET",
+            )
+
+            for instance in res["TraceInstances"]:
+                if instance["TraceInstanceName"] == trace_instance_name:
+                    return instance
+        except KeyError:
+            raise ValueError(f"Failed to create log project: {res}")
+
+    def create_tracing_instance(self, log_project_id: str, trace_instance_name: str):
+        """Create a tracing instance in TLS.
+
+        Args:
+            instance_name (str): The name of the tracing instance to create.
+
+        Returns:
+            dict: The tracing instance.
+        """
+        logger.info(f"Creating tracing instance '{trace_instance_name}' in TLS...")
+
+        request_body = {
+            "ProjectId": log_project_id,
+            "TraceInstanceName": trace_instance_name,
+            "Description": "Created by Volcengine Agent Development Kit (VeADK)",
+        }
+
+        try:
+            res = None
+            res = self._ve_tls_request(
+                client=self._client,
+                api="CreateTraceInstance",
+                body=request_body,
+            )
+
+            if res["ErrorCode"] == "TopicAlreadyExist":
+                logger.debug(
+                    f"Log project '{trace_instance_name}' already exists. Check its ID."
+                )
+                return self.get_trace_instance_by_name(
+                    log_project_id, trace_instance_name
+                )
+
+            # after creation, get the trace instance details
+            res = self._ve_tls_request(
+                client=self._client,
+                api="DescribeTraceInstance",
+                body={"TraceInstanceID": res["TraceInstanceID"]},
+                method="GET",
+            )
+
+            return res
+        except KeyError:
+            raise ValueError(f"Failed to create tracing instance: {res}")
+
+    def get_trace_topic_id(self):
+        """Get the trace topic ID under VeADK default names.
+
+        This method is a tool function just designed for `veadk/config.py`.
+
+        Returns:
+            str: The trace topic ID.
+        """
+        logger.info("Getting trace topic ID for tracing instance in TLS...")
+
+        log_project_id = self.create_log_project(DEFAULT_TLS_LOG_PROJECT_NAME)
+
+        instance = self.create_tracing_instance(
+            log_project_id, DEFAULT_TLS_TRACING_INSTANCE_NAME
+        )
+
+        if not instance:
+            raise ValueError("None instance")
+
+        logger.info(f"Fetched trace topic id: {instance['TraceTopicId']}")
+
+        return instance["TraceTopicId"]

veadk/integrations/ve_tos/ve_tos.py

@@ -12,75 +12,73 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
-from veadk.config import getenv
-from veadk.utils.logger import get_logger
 import asyncio
-from typing import Union
-from pydantic import BaseModel, Field
-from typing import Any
-from urllib.parse import urlparse
+import os
 from datetime import datetime
+from typing import TYPE_CHECKING, Union
+from urllib.parse import urlparse
+
+from veadk.consts import DEFAULT_TOS_BUCKET_NAME
+from veadk.utils.logger import get_logger
+from veadk.utils.misc import getenv
+
+if TYPE_CHECKING:
+    pass
+
 
 # Initialize logger before using it
 logger = get_logger(__name__)
 
-# Try to import tos module, and provide helpful error message if it fails
-try:
-    import tos
-except ImportError as e:
-    logger.error(
-        "Failed to import 'tos' module. Please install it using: pip install tos\n"
-    )
-    raise ImportError(
-        "Missing 'tos' module. Please install it using: pip install tos\n"
-    ) from e
-
-
-class TOSConfig(BaseModel):
-    region: str = Field(
-        default_factory=lambda: getenv("DATABASE_TOS_REGION"),
-        description="TOS region",
-    )
-    ak: str = Field(
-        default_factory=lambda: getenv("VOLCENGINE_ACCESS_KEY"),
-        description="Volcengine access key",
-    )
-    sk: str = Field(
-        default_factory=lambda: getenv("VOLCENGINE_SECRET_KEY"),
-        description="Volcengine secret key",
-    )
-    bucket_name: str = Field(
-        default_factory=lambda: getenv("DATABASE_TOS_BUCKET"),
-        description="TOS bucket name",
-    )
-
-
-class VeTOS(BaseModel):
-    config: TOSConfig = Field(default_factory=TOSConfig)
-
-    def model_post_init(self, __context: Any) -> None:
+
+class VeTOS:
+    def __init__(
+        self,
+        ak: str = "",
+        sk: str = "",
+        region: str = "cn-beijing",
+        bucket_name: str = DEFAULT_TOS_BUCKET_NAME,
+    ) -> None:
+        self.ak = ak if ak else os.getenv("VOLCENGINE_ACCESS_KEY", "")
+        self.sk = sk if sk else os.getenv("VOLCENGINE_SECRET_KEY", "")
+        self.region = region
+        self.bucket_name = (
+            bucket_name if bucket_name else getenv("", DEFAULT_TOS_BUCKET_NAME)
+        )
+        self._tos_module = None
+
+        try:
+            import tos
+
+            self._tos_module = tos
+        except ImportError as e:
+            logger.error(
+                "Failed to import 'tos' module. Please install it using: pip install tos\n"
+            )
+            raise ImportError(
+                "Missing 'tos' module. Please install it using: pip install tos\n"
+            ) from e
+
+        self._client = None
         try:
-            self._client = tos.TosClientV2(
-                self.config.ak,
-                self.config.sk,
-                endpoint=f"tos-{self.config.region}.volces.com",
-                region=self.config.region,
+            self._client = self._tos_module.TosClientV2(
+                ak=self.ak,
+                sk=self.sk,
+                endpoint=f"tos-{self.region}.volces.com",
+                region=self.region,
             )
-            logger.info("Connected to TOS successfully.")
+            logger.info("Init TOS client.")
         except Exception as e:
             logger.error(f"Client initialization failed:{e}")
-            self._client = None
 
     def _refresh_client(self):
         try:
             if self._client:
                 self._client.close()
-            self._client = tos.TosClientV2(
-                self.config.ak,
-                self.config.sk,
-                endpoint=f"tos-{self.config.region}.volces.com",
-                region=self.config.region,
+            self._client = self._tos_module.TosClientV2(
+                self.ak,
+                self.sk,
+                endpoint=f"tos-{self.region}.volces.com",
+                region=self.region,
             )
             logger.info("refreshed client successfully.")
         except Exception as e:
@@ -93,19 +91,17 @@ class VeTOS(BaseModel):
             logger.error("TOS client is not initialized")
             return False
         try:
-            self._client.head_bucket(self.config.bucket_name)
-            logger.info(f"Bucket {self.config.bucket_name} already exists")
-        except tos.exceptions.TosServerError as e:
+            self._client.head_bucket(self.bucket_name)
+            logger.info(f"Bucket {self.bucket_name} already exists")
+        except self._tos_module.exceptions.TosServerError as e:
             if e.status_code == 404:
                 try:
                     self._client.create_bucket(
-                        bucket=self.config.bucket_name,
-                        storage_class=tos.StorageClassType.Storage_Class_Standard,
-                        acl=tos.ACLType.ACL_Public_Read,  # 公开读
-                    )
-                    logger.info(
-                        f"Bucket {self.config.bucket_name} created successfully"
+                        bucket=self.bucket_name,
+                        storage_class=self._tos_module.StorageClassType.Storage_Class_Standard,
+                        acl=self._tos_module.ACLType.ACL_Public_Read,
                     )
+                    logger.info(f"Bucket {self.bucket_name} created successfully")
                     self._refresh_client()
                 except Exception as create_error:
                     logger.error(f"Bucket creation failed: {str(create_error)}")
@@ -117,7 +113,7 @@ class VeTOS(BaseModel):
             logger.error(f"Bucket check failed: {str(e)}")
             return False
 
-        # 确保在所有路径上返回布尔值
+        # ensure return bool type
         return self._set_cors_rules()
 
     def _set_cors_rules(self) -> bool:
@@ -125,20 +121,18 @@ class VeTOS(BaseModel):
             logger.error("TOS client is not initialized")
             return False
         try:
-            rule = tos.models2.CORSRule(
+            rule = self._tos_module.models2.CORSRule(
                 allowed_origins=["*"],
                 allowed_methods=["GET", "HEAD"],
                 allowed_headers=["*"],
                 max_age_seconds=1000,
             )
-            self._client.put_bucket_cors(self.config.bucket_name, [rule])
-            logger.info(
-                f"CORS rules for bucket {self.config.bucket_name} set successfully"
-            )
+            self._client.put_bucket_cors(self.bucket_name, [rule])
+            logger.info(f"CORS rules for bucket {self.bucket_name} set successfully")
             return True
         except Exception as e:
             logger.error(
-                f"Failed to set CORS rules for bucket {self.config.bucket_name}: {str(e)}"
+                f"Failed to set CORS rules for bucket {self.bucket_name}: {str(e)}"
             )
             return False
 
@@ -155,7 +149,9 @@ class VeTOS(BaseModel):
 
         timestamp: str = datetime.now().strftime("%Y%m%d%H%M%S%f")[:-3]
         object_key: str = f"{app_name}-{user_id}-{session_id}/{timestamp}-{file_name}"
-        tos_url: str = f"https://{self.config.bucket_name}.tos-{self.config.region}.volces.com/{object_key}"
+        tos_url: str = (
+            f"https://{self.bucket_name}.tos-{self.region}.volces.com/{object_key}"
+        )
 
         return object_key, tos_url
 
@@ -182,9 +178,9 @@ class VeTOS(BaseModel):
             if not self.create_bucket():
                 return
             self._client.put_object(
-                bucket=self.config.bucket_name, key=object_key, content=data
+                bucket=self.bucket_name, key=object_key, content=data
             )
-            logger.debug(f"Upload success, object_key: {object_key}")
+            logger.debug(f"Upload success, url: {object_key}")
             self._close()
             return
         except Exception as e:
@@ -199,7 +195,7 @@ class VeTOS(BaseModel):
             if not self.create_bucket():
                 return
             self._client.put_object_from_file(
-                bucket=self.config.bucket_name, key=object_key, file_path=file_path
+                bucket=self.bucket_name, key=object_key, file_path=file_path
             )
             self._close()
             logger.debug(f"Upload success, object_key: {object_key}")
@@ -215,7 +211,7 @@ class VeTOS(BaseModel):
             logger.error("TOS client is not initialized")
             return False
         try:
-            object_stream = self._client.get_object(self.config.bucket_name, object_key)
+            object_stream = self._client.get_object(self.bucket_name, object_key)
 
             save_dir = os.path.dirname(save_path)
             if save_dir and not os.path.exists(save_dir):

veadk/knowledgebase/backends/__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/knowledgebase/backends/base_backend.py

@@ -0,0 +1,59 @@
+# 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.
+
+from abc import ABC, abstractmethod
+
+from pydantic import BaseModel
+
+
+class BaseKnowledgebaseBackend(ABC, BaseModel):
+    index: str
+    """Index or collection name of the vector storage."""
+
+    @abstractmethod
+    def precheck_index_naming(self) -> None:
+        """Check the index name is valid or not.
+
+        If index naming is not valid, raise an exception.
+        """
+
+    @abstractmethod
+    def add_from_directory(self, directory: str, **kwargs) -> bool:
+        """Add knowledge from file path to knowledgebase"""
+
+    @abstractmethod
+    def add_from_files(self, files: list[str], **kwargs) -> bool:
+        """Add knowledge (e.g, documents, strings, ...) to knowledgebase"""
+
+    @abstractmethod
+    def add_from_text(self, text: str | list[str], **kwargs) -> bool:
+        """Add knowledge from text to knowledgebase"""
+
+    @abstractmethod
+    def search(self, **kwargs) -> list:
+        """Search knowledge from knowledgebase"""
+
+    # Optional methods for future use:
+    # - `delete`: Delete collection or documents
+    # - `list_docs`: List original documents
+    # - `list_chunks`: List embedded document chunks
+
+    # def delete(self, **kwargs) -> bool:
+    #     """Delete knowledge from knowledgebase"""
+
+    # def list_docs(self, **kwargs) -> None:
+    #     """List original documents in knowledgebase"""
+
+    # def list_chunks(self, **kwargs) -> None:
+    #     """List embeded document chunks in knowledgebase"""

veadk/knowledgebase/backends/in_memory_backend.py

@@ -0,0 +1,82 @@
+# 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.
+
+from llama_index.core import Document, SimpleDirectoryReader, VectorStoreIndex
+from llama_index.core.schema import BaseNode
+from llama_index.embeddings.openai_like import OpenAILikeEmbedding
+from pydantic import Field
+from typing_extensions import Any, override
+
+from veadk.configs.model_configs import EmbeddingModelConfig, NormalEmbeddingModelConfig
+from veadk.knowledgebase.backends.base_backend import BaseKnowledgebaseBackend
+from veadk.knowledgebase.backends.utils import get_llama_index_splitter
+
+
+class InMemoryKnowledgeBackend(BaseKnowledgebaseBackend):
+    embedding_config: NormalEmbeddingModelConfig | EmbeddingModelConfig = Field(
+        default_factory=EmbeddingModelConfig
+    )
+    """Embedding model configs"""
+
+    def model_post_init(self, __context: Any) -> None:
+        self._embed_model = OpenAILikeEmbedding(
+            model_name=self.embedding_config.name,
+            api_key=self.embedding_config.api_key,
+            api_base=self.embedding_config.api_base,
+        )
+        self._vector_index = VectorStoreIndex([], embed_model=self._embed_model)
+
+    @override
+    def precheck_index_naming(self) -> None:
+        # Checking is not needed
+        pass
+
+    @override
+    def add_from_directory(self, directory: str) -> bool:
+        documents = SimpleDirectoryReader(input_dir=directory).load_data()
+        nodes = self._split_documents(documents)
+        self._vector_index.insert_nodes(nodes)
+        return True
+
+    @override
+    def add_from_files(self, files: list[str]) -> bool:
+        documents = SimpleDirectoryReader(input_files=files).load_data()
+        nodes = self._split_documents(documents)
+        self._vector_index.insert_nodes(nodes)
+        return True
+
+    @override
+    def add_from_text(self, text: str | list[str]) -> bool:
+        if isinstance(text, str):
+            documents = [Document(text=text)]
+        else:
+            documents = [Document(text=t) for t in text]
+        nodes = self._split_documents(documents)
+        self._vector_index.insert_nodes(nodes)
+        return True
+
+    @override
+    def search(self, query: str, top_k: int = 5) -> list[str]:
+        _retriever = self._vector_index.as_retriever(similarity_top_k=top_k)
+        retrieved_nodes = _retriever.retrieve(query)
+        return [node.text for node in retrieved_nodes]
+
+    def _split_documents(self, documents: list[Document]) -> list[BaseNode]:
+        """Split document into chunks"""
+        nodes = []
+        for document in documents:
+            splitter = get_llama_index_splitter(document.metadata.get("file_path", ""))
+            _nodes = splitter.get_nodes_from_documents([document])
+            nodes.extend(_nodes)
+        return nodes