veadk-python 0.2.5__py3-none-any.whl → 0.2.7__py3-none-any.whl

veadk/integrations/ve_tos/ve_tos.py

@@ -12,72 +12,125 @@
 # 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 tos
 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
+
+if TYPE_CHECKING:
+    pass
+
 
+# Initialize logger before using it
 logger = get_logger(__name__)
 
 
-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
+        self._tos_module = None
+
         try:
-            self._client = tos.TosClientV2(
-                self.config.ak,
-                self.config.sk,
-                endpoint=f"tos-{self.config.region}.volces.com",
-                region=self.config.region,
+            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"
             )
-            logger.info("Connected to TOS successfully.")
+            raise ImportError(
+                "Missing 'tos' module. Please install it using: pip install tos\n"
+            ) from e
+
+        self._client = None
+        try:
+            self._client = self._tos_module.TosClientV2(
+                ak=self.ak,
+                sk=self.sk,
+                endpoint=f"tos-{self.region}.volces.com",
+                region=self.region,
+            )
+            logger.info("Init TOS client.")
         except Exception as e:
             logger.error(f"Client initialization failed:{e}")
-            return None
+
+    def _refresh_client(self):
+        try:
+            if self._client:
+                self._client.close()
+            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:
+            logger.error(f"Failed to refresh client: {str(e)}")
+            self._client = None
 
     def create_bucket(self) -> bool:
-        """If the bucket does not exist, create it"""
+        """If the bucket does not exist, create it and set CORS rules"""
+        if not self._client:
+            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")
-            return True
-        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:
-                self._client.create_bucket(
-                    bucket=self.config.bucket_name,
-                    storage_class=tos.StorageClassType.Storage_Class_Standard,
-                    acl=tos.ACLType.ACL_Private,
-                )
-                logger.info(f"Bucket {self.config.bucket_name} created successfully")
-                return True
+                try:
+                    self._client.create_bucket(
+                        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)}")
+                    return False
+            else:
+                logger.error(f"Bucket check failed: {str(e)}")
+                return False
+        except Exception as e:
+            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:
+        if not self._client:
+            logger.error("TOS client is not initialized")
+            return False
+        try:
+            rule = self._tos_module.models2.CORSRule(
+                allowed_origins=["*"],
+                allowed_methods=["GET", "HEAD"],
+                allowed_headers=["*"],
+                max_age_seconds=1000,
+            )
+            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"Bucket creation failed: {str(e)}")
+            logger.error(
+                f"Failed to set CORS rules for bucket {self.bucket_name}: {str(e)}"
+            )
             return False
 
     def build_tos_url(
@@ -93,7 +146,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
 
@@ -103,57 +158,57 @@ class VeTOS(BaseModel):
         data: Union[str, bytes],
     ):
         if isinstance(data, str):
-            data_type = "file"
+            # data is a file path
+            return asyncio.to_thread(self._do_upload_file, object_key, data)
         elif isinstance(data, bytes):
-            data_type = "bytes"
+            # data is bytes content
+            return asyncio.to_thread(self._do_upload_bytes, object_key, data)
         else:
             error_msg = f"Upload failed: data type error. Only str (file path) and bytes are supported, got {type(data)}"
             logger.error(error_msg)
             raise ValueError(error_msg)
-        if data_type == "file":
-            return asyncio.to_thread(self._do_upload_file, object_key, data)
-        elif data_type == "bytes":
-            return asyncio.to_thread(self._do_upload_bytes, object_key, data)
 
-    def _do_upload_bytes(self, object_key: str, bytes: bytes) -> bool:
+    def _do_upload_bytes(self, object_key: str, data: bytes) -> None:
         try:
             if not self._client:
-                return False
+                return
             if not self.create_bucket():
-                return False
+                return
             self._client.put_object(
-                bucket=self.config.bucket_name, key=object_key, content=bytes
+                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 True
+            return
         except Exception as e:
             logger.error(f"Upload failed: {e}")
             self._close()
-            return False
+            return
 
-    def _do_upload_file(self, object_key: str, file_path: str) -> bool:
+    def _do_upload_file(self, object_key: str, file_path: str) -> None:
         try:
             if not self._client:
-                return False
+                return
             if not self.create_bucket():
-                return False
-
+                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}")
-            return True
+            return
         except Exception as e:
             logger.error(f"Upload failed: {e}")
             self._close()
