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

veadk/database/kv/redis_database.py

@@ -1,159 +0,0 @@
-# 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 __future__ import annotations
-
-from typing import Any
-
-import redis
-from pydantic import BaseModel, Field
-from typing_extensions import override
-
-from veadk.config import getenv
-from veadk.utils.logger import get_logger
-
-from ..base_database import BaseDatabase
-
-logger = get_logger(__name__)
-
-
-class RedisDatabaseConfig(BaseModel):
-    host: str = Field(
-        default_factory=lambda: getenv("DATABASE_REDIS_HOST"),
-        description="Redis host",
-    )
-    port: int = Field(
-        default_factory=lambda: int(getenv("DATABASE_REDIS_PORT")),
-        description="Redis port",
-    )
-    db: int = Field(
-        default_factory=lambda: int(getenv("DATABASE_REDIS_DB")),
-        description="Redis db",
-    )
-    password: str = Field(
-        default_factory=lambda: getenv("DATABASE_REDIS_PASSWORD"),
-        description="Redis password",
-    )
-    decode_responses: bool = Field(
-        default=True,
-        description="Redis decode responses",
-    )
-
-
-class RedisDatabase(BaseModel, BaseDatabase):
-    config: RedisDatabaseConfig = Field(default_factory=RedisDatabaseConfig)
-
-    def model_post_init(self, context: Any, /) -> None:
-        try:
-            self._client = redis.StrictRedis(
-                host=self.config.host,
-                port=self.config.port,
-                db=self.config.db,
-                password=self.config.password,
-                decode_responses=self.config.decode_responses,
-            )
-
-            self._client.ping()
-            logger.info("Connected to Redis successfully.")
-        except Exception as e:
-            logger.error(f"Failed to connect to Redis: {e}")
-            raise e
-
-    @override
-    def add(self, key: str, value: str, **kwargs):
-        try:
-            self._client.rpush(key, value)
-        except Exception as e:
-            logger.error(f"Failed to add value to Redis list key `{key}`: {e}")
-            raise e
-
-    @override
-    def query(self, key: str, query: str = "", **kwargs) -> list:
-        try:
-            result = self._client.lrange(key, 0, -1)
-            return result  # type: ignore
-        except Exception as e:
-            logger.error(f"Failed to search from Redis list key '{key}': {e}")
-            raise e
-
-    @override
-    def delete(self, **kwargs):
-        """Delete Redis list key based on app_name, user_id and session_id, or directly by key."""
-        key = kwargs.get("key")
-        if key is None:
-            app_name = kwargs.get("app_name")
-            user_id = kwargs.get("user_id")
-            session_id = kwargs.get("session_id")
-            key = f"{app_name}:{user_id}:{session_id}"
-
-        try:
-            # For simple key deletion
-            # We use sync Redis client to delete the key
-            # so the result will be `int`
-            result = self._client.delete(key)
-
-            if result > 0:  # type: ignore
-                logger.info(f"Deleted key `{key}` from Redis.")
-            else:
-                logger.info(f"Key `{key}` not found in Redis. Skipping deletion.")
-        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

@@ -1,62 +0,0 @@
-# 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 typing import Any
-
-from .base_database import BaseDatabase
-
-
-class LocalDataBase(BaseDatabase):
-    """This database is only for basic demonstration.
-    It does not support the vector search function, and the `search` function will return all data.
-    """
-
-    def __init__(self, **kwargs):
-        super().__init__()
-        self.data = {}
-        self._type = "local"
-        self._next_id = 0  # Used to generate unique IDs
-
-    def add_texts(self, texts: list[str], **kwargs):
-        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 list(self.data.values())
-
-    def delete(self, **kwargs: Any):
-        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

