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

veadk/consts.py

@@ -12,12 +12,32 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import time
+
+from veadk.utils.misc import getenv
 from veadk.version import VERSION
 
+DEFAULT_AGENT_NAME = "veAgent"
+
 DEFAULT_MODEL_AGENT_NAME = "doubao-seed-1-6-250615"
 DEFAULT_MODEL_AGENT_PROVIDER = "openai"
 DEFAULT_MODEL_AGENT_API_BASE = "https://ark.cn-beijing.volces.com/api/v3/"
-DEFAULT_MODEL_EXTRA_HEADERS = {"veadk-source": "veadk", "veadk-version": VERSION}
+DEFAULT_MODEL_EXTRA_CONFIG = {
+    "extra_headers": {
+        "x-is-encrypted": getenv("MODEL_AGENT_ENCRYPTED", "true"),
+        "veadk-source": "veadk",
+        "veadk-version": VERSION,
+    },
+    "extra_body": {
+        "caching": {
+            "type": getenv("MODEL_AGENT_CACHING", "enabled"),
+        },
+        "thinking": {
+            "type": "disabled"
+        },  # disable thinking for token saving and shorter TTFT by default
+        "expire_at": int(time.time()) + 3600,  # expire after 1 hour
+    },
+}
 
 DEFAULT_APMPLUS_OTEL_EXPORTER_ENDPOINT = "http://apmplus-cn-beijing.volces.com:4317"
 DEFAULT_APMPLUS_OTEL_EXPORTER_SERVICE_NAME = "veadk_tracing"