-            return False
+            return
 
     def download(self, object_key: str, save_path: str) -> bool:
         """download image from TOS"""
+        if not self._client:
+            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/knowledgebase.py

@@ -11,11 +11,15 @@
 # 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 io
+import os.path
+from typing import Any, BinaryIO, Literal, TextIO
 
-from typing import BinaryIO, Literal, TextIO
+from pydantic import BaseModel
 
 from veadk.database.database_adapter import get_knowledgebase_database_adapter
 from veadk.database.database_factory import DatabaseFactory
+from veadk.utils.misc import formatted_timestamp
 from veadk.utils.logger import get_logger
 
 logger = get_logger(__name__)
@@ -25,23 +29,23 @@ def build_knowledgebase_index(app_name: str):
     return f"veadk_kb_{app_name}"
 
 
-class KnowledgeBase:
-    def __init__(
-        self,
-        backend: Literal["local", "opensearch", "viking", "redis", "mysql"] = "local",
-        top_k: int = 10,
-        db_config=None,
-    ):
-        logger.info(f"Initializing knowledgebase: backend={backend} top_k={top_k}")
+class KnowledgeBase(BaseModel):
+    backend: Literal["local", "opensearch", "viking", "redis", "mysql"] = "local"
+    top_k: int = 10
+    db_config: Any | None = None
 
-        self.backend = backend
-        self.top_k = top_k
+    def model_post_init(self, __context: Any) -> None:
+        logger.info(
+            f"Initializing knowledgebase: backend={self.backend} top_k={self.top_k}"
+        )
 
-        self.db_client = DatabaseFactory.create(backend=backend, config=db_config)
-        self.adapter = get_knowledgebase_database_adapter(self.db_client)
+        self._db_client = DatabaseFactory.create(
+            backend=self.backend, config=self.db_config
+        )
+        self._adapter = get_knowledgebase_database_adapter(self._db_client)
 
         logger.info(
-            f"Initialized knowledgebase: db_client={self.db_client.__class__.__name__} adapter={self.adapter}"
+            f"Initialized knowledgebase: db_client={self._db_client.__class__.__name__} adapter={self._adapter}"
         )
 
     def add(
@@ -52,9 +56,16 @@ class KnowledgeBase:
     ):
         """
         Add documents to the vector database.
-        You can only upload files or file characters when the adapter type used is vikingdb.
-        In addition, if you upload data of the bytes type,
-            for example, if you read the file stream of a pdf, then you need to pass an additional parameter file_ext = '.pdf'.
+        Args:
+            data (str | list[str] | TextIO | BinaryIO | bytes): The data to be added.
+                - str: A single file path. (viking only)
+                - list[str]: A list of file paths.
+                - TextIO: A file object (TextIO). (viking only) file descriptor
+                - BinaryIO: A file object (BinaryIO). (viking only) file descriptor
+                - bytes: Binary data. (viking only) binary data (f.read())
+            app_name: index name
+            **kwargs: Additional keyword arguments.
+                - file_name (str | list[str]): The file name or a list of file names (including suffix). (viking only)
         """
         if self.backend != "viking" and not (
             isinstance(data, str) or isinstance(data, list)
@@ -64,10 +75,68 @@ class KnowledgeBase:
             )
 
         index = build_knowledgebase_index(app_name)
-
         logger.info(f"Adding documents to knowledgebase: index={index}")
 