@@ -1,173 +0,0 @@
-# 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 __future__ import annotations
-
-from typing import Any
-
-import pymysql
-from pydantic import BaseModel, Field
-from typing_extensions import override
-
-from veadk.config import getenv
-from veadk.utils.logger import get_logger
-
-from ..base_database import BaseDatabase
-
-logger = get_logger(__name__)
-
-
-class MysqlDatabaseConfig(BaseModel):
-    host: str = Field(
-        default_factory=lambda: getenv("DATABASE_MYSQL_HOST"),
-        description="Mysql host",
-    )
-    user: str = Field(
-        default_factory=lambda: getenv("DATABASE_MYSQL_USER"),
-        description="Mysql user",
-    )
-    password: str = Field(
-        default_factory=lambda: getenv("DATABASE_MYSQL_PASSWORD"),
-        description="Mysql password",
-    )
-    database: str = Field(
-        default_factory=lambda: getenv("DATABASE_MYSQL_DATABASE"),
-        description="Mysql database",
-    )
-    charset: str = Field(
-        default_factory=lambda: getenv("DATABASE_MYSQL_CHARSET", "utf8mb4"),
-        description="Mysql charset",
-    )
-
-
-class MysqlDatabase(BaseModel, BaseDatabase):
-    config: MysqlDatabaseConfig = Field(default_factory=MysqlDatabaseConfig)
-
-    def model_post_init(self, context: Any, /) -> None:
-        self._connection = pymysql.connect(
-            host=self.config.host,
-            user=self.config.user,
-            password=self.config.password,
-            database=self.config.database,
-            charset=self.config.charset,
-            cursorclass=pymysql.cursors.DictCursor,
-        )
-        self._connection.ping()
-        logger.info("Connected to MySQL successfully.")
-
-        self._type = "mysql"
-
-    def table_exists(self, table: str) -> bool:
-        with self._connection.cursor() as cursor:
-            cursor.execute(
-                "SELECT TABLE_NAME FROM information_schema.TABLES WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s",
-                (self.config.database, table),
-            )
-            result = cursor.fetchone()
-            return result is not None
-
-    @override
-    def add(self, sql: str, params=None, **kwargs):
-        with self._connection.cursor() as cursor:
-            cursor.execute(sql, params)
-            self._connection.commit()
-
-    @override
-    def query(self, sql: str, params=None, **kwargs) -> tuple[dict[str, Any], ...]:
-        with self._connection.cursor() as cursor:
-            cursor.execute(sql, params)
-            return cursor.fetchall()
-
-    @override
-    def delete(self, **kwargs):
-        table = kwargs.get("table")
-        if table is None:
-            app_name = kwargs.get("app_name", "default")
-            table = app_name
-
-        if not self.table_exists(table):
-            logger.warning(f"Table {table} does not exist. Skipping delete operation.")
-            return
-
-        try:
-            with self._connection.cursor() as cursor:
-                # Drop the table directly
-                sql = f"DROP TABLE `{table}`"
-                cursor.execute(sql)
-                self._connection.commit()
-                logger.info(f"Dropped table {table}")
-        except Exception as e:
-            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