@@ -28,3 +48,14 @@ DEFAULT_COZELOOP_OTEL_EXPORTER_ENDPOINT = (
 
 DEFAULT_TLS_OTEL_EXPORTER_ENDPOINT = "https://tls-cn-beijing.volces.com:4318/v1/traces"
 DEFAULT_TLS_OTEL_EXPORTER_REGION = "cn-beijing"
+
+DEFAULT_CR_INSTANCE_NAME = "veadk-user-instance"
+DEFAULT_CR_NAMESPACE_NAME = "veadk-user-namespace"
+DEFAULT_CR_REPO_NAME = "veadk-user-repo"
+
+DEFAULT_TLS_LOG_PROJECT_NAME = "veadk-logs"
+DEFAULT_TLS_TRACING_INSTANCE_NAME = "veadk-tracing"
+
+DEFAULT_TOS_BUCKET_NAME = "veadk-default-bucket"
+
+DEFAULT_COZELOOP_SPACE_NAME = "VeADK Space"

veadk/database/database_adapter.py

@@ -28,7 +28,26 @@ class KVDatabaseAdapter:
 
         self.client: RedisDatabase = client
 
-    def add(self, data: list[str], index: str):
+    def index_exists(self, index: str) -> bool:
+        """
+        Check if the index (key) exists in Redis.
+
+        Args:
+            index: The Redis key to check
+
+        Returns:
+            bool: True if the key exists, False otherwise
+        """
+        try:
+            # Use Redis EXISTS command to check if key exists
+            return bool(self.client._client.exists(index))
+        except Exception as e:
+            logger.error(
+                f"Failed to check if key exists in Redis: index={index} error={e}"
+            )
+            return False
+
+    def add(self, data: list[str], index: str, **kwargs):
         logger.debug(f"Adding documents to Redis database: index={index}")
 
         try:
@@ -54,6 +73,44 @@ class KVDatabaseAdapter:
             logger.error(f"Failed to search from Redis: index={index} error={e}")
             raise e
 
+    def delete(self, index: str) -> bool:
+        logger.debug(f"Deleting key from Redis database: index={index}")
+        try:
+            self.client.delete(key=index)
+            return True
+        except Exception as e:
+            logger.error(
+                f"Failed to delete key from Redis database: index={index} error={e}"
+            )
+            return False
+
+    def delete_doc(self, index: str, id: str) -> bool:
+        logger.debug(f"Deleting document from Redis database: index={index} id={id}")
+        try:
+            # For Redis, we need to handle deletion differently since RedisDatabase.delete_doc
+            # takes a key and a single id
+            result = self.client.delete_doc(key=index, id=id)
+            return result
+        except Exception as e:
+            logger.error(
+                f"Failed to delete document from Redis database: index={index} id={id} error={e}"
+            )
+            return False
+
+    def list_chunks(self, index: str, offset: int = 0, limit: int = 100) -> list[dict]:
+        logger.debug(f"Listing documents from Redis database: index={index}")
+        try:
+            # Get all documents from Redis
+            docs = self.client.list_docs(key=index)
+
+            # Apply offset and limit for pagination
+            return docs[offset : offset + limit]
+        except Exception as e:
+            logger.error(
+                f"Failed to list documents from Redis database: index={index} error={e}"
+            )
+            return []
+
 
 class RelationalDatabaseAdapter:
     def __init__(self, client):
@@ -61,6 +118,24 @@ class RelationalDatabaseAdapter:
 
         self.client: MysqlDatabase = client
 
+    def index_exists(self, index: str) -> bool:
+        """
+        Check if the table (index) exists in MySQL database.
+
+        Args:
+            index: The table name to check
+
+        Returns:
+            bool: True if the table exists, False otherwise
+        """
+        try:
+            return self.client.table_exists(index)
+        except Exception as e:
+            logger.error(
+                f"Failed to check if table exists in MySQL: index={index} error={e}"
+            )
+            return False
+
     def create_table(self, table_name: str):
         logger.debug(f"Creating table for SQL database: table_name={table_name}")
 
@@ -73,7 +148,7 @@ class RelationalDatabaseAdapter:
         """
         self.client.add(sql)
 
-    def add(self, data: list[str], index: str):
+    def add(self, data: list[str], index: str, **kwargs):
         logger.debug(
             f"Adding documents to SQL database: table_name={index} data_len={len(data)}"
         )
@@ -108,6 +183,39 @@ class RelationalDatabaseAdapter:
 
         return [item["data"] for item in results]
 
+    def delete(self, index: str) -> bool:
+        logger.debug(f"Deleting table from SQL database: table_name={index}")
+        try:
+            self.client.delete(table=index)
+            return True
+        except Exception as e:
+            logger.error(
+                f"Failed to delete table from SQL database: table_name={index} error={e}"
+            )
+            return False
+
+    def delete_doc(self, index: str, id: str) -> bool:
+        logger.debug(f"Deleting document from SQL database: table_name={index} id={id}")
+        try:
+            # Convert single id to list for the client method
+            result = self.client.delete_doc(table=index, ids=[int(id)])
+            return result
+        except Exception as e:
+            logger.error(
+                f"Failed to delete document from SQL database: table_name={index} id={id} error={e}"
+            )
+            return False
+
+    def list_docs(self, index: str, offset: int = 0, limit: int = 100) -> list[dict]:
+        logger.debug(f"Listing documents from SQL database: table_name={index}")
+        try:
+            return self.client.list_docs(table=index, offset=offset, limit=limit)
+        except Exception as e:
+            logger.error(
+                f"Failed to list documents from SQL database: table_name={index} error={e}"
+            )
+            return []
+
 
 class VectorDatabaseAdapter:
     def __init__(self, client):
@@ -117,6 +225,25 @@ class VectorDatabaseAdapter:
 
         self.client: OpenSearchVectorDatabase = client
 
+    def index_exists(self, index: str) -> bool:
+        """
+        Check if the collection (index) exists in OpenSearch.
+
+        Args:
+            index: The collection name to check
+
+        Returns:
+            bool: True if the collection exists, False otherwise
+        """
+        try:
+            self._validate_index(index)
+            return self.client.collection_exists(index)
+        except Exception as e:
+            logger.error(
+                f"Failed to check if collection exists in OpenSearch: index={index} error={e}"
+            )
+            return False
+
     def _validate_index(self, index: str):
         """
         Verify whether the string conforms to the naming rules of index_name in OpenSearch.
@@ -132,7 +259,7 @@ class VectorDatabaseAdapter:
                 "The index name does not conform to the naming rules of OpenSearch"
             )
 