-        self.adapter.add(data=data, index=index)
+        if self.backend == "viking":
+            # Case 1: Handling file paths or lists of file paths (str)
+            if isinstance(data, str) and os.path.isfile(data):
+                # Get the file name (including the suffix)
+                if "file_name" not in kwargs or not kwargs["file_name"]:
+                    kwargs["file_name"] = os.path.basename(data)
+                return self._adapter.add(data=data, index=index, **kwargs)
+            # Case 2: Handling when list[str] is a full path  (list[str])
+            if isinstance(data, list):
+                if all(isinstance(item, str) for item in data):
+                    all_paths = all(os.path.isfile(item) for item in data)
+                    all_not_paths = all(not os.path.isfile(item) for item in data)
+                    if all_paths:
+                        if "file_name" not in kwargs or not kwargs["file_name"]:
+                            kwargs["file_name"] = [
+                                os.path.basename(item) for item in data
+                            ]
+                        return self._adapter.add(data=data, index=index, **kwargs)
+                    elif (
+                        not all_not_paths
+                    ):  # Prevent the occurrence of non-existent paths
+                        # There is a mixture of paths and non-paths
+                        raise ValueError(
+                            "Mixed file paths and content strings in list are not allowed"
+                        )
+            # Case 3: Handling strings or string arrays (content)  (str or list[str])
+            if isinstance(data, str) or (
+                isinstance(data, list) and all(isinstance(item, str) for item in data)
+            ):
+                if "file_name" not in kwargs or not kwargs["file_name"]:
+                    if isinstance(data, str):
+                        kwargs["file_name"] = f"{formatted_timestamp()}.txt"
+                    else:  # list[str] without file_names
+                        prefix_file_name = formatted_timestamp()
+                        kwargs["file_name"] = [
+                            f"{prefix_file_name}_{i}.txt" for i in range(len(data))
+                        ]
+                return self._adapter.add(data=data, index=index, **kwargs)
+
+            # Case 4: Handling binary data (bytes)
+            if isinstance(data, bytes):
+                # user must give file_name
+                if "file_name" not in kwargs:
+                    raise ValueError("file_name must be provided for binary data")
+                return self._adapter.add(data=data, index=index, **kwargs)
+
+            # Case 5: Handling file objects TextIO or BinaryIO
+            if isinstance(data, (io.TextIOWrapper, io.BufferedReader)):
+                if not kwargs.get("file_name") and hasattr(data, "name"):
+                    kwargs["file_name"] = os.path.basename(data.name)
+                return self._adapter.add(data=data, index=index, **kwargs)
+            # Case6: Unsupported data type
+            raise TypeError(f"Unsupported data type: {type(data)}")
+
+        if not isinstance(data, list):
+            raise TypeError(
+                f"Unsupported data type: {type(data)}. Only viking support file_path and file bytes"
+            )
+        # not viking
+        return self._adapter.add(data=data, index=index, **kwargs)
 
     def search(self, query: str, app_name: str, top_k: int | None = None) -> list[str]:
         top_k = self.top_k if top_k is None else top_k
@@ -76,7 +145,34 @@ class KnowledgeBase:
             f"Searching knowledgebase: app_name={app_name} query={query} top_k={top_k}"
         )
         index = build_knowledgebase_index(app_name)
-        result = self.adapter.query(query=query, index=index, top_k=top_k)
+        result = self._adapter.query(query=query, index=index, top_k=top_k)
         if len(result) == 0:
             logger.warning(f"No documents found in knowledgebase. Query: {query}")
         return result
+
+    def delete(self, app_name: str) -> bool:
+        index = build_knowledgebase_index(app_name)
+        return self._adapter.delete(index=index)
+
+    def delete_doc(self, app_name: str, id: str) -> bool:
+        index = build_knowledgebase_index(app_name)
+        return self._adapter.delete_doc(index=index, id=id)
+
+    def list_chunks(
+        self, app_name: str, offset: int = 0, limit: int = 100
+    ) -> list[dict]:
+        index = build_knowledgebase_index(app_name)
+        return self._adapter.list_chunks(index=index, offset=offset, limit=limit)
+
+    def list_docs(self, app_name: str, offset: int = 0, limit: int = 100) -> list[dict]:
+        if self.backend == "viking":
+            index = build_knowledgebase_index(app_name)
+            return self._adapter.list_docs(index=index, offset=offset, limit=limit)
+        else:
+            raise NotImplementedError(
+                f"list_docs not supported for {self.backend}, only viking support list_docs"
+            )
+
+    def exists(self, app_name: str) -> bool:
+        index = build_knowledgebase_index(app_name)
+        return self._adapter.index_exists(index=index)

veadk/memory/long_term_memory.py

@@ -13,8 +13,9 @@
 # limitations under the License.
 
 # adapted from Google ADK memory service adk-python/src/google/adk/memory/vertex_ai_memory_bank_service.py at 0a9e67dbca67789247e882d16b139dbdc76a329a · google/adk-python
+
 import json
