veadk-python 0.1.0__py3-none-any.whl

veadk/database/kv/redis_database.py

@@ -0,0 +1,109 @@
+# 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, PrivateAttr
+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=getenv("DATABASE_REDIS_HOST"),
+        description="Redis host",
+    )
+    port: int = Field(
+        default=getenv("DATABASE_REDIS_PORT"),
+        description="Redis port",
+    )
+    db: int = Field(
+        default=getenv("DATABASE_REDIS_DB"),
+        description="Redis db",
+    )
+    password: str = Field(
+        default=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)
+    _client: redis.Redis = PrivateAttr(default=None)
+
+    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[str]:
+        try:
+            result = self._client.lrange(key, 0, -1)
+            return result
+        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
+            result = self._client.delete(key)
+            if result > 0:
+                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

veadk/database/local_database.py

@@ -0,0 +1,43 @@
+# 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"
+
+    def add_texts(self, texts: list[str], **kwargs):
+        self.data.extend(texts)
+
+    def is_empty(self):
+        return len(self.data) == 0
+
+    def query(self, query: str, **kwargs: Any) -> list[str]:
+        return self.data
+
+    def delete(self, **kwargs: Any):
+        self.data = []
+
+    def add(self, texts: list[str], **kwargs: Any):
+        return self.add_texts(texts)

veadk/database/relational/__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/database/relational/mysql_database.py

@@ -0,0 +1,114 @@
+# 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, PrivateAttr
+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=getenv("DATABASE_MYSQL_HOST"),
+        description="Mysql host",
+    )
+    user: str = Field(
+        default=getenv("DATABASE_MYSQL_USER"),
+        description="Mysql user",
+    )
+    password: str = Field(
+        default=getenv("DATABASE_MYSQL_PASSWORD"),
+        description="Mysql password",
+    )
+    database: str = Field(
+        default=getenv("DATABASE_MYSQL_DATABASE"),
+        description="Mysql database",
+    )
+    charset: str = Field(
+        default=getenv("DATABASE_MYSQL_CHARSET", "utf8mb4"),
+        description="Mysql charset",
+    )
+
+
+class MysqlDatabase(BaseModel, BaseDatabase):
+    config: MysqlDatabaseConfig = Field(default_factory=MysqlDatabaseConfig)
+
+    _connection: pymysql.Connection = PrivateAttr(default=None)
+
+    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._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) -> list[str]:
+        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 is_empty(self):
+        pass

veadk/database/vector/__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/database/vector/opensearch_vector_database.py

@@ -0,0 +1,205 @@
+# 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=getenv("DATABASE_OPENSEARCH_HOST"),
+        description="OpenSearch host",
+    )
+    port: str | int = Field(
+        default=getenv("DATABASE_OPENSEARCH_PORT"),
+        description="OpenSearch port",
+    )
+    username: Optional[str] = Field(
+        default=getenv("DATABASE_OPENSEARCH_USERNAME"),
+        description="OpenSearch username",
+    )
+    password: Optional[str] = Field(
+        default=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. retun 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):
+        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):
+        response = self._opensearch_client.count(index=collection_name)
+        return response["count"] == 0

veadk/database/vector/type.py

@@ -0,0 +1,50 @@
+# 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 requests
+
+from veadk.config import getenv
+
+
+class Embeddings:
+    def __init__(
+        self,
+        model: str = getenv("MODEL_EMBEDDING_NAME"),
+        api_base: str = getenv("MODEL_EMBEDDING_API_BASE"),
+        api_key: str = getenv("MODEL_EMBEDDING_API_KEY"),
+        dim: int = int(getenv("MODEL_EMBEDDING_DIM")),
+    ):
+        self.model = model
+        self.url = api_base
+        self.api_key = api_key
+        self.dim = dim
+
+        self.headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {self.api_key}",
+        }
+
+    def embed_documents(self, texts: list[str]) -> list[list[float]]:
+        MAX_CHARS = 4000
+        data = {"model": self.model, "input": [text[:MAX_CHARS] for text in texts]}
+        response = requests.post(self.url, headers=self.headers, json=data)
+        response.raise_for_status()
+        result = response.json()
+        return [item["embedding"] for item in result["data"]]
+
+    def embed_query(self, text: str) -> list[float]:
+        return self.embed_documents([text])[0]
+
+    def get_embedding_dim(self) -> int:
+        return self.dim

veadk/database/viking/__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.