-    def add(self, data: list[str], index: str):
+    def add(self, data: list[str], index: str, **kwargs):
         self._validate_index(index)
 
         logger.debug(
@@ -152,6 +279,35 @@ class VectorDatabaseAdapter:
             top_k=top_k,
         )
 
+    def delete(self, index: str) -> bool:
+        self._validate_index(index)
+        logger.debug(f"Deleting collection from vector database: index={index}")
+        try:
+            self.client.delete(collection_name=index)
+            return True
+        except Exception as e:
+            logger.error(
+                f"Failed to delete collection from vector database: index={index} error={e}"
+            )
+            return False
+
+    def delete_doc(self, index: str, id: str) -> bool:
+        self._validate_index(index)
+        logger.debug(f"Deleting documents from vector database: index={index} id={id}")
+        try:
+            self.client.delete_by_id(collection_name=index, id=id)
+            return True
+        except Exception as e:
+            logger.error(
+                f"Failed to delete document from vector database: index={index} id={id} error={e}"
+            )
+            return False
+
+    def list_chunks(self, index: str, offset: int = 0, limit: int = 1000) -> list[dict]:
+        self._validate_index(index)
+        logger.debug(f"Listing documents from vector database: index={index}")
+        return self.client.list_docs(collection_name=index, offset=offset, limit=limit)
+
 
 class VikingDatabaseAdapter:
     def __init__(self, client):
@@ -159,6 +315,25 @@ class VikingDatabaseAdapter:
 
         self.client: VikingDatabase = client
 