-from typing import Literal
+from typing import Any, Literal
 
 from google.adk.events.event import Event
 from google.adk.memory.base_memory_service import (
@@ -24,6 +25,7 @@ from google.adk.memory.base_memory_service import (
 from google.adk.memory.memory_entry import MemoryEntry
 from google.adk.sessions import Session
 from google.genai import types
+from pydantic import BaseModel
 from typing_extensions import override
 
 from veadk.database import DatabaseFactory
@@ -37,33 +39,30 @@ def build_long_term_memory_index(app_name: str, user_id: str):
     return f"{app_name}_{user_id}"
 
 
-class LongTermMemory(BaseMemoryService):
-    def __init__(
-        self,
-        backend: Literal[
-            "local", "opensearch", "redis", "mysql", "viking", "viking_mem"
-        ] = "opensearch",
-        top_k: int = 5,
-    ):
-        if backend == "viking":
+class LongTermMemory(BaseMemoryService, BaseModel):
+    backend: Literal[
+        "local", "opensearch", "redis", "mysql", "viking", "viking_mem"
+    ] = "opensearch"
+    top_k: int = 5
+
+    def model_post_init(self, __context: Any) -> None:
+        if self.backend == "viking":
             logger.warning(
                 "`viking` backend is deprecated, switching to `viking_mem` backend."
             )
-            backend = "viking_mem"
-        self.top_k = top_k
-        self.backend = backend
+            self.backend = "viking_mem"
 
         logger.info(
             f"Initializing long term memory: backend={self.backend} top_k={self.top_k}"
         )
 
-        self.db_client = DatabaseFactory.create(
-            backend=backend,
+        self._db_client = DatabaseFactory.create(
+            backend=self.backend,
         )
-        self.adapter = get_long_term_memory_database_adapter(self.db_client)
+        self._adapter = get_long_term_memory_database_adapter(self._db_client)
 
         logger.info(
-            f"Initialized long term memory: db_client={self.db_client.__class__.__name__} adapter={self.adapter}"
+            f"Initialized long term memory: db_client={self._db_client.__class__.__name__} adapter={self._adapter}"
         )
 
     def _filter_and_convert_events(self, events: list[Event]) -> list[str]:
@@ -101,9 +100,9 @@ class LongTermMemory(BaseMemoryService):
 
         # check if viking memory database, should give a user id： if/else
         if self.backend == "viking_mem":
-            self.adapter.add(data=event_strings, index=index, user_id=session.user_id)
+            self._adapter.add(data=event_strings, index=index, user_id=session.user_id)
         else:
-            self.adapter.add(data=event_strings, index=index)
+            self._adapter.add(data=event_strings, index=index)
 
         logger.info(
             f"Added {len(event_strings)} events to long term memory: index={index}"
@@ -119,11 +118,11 @@ class LongTermMemory(BaseMemoryService):
 
         # user id if viking memory db
         if self.backend == "viking_mem":
-            memory_chunks = self.adapter.query(
+            memory_chunks = self._adapter.query(
                 query=query, index=index, top_k=self.top_k, user_id=user_id
             )
         else:
-            memory_chunks = self.adapter.query(
+            memory_chunks = self._adapter.query(
                 query=query, index=index, top_k=self.top_k
             )
 

veadk/memory/short_term_memory_processor.py

@@ -20,7 +20,12 @@ from google.adk.sessions import Session
 from google.genai.types import Content, Part
 from litellm import completion
 
-from veadk.config import getenv
+from veadk.config import settings
+from veadk.consts import (
+    DEFAULT_MODEL_AGENT_API_BASE,
+    DEFAULT_MODEL_AGENT_NAME,
+    DEFAULT_MODEL_AGENT_PROVIDER,
+)
 from veadk.prompts.prompt_memory_processor import render_prompt
 from veadk.utils.logger import get_logger
 
@@ -62,9 +67,9 @@ class ShortTermMemoryProcessor:
         prompt = render_prompt(messages=messages)
 
         res = completion(
-            model=getenv("MODEL_AGENT_PROVIDER") + "/" + getenv("MODEL_AGENT_NAME"),
-            base_url=getenv("MODEL_AGENT_API_BASE"),
-            api_key=getenv("MODEL_AGENT_API_KEY"),
+            model=DEFAULT_MODEL_AGENT_PROVIDER + "/" + DEFAULT_MODEL_AGENT_NAME,
+            base_url=DEFAULT_MODEL_AGENT_API_BASE,
+            api_key=settings.model.api_key,
             messages=[
                 {
                     "role": "user",