@@ -1,263 +0,0 @@
-# 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 __future__ import annotations
-
-import os
-from typing import Any, Literal, Optional
-
-from opensearchpy import OpenSearch, Urllib3HttpConnection, helpers
-from pydantic import BaseModel, Field, PrivateAttr
-from typing_extensions import override
-
-from veadk.config import getenv
-from veadk.utils.logger import get_logger
-
-from ..base_database import BaseDatabase
-from .type import Embeddings
-
-logger = get_logger(__name__)
-
-
-class OpenSearchVectorDatabaseConfig(BaseModel):
-    host: str = Field(
-        default_factory=lambda: getenv("DATABASE_OPENSEARCH_HOST"),
-        description="OpenSearch host",
-    )
-
-    port: str | int = Field(
-        default_factory=lambda: getenv("DATABASE_OPENSEARCH_PORT"),
-        description="OpenSearch port",
-    )
-
-    username: Optional[str] = Field(
-        default_factory=lambda: getenv("DATABASE_OPENSEARCH_USERNAME"),
-        description="OpenSearch username",
-    )
-
-    password: Optional[str] = Field(
-        default_factory=lambda: getenv("DATABASE_OPENSEARCH_PASSWORD"),
-        description="OpenSearch password",
-    )
-
-    secure: bool = Field(default=True, description="Whether enable SSL")
-
-    verify_certs: bool = Field(default=False, description="Whether verify SSL certs")
-
-    auth_method: Literal["basic", "aws_managed_iam"] = Field(
-        default="basic", description="OpenSearch auth method"
-    )
-
-    def to_opensearch_params(self) -> dict[str, Any]:
-        params = {
-            "hosts": [{"host": self.host, "port": int(self.port)}],
-            "use_ssl": self.secure,
-            "verify_certs": self.verify_certs,
-            "connection_class": Urllib3HttpConnection,
-            "pool_maxsize": 20,
-        }
-        ca_cert_path = os.getenv("OPENSEARCH_CA_CERT")
-        if self.verify_certs and ca_cert_path:
-            params["ca_certs"] = ca_cert_path
-
-        params["http_auth"] = (self.username, self.password)
-
-        return params
-
-
-class OpenSearchVectorDatabase(BaseModel, BaseDatabase):
-    config: OpenSearchVectorDatabaseConfig = Field(
-        default_factory=OpenSearchVectorDatabaseConfig
-    )
-
-    _embedding_client: Embeddings = PrivateAttr()
-    _opensearch_client: OpenSearch = PrivateAttr()
-
-    def model_post_init(self, context: Any, /) -> None:
-        self._embedding_client = Embeddings()
-        self._opensearch_client = OpenSearch(**self.config.to_opensearch_params())
-
-        self._type = "opensearch"
-
-    def _get_settings(self) -> dict:
-        settings = {"index": {"knn": True}}
-        return settings
-
-    def _get_mappings(self, dim: int = 2560) -> dict:
-        mappings = {
-            "properties": {
-                "page_content": {
-                    "type": "text",
-                },
-                "vector": {
-                    "type": "knn_vector",
-                    "dimension": dim,
-                    "method": {
-                        "name": "hnsw",
-                        "space_type": "l2",
-                        "engine": "faiss",
-                        "parameters": {"ef_construction": 64, "m": 8},
-                    },
-                },
-            }
-        }
-        return mappings
-
-    def create_collection(
-        self,
-        collection_name: str,
-        embedding_dim: int,
-    ):
-        if not self._opensearch_client.indices.exists(index=collection_name):
-            self._opensearch_client.indices.create(
-                index=collection_name,
-                body={
-                    "mappings": self._get_mappings(dim=embedding_dim),
-                    "settings": self._get_settings(),
-                },
-            )
-        else:
-            logger.warning(f"Collection {collection_name} already exists.")
-
-        self._opensearch_client.indices.refresh(index=collection_name)
-        return
-
-    def _search_by_vector(
-        self, collection_name: str, query_vector: list[float], **kwargs: Any
-    ) -> list[str]:
-        top_k = kwargs.get("top_k", 5)
-        query = {
-            "size": top_k,
-            "query": {"knn": {"vector": {"vector": query_vector, "k": top_k}}},
-        }
-        response = self._opensearch_client.search(index=collection_name, body=query)
-
-        result_list = []
-        for hit in response["hits"]["hits"]:
-            result_list.append(hit["_source"]["page_content"])
-
-        return result_list
-
-    def get_health(self):
-        response = self._opensearch_client.cat.health()
-        logger.info(response)
-
-    def add(self, texts: list[str], **kwargs):
-        collection_name = kwargs.get("collection_name")
-        assert collection_name is not None, "Collection name is required."
-        if not self._opensearch_client.indices.exists(index=collection_name):
-            self.create_collection(
-                embedding_dim=self._embedding_client.get_embedding_dim(),
-                collection_name=collection_name,
-            )
-
-        actions = []
-        embeddings = self._embedding_client.embed_documents(texts)
-        for i in range(len(texts)):
-            action = {
-                "_op_type": "index",
-                "_index": collection_name,
-                "_source": {
-                    "page_content": texts[i],
-                    "vector": embeddings[i],
-                },
-            }
-            actions.append(action)
-
-        helpers.bulk(
-            client=self._opensearch_client,
-            actions=actions,
-            timeout=30,
-            max_retries=3,
-        )
-
-        self._opensearch_client.indices.refresh(index=collection_name)
-        return
-
-    @override
-    def query(self, query: str, **kwargs: Any) -> list[str]:
-        collection_name = kwargs.get("collection_name")
-        top_k = kwargs.get("top_k", 5)
-        assert collection_name is not None, "Collection name is required."
-        if not self._opensearch_client.indices.exists(index=collection_name):
-            logger.warning(
-                f"querying {query}, but collection {collection_name} does not exist. return a empty list."
-            )
-            return []
-        query_vector = self._embedding_client.embed_query(query)
-        return self._search_by_vector(
-            collection_name=collection_name, query_vector=query_vector, top_k=top_k
-        )
-
-    @override
-    def delete(self, collection_name: str, **kwargs: Any):
-        """drop index"""
-        if not self._opensearch_client.indices.exists(index=collection_name):
-            raise ValueError(f"Collection {collection_name} does not exist.")
-        self._opensearch_client.indices.delete(index=collection_name)
-
-    def is_empty(self, collection_name: str) -> bool:
-        response = self._opensearch_client.count(index=collection_name)
-        return response["count"] == 0
-
-    def collection_exists(self, collection_name: str) -> bool:
-        return self._opensearch_client.indices.exists(index=collection_name)
-
-    def list_all_collection(self) -> list:
-        """List all index name of OpenSearch."""
-        response = self._opensearch_client.indices.get_alias()
-        return list(response.keys())
-
-    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(
-                f"Get all docs, but collection {collection_name} does not exist. return a empty list."
-            )
-            return []
-
-        query = {"size": limit, "from": offset, "query": {"match_all": {}}}
-        response = self._opensearch_client.search(index=collection_name, body=query)
-        return [
-            {
-                "id": hit["_id"],
-                "content": hit["_source"]["page_content"],
-                "metadata": {},
-            }
-            for hit in response["hits"]["hits"]
-        ]
-
-    def delete_by_query(self, collection_name: str, query: str) -> Any:
-        """Delete docs by query in one index of OpenSearch"""
-        if not self.collection_exists(collection_name):
-            raise ValueError(f"Collection {collection_name} does not exist.")
-
-        query_payload = {"query": {"match": {"page_content": query}}}
-        response = self._opensearch_client.delete_by_query(
-            index=collection_name, body=query_payload
-        )
-
-        self._opensearch_client.indices.refresh(index=collection_name)
-        return response
-
-    def delete_by_id(self, collection_name: str, id: str):
-        """Delete docs by id in index of OpenSearch"""
-        if not self.collection_exists(collection_name):
-            raise ValueError(f"Collection {collection_name} does not exist.")
-
-        response = self._opensearch_client.delete(index=collection_name, id=id)
-        self._opensearch_client.indices.refresh(index=collection_name)
-        return response