+    def index_exists(self, index: str) -> bool:
+        """
+        Check if the collection (index) exists in VikingDB.
+
+        Args:
+            index: The collection name to check
+
+        Returns:
+            bool: True if the collection exists, False otherwise
+        """
+        try:
+            self._validate_index(index)
+            return self.client.collection_exists(index)
+        except Exception as e:
+            logger.error(
+                f"Failed to check if collection exists in VikingDB: index={index} error={e}"
+            )
+            return False
+
     def _validate_index(self, index: str):
         """
         Only English letters, numbers, and underscores (_) are allowed.
@@ -212,6 +387,28 @@ class VikingDatabaseAdapter:
 
         return self.client.query(query, collection_name=index, top_k=top_k)
 
+    def delete(self, index: str) -> bool:
+        self._validate_index(index)
+        logger.debug(f"Deleting collection from Viking database: index={index}")
+        return self.client.delete(name=index)
+
+    def delete_doc(self, index: str, id: str) -> bool:
+        self._validate_index(index)
+        logger.debug(f"Deleting documents from vector database: index={index} id={id}")
+        return self.client.delete_by_id(collection_name=index, id=id)
+
+    def list_chunks(self, index: str, offset: int, limit: int) -> list[dict]:
+        self._validate_index(index)
+        logger.debug(f"Listing documents from vector database: index={index}")
+        return self.client.list_chunks(
+            collection_name=index, offset=offset, limit=limit
+        )
+
+    def list_docs(self, index: str, offset: int, limit: int) -> list[dict]:
+        self._validate_index(index)
+        logger.debug(f"Listing documents from vector database: index={index}")
+        return self.client.list_docs(collection_name=index, offset=offset, limit=limit)
+
 
 class VikingMemoryDatabaseAdapter:
     def __init__(self, client):
@@ -219,6 +416,25 @@ class VikingMemoryDatabaseAdapter:
 
         self.client: VikingMemoryDatabase = client
 
+    def index_exists(self, index: str) -> bool:
+        """
+        Check if the collection (index) exists in VikingMemoryDB.
+
+        Note:
+            VikingMemoryDatabase does not support checking if a collection exists.
+            This method always returns False.
+
+        Args:
+            index: The collection name to check
+
+        Returns:
+            bool: Always returns False as VikingMemoryDatabase does not support this functionality
+        """
+        logger.warning(
+            "VikingMemoryDatabase does not support checking if a collection exists"
+        )
+        raise NotImplementedError("VikingMemoryDatabase does not support index_exists")
+
     def _validate_index(self, index: str):
         if not (
             isinstance(index, str)
@@ -248,6 +464,17 @@ class VikingMemoryDatabaseAdapter:
         result = self.client.query(query, collection_name=index, top_k=top_k, **kwargs)
         return result
 
+    def delete(self, index: str) -> bool:
+        self._validate_index(index)
+        logger.debug(f"Deleting collection from Viking database memory: index={index}")
+        raise NotImplementedError("VikingMemoryDatabase does not support delete")
+
+    def delete_docs(self, index: str, ids: list[int]):
+        raise NotImplementedError("VikingMemoryDatabase does not support delete_docs")
+
+    def list_chunks(self, index: str):
+        raise NotImplementedError("VikingMemoryDatabase does not support list_docs")
+
 
 class LocalDatabaseAdapter:
     def __init__(self, client):
@@ -255,12 +482,38 @@ class LocalDatabaseAdapter:
 
         self.client: LocalDataBase = client
 
+    def index_exists(self, index: str) -> bool:
+        """
+        Check if the index exists in LocalDataBase.
+
+        Note:
+            LocalDataBase does not support checking if an index exists.
+            This method always returns False.
+
+        Args:
+            index: The index name to check (not used in LocalDataBase)
+
+        Returns:
+            bool: Always returns False as LocalDataBase does not support this functionality
+        """
+        logger.warning("LocalDataBase does not support checking if an index exists")
+        return True
+
     def add(self, data: list[str], **kwargs):
         self.client.add(data)
 
     def query(self, query: str, **kwargs):
         return self.client.query(query, **kwargs)
 
+    def delete(self, index: str) -> bool:
+        return self.client.delete()
+
+    def delete_doc(self, index: str, id: str) -> bool:
+        return self.client.delete_doc(id)
+
+    def list_chunks(self, index: str, offset: int = 0, limit: int = 100) -> list[dict]:
+        return self.client.list_docs(offset=offset, limit=limit)
+
 
 MAPPING = {
     "RedisDatabase": KVDatabaseAdapter,

veadk/database/kv/redis_database.py

@@ -110,3 +110,50 @@ class RedisDatabase(BaseModel, BaseDatabase):
         except Exception as e:
             logger.error(f"Failed to delete key `{key}`: {e}")
             raise e
+
+    def delete_doc(self, key: str, id: str) -> bool:
+        """Delete a specific document by ID from a Redis list.
+
+        Args:
+            key: The Redis key (list) to delete from
+            id: The ID of the document to delete
+
+        Returns:
+            bool: True if deletion was successful, False otherwise
+        """
+        try:
+            # Get all items in the list
+            items = self._client.lrange(key, 0, -1)
+
+            # Find the index of the item to delete
+            for i, item in enumerate(items):
+                # Assuming the item is stored as a JSON string with an 'id' field
+                # If it's just the content, we'll use the list index as ID
+                if str(i) == id:
+                    self._client.lrem(key, 1, item)
+                    return True
+
+            logger.warning(f"Document with id {id} not found in key {key}")
+            return False
+        except Exception as e:
+            logger.error(f"Failed to delete document with id {id} from key {key}: {e}")
+            return False
+
+    def list_docs(self, key: str) -> list[dict]:
+        """List all documents in a Redis list.
+
+        Args:
+            key: The Redis key (list) to list documents from
+
+        Returns:
+            list[dict]: List of documents with id and content
+        """
+        try:
+            items = self._client.lrange(key, 0, -1)
+            return [
+                {"id": str(i), "content": item, "metadata": {}}
+                for i, item in enumerate(items)
+            ]
+        except Exception as e:
+            logger.error(f"Failed to list documents from key {key}: {e}")
+            return []

veadk/database/local_database.py

@@ -24,20 +24,39 @@ class LocalDataBase(BaseDatabase):
 
     def __init__(self, **kwargs):
         super().__init__()
-        self.data = []
+        self.data = {}
         self._type = "local"
+        self._next_id = 0  # Used to generate unique IDs
 
     def add_texts(self, texts: list[str], **kwargs):
-        self.data.extend(texts)
+        for text in texts:
+            self.data[str(self._next_id)] = text
+            self._next_id += 1
 
     def is_empty(self):
         return len(self.data) == 0
 
     def query(self, query: str, **kwargs: Any) -> list[str]:
-        return self.data
+        return list(self.data.values())
 
     def delete(self, **kwargs: Any):
-        self.data = []
+        self.data = {}
+        return True
 
     def add(self, texts: list[str], **kwargs: Any):
         return self.add_texts(texts)
+
+    def list_docs(self, **kwargs: Any) -> list[dict]:
+        return [
+            {"id": id, "content": content, "metadata": {}}
+            for id, content in self.data.items()
+        ]
+
+    def delete_doc(self, id: str, **kwargs: Any):
+        if id not in self.data:
+            raise ValueError(f"id {id} not found")
+        try:
+            del self.data[id]
+            return True
+        except Exception:
+            return False

veadk/database/relational/mysql_database.py

@@ -111,5 +111,63 @@ class MysqlDatabase(BaseModel, BaseDatabase):
             logger.error(f"Failed to drop table {table}: {e}")
             raise e
 
+    def delete_doc(self, table: str, ids: list[int]) -> bool:
+        """Delete documents by IDs from a MySQL table.
+
+        Args:
+            table: The table name to delete from
+            ids: List of document IDs to delete
+
+        Returns:
+            bool: True if deletion was successful, False otherwise
+        """
+        if not self.table_exists(table):
+            logger.warning(f"Table {table} does not exist. Skipping delete operation.")
+            return False
+
+        if not ids:
+            return True  # Nothing to delete
+
+        try:
+            with self._connection.cursor() as cursor:
+                # Create placeholders for the IDs
+                placeholders = ",".join(["%s"] * len(ids))
+                sql = f"DELETE FROM `{table}` WHERE id IN ({placeholders})"
+                cursor.execute(sql, ids)
+                self._connection.commit()
+                logger.info(f"Deleted {cursor.rowcount} documents from table {table}")
+                return True
+        except Exception as e:
+            logger.error(f"Failed to delete documents from table {table}: {e}")
+            return False
+
+    def list_docs(self, table: str, offset: int = 0, limit: int = 100) -> list[dict]:
+        """List documents from a MySQL table.
+
+        Args:
+            table: The table name to list documents from
+            offset: Offset for pagination
+            limit: Limit for pagination
+
+        Returns:
+            list[dict]: List of documents with id and content
+        """
+        if not self.table_exists(table):
+            logger.warning(f"Table {table} does not exist. Returning empty list.")
+            return []
+
+        try:
+            with self._connection.cursor() as cursor:
+                sql = f"SELECT id, data FROM `{table}` ORDER BY created_at DESC LIMIT %s OFFSET %s"
+                cursor.execute(sql, (limit, offset))
+                results = cursor.fetchall()
+                return [
+                    {"id": str(row["id"]), "content": row["data"], "metadata": {}}
+                    for row in results
+                ]
+        except Exception as e:
+            logger.error(f"Failed to list documents from table {table}: {e}")
+            return []
+
     def is_empty(self):
         pass

veadk/database/vector/opensearch_vector_database.py

@@ -219,7 +219,9 @@ class OpenSearchVectorDatabase(BaseModel, BaseDatabase):
         response = self._opensearch_client.indices.get_alias()
         return list(response.keys())
 
-    def get_all_docs(self, collection_name: str, size: int = 10000) -> list[dict]:
+    def list_docs(
+        self, collection_name: str, offset: int = 0, limit: int = 10000
+    ) -> list[dict]:
         """Match all docs in one index of OpenSearch"""
         if not self.collection_exists(collection_name):
             logger.warning(
@@ -227,12 +229,13 @@ class OpenSearchVectorDatabase(BaseModel, BaseDatabase):
             )
             return []
 
-        query = {"size": size, "query": {"match_all": {}}}
+        query = {"size": limit, "from": offset, "query": {"match_all": {}}}
         response = self._opensearch_client.search(index=collection_name, body=query)
         return [
             {
                 "id": hit["_id"],
-                "page_content": hit["_source"]["page_content"],
+                "content": hit["_source"]["page_content"],
+                "metadata": {},
             }
             for hit in response["hits"]["hits"]
         ]