viking_database.py

# 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 io
import json
import os
import uuid
from typing import Any, BinaryIO, Literal, Optional, TextIO

import requests
import tos
from pydantic import BaseModel, Field
from volcengine.auth.SignerV4 import SignerV4
from volcengine.base.Request import Request
from volcengine.Credentials import Credentials

from veadk.config import getenv
from veadk.database.base_database import BaseDatabase
from veadk.utils.logger import get_logger

logger = get_logger(__name__)

# knowledge base domain
g_knowledge_base_domain = "api-knowledgebase.mlp.cn-beijing.volces.com"
# paths
create_collection_path = "/api/knowledge/collection/create"
search_knowledge_path = "/api/knowledge/collection/search_knowledge"
list_collections_path = "/api/knowledge/collection/list"
get_collections_path = "/api/knowledge/collection/info"
doc_add_path = "/api/knowledge/doc/add"
doc_info_path = "/api/knowledge/doc/info"
doc_del_path = "/api/collection/drop"


class VolcengineTOSConfig(BaseModel):
    endpoint: Optional[str] = Field(
        default=getenv("DATABASE_TOS_ENDPOINT", "tos-cn-beijing.volces.com"),
        description="VikingDB TOS endpoint",
    )
    region: Optional[str] = Field(
        default=getenv("DATABASE_TOS_REGION", "cn-beijing"),
        description="VikingDB TOS region",
    )
    bucket: Optional[str] = Field(
        default=getenv("DATABASE_TOS_BUCKET"),
        description="VikingDB TOS bucket",
    )
    base_key: Optional[str] = Field(
        default="veadk",
        description="VikingDB TOS base key",
    )


class VikingDatabaseConfig(BaseModel):
    volcengine_ak: Optional[str] = Field(
        default=getenv("VOLCENGINE_ACCESS_KEY"),
        description="VikingDB access key",
    )
    volcengine_sk: Optional[str] = Field(
        default=getenv("VOLCENGINE_SECRET_KEY"),
        description="VikingDB secret key",
    )
    project: Optional[str] = Field(
        default=getenv("DATABASE_VIKING_PROJECT"),
        description="VikingDB project name",
    )
    region: Optional[str] = Field(
        default=getenv("DATABASE_VIKING_REGION"),
        description="VikingDB region",
    )
    tos: Optional[VolcengineTOSConfig] = Field(
        default_factory=VolcengineTOSConfig,
        description="VikingDB TOS configuration",
    )


def prepare_request(
    method, path, config: VikingDatabaseConfig, params=None, data=None, doseq=0
):
    ak = config.volcengine_ak
    sk = config.volcengine_sk

    if params:
        for key in params:
            if (
                type(params[key]) is int
                or type(params[key]) is float
                or type(params[key]) is bool
            ):
                params[key] = str(params[key])
            elif type(params[key]) is list:
                if not doseq:
                    params[key] = ",".join(params[key])
    r = Request()
    r.set_shema("https")
    r.set_method(method)
    r.set_connection_timeout(10)
    r.set_socket_timeout(10)
    mheaders = {
        "Accept": "application/json",
        "Content-Type": "application/json",
    }
    r.set_headers(mheaders)
    if params:
        r.set_query(params)
    r.set_path(path)
    if data is not None:
        r.set_body(json.dumps(data))
    credentials = Credentials(ak, sk, "air", config.region)
    SignerV4.sign(r, credentials)
    return r


class VikingDatabase(BaseModel, BaseDatabase):
    config: VikingDatabaseConfig = Field(
        default_factory=VikingDatabaseConfig,
        description="VikingDB configuration",
    )

    def _upload_to_tos(
        self,
        data: str | list[str] | TextIO | BinaryIO | bytes,
        **kwargs: Any,
    ):
        file_ext = kwargs.get(
            "file_ext", ".pdf"
        )  # when bytes data, file_ext is required
        ak = self.config.volcengine_ak
        sk = self.config.volcengine_sk
        tos_bucket = self.config.tos.bucket
        tos_endpoint = self.config.tos.endpoint
        tos_region = self.config.tos.region
        tos_key = self.config.tos.base_key
        client = tos.TosClientV2(ak, sk, tos_endpoint, tos_region, max_connections=1024)
        if isinstance(data, str) and os.path.isfile(data):  # Process file path
            file_ext = os.path.splitext(data)[1]
            new_key = f"{tos_key}/{str(uuid.uuid4())}{file_ext}"
            with open(data, "rb") as f:
                upload_data = f.read()

        elif isinstance(
            data,
            (io.TextIOWrapper, io.BufferedReader),  # file type: TextIO | BinaryIO
        ):  # Process file stream
            # Try to get the file extension from the file name, and use the default value if there is none
            file_ext = ".unknown"
            if hasattr(data, "name"):
                _, file_ext = os.path.splitext(data.name)
            new_key = f"{tos_key}/{str(uuid.uuid4())}{file_ext}"
            if isinstance(data, TextIO):
                # Encode the text stream content into bytes
                upload_data = data.read().encode("utf-8")
            else:
                # Read the content of the binary stream
                upload_data = data.read()

        elif isinstance(data, str):  # Process ordinary strings
            new_key = f"{tos_key}/{str(uuid.uuid4())}.txt"
            upload_data = data.encode("utf-8")  # Encode as byte type

        elif isinstance(data, list):  # Process list of strings
            new_key = f"{tos_key}/{str(uuid.uuid4())}.txt"
            # Join the strings in the list with newlines and encode as byte type
            upload_data = "\n".join(data).encode("utf-8")

        elif isinstance(data, bytes):  # Process bytes data
            new_key = f"{tos_key}/{str(uuid.uuid4())}{file_ext}"
            upload_data = data

        else:
            raise ValueError(f"Unsupported data type: {type(data)}")

        resp = client.put_object(tos_bucket, new_key, content=upload_data)
        tos_url = f"{tos_bucket}/{new_key}"

        return resp.resp.status, tos_url

    def _add_doc(self, collection_name: str, tos_url: str, doc_id: str, **kwargs: Any):
        request_params = {
            "collection_name": collection_name,
            "project": self.config.project,
            "add_type": "tos",
            "doc_id": doc_id,
            "tos_path": tos_url,
        }

        doc_add_req = prepare_request(
            method="POST", path=doc_add_path, config=self.config, data=request_params
        )
        rsp = requests.request(
            method=doc_add_req.method,
            url="https://{}{}".format(g_knowledge_base_domain, doc_add_req.path),
            headers=doc_add_req.headers,
            data=doc_add_req.body,
        )

        result = rsp.json()
        if result["code"] != 0:
            logger.error(f"Error in add_doc: {result['message']}")
            return {"error": result["message"]}

        doc_add_data = result["data"]
        if not doc_add_data:
            raise ValueError(f"doc {doc_id} has no data.")

        return doc_id

    def add(
        self,
        data: str | list[str] | TextIO | BinaryIO | bytes,
        collection_name: str,
        **kwargs,
    ):
        """
        Args:
            data: str, file path or file stream:  Both file or file.read() are acceptable.
            **kwargs: collection_name(required)
        Returns:
            {
                "tos_url": "tos://<bucket>/<key>",
                "doc_id": "<doc_id>",
            }
        """

        status, tos_url = self._upload_to_tos(data=data, **kwargs)
        if status != 200:
            raise ValueError(f"Error in upload_to_tos: {status}")
        doc_id = self._add_doc(
            collection_name=collection_name,
            tos_url=tos_url,
            doc_id=str(uuid.uuid4()),
        )
        return {
            "tos_url": f"tos://{tos_url}",
            "doc_id": doc_id,
        }

    def delete(self, **kwargs: Any):
        collection_name = kwargs.get("collection_name")
        resource_id = kwargs.get("resource_id")
        request_param = {"collection_name": collection_name, "resource_id": resource_id}
        doc_del_req = prepare_request(
            method="POST", path=doc_del_path, config=self.config, data=request_param
        )
        rsp = requests.request(
            method=doc_del_req.method,
            url="http://{}{}".format(g_knowledge_base_domain, doc_del_req.path),
            headers=doc_del_req.headers,
            data=doc_del_req.body,
        )
        result = rsp.json()
        if result["code"] != 0:
            logger.error(f"Error in add_doc: {result['message']}")
            return {"error": result["message"]}
        return {}

    def query(self, query: str, **kwargs: Any) -> list[str]:
        """
        Args:
            query:  query text
            **kwargs: collection_name(required), top_k(optional, default 5)

        Returns: list of str, the search result
        """
        collection_name = kwargs.get("collection_name")
        assert collection_name is not None, "collection_name is required"
        request_params = {
            "query": query,
            "limit": int(kwargs.get("top_k", 5)),
            "name": collection_name,
            "project": self.config.project,
        }
        search_req = prepare_request(
            method="POST",
            path=search_knowledge_path,
            config=self.config,
            data=request_params,
        )
        resp = requests.request(
            method=search_req.method,
            url="https://{}{}".format(g_knowledge_base_domain, search_req.path),
            headers=search_req.headers,
            data=search_req.body,
        )

        result = resp.json()
        if result["code"] != 0:
            logger.error(f"Error in search_knowledge: {result['message']}")
            raise ValueError(f"Error in search_knowledge: {result['message']}")

        if not result["data"]["result_list"]:
            raise ValueError(f"No results found for collection {collection_name}")

        chunks = result["data"]["result_list"]

        search_result = []

        for chunk in chunks:
            search_result.append(chunk["content"])

        return search_result

    def create_collection(
        self,
        collection_name: str,
        description: str = "",
        version: Literal[2, 4] = 4,
        data_type: Literal[
            "unstructured_data", "structured_data"
        ] = "unstructured_data",
        chunking_strategy: Literal["custom_balance", "custom"] = "custom_balance",
        chunk_length: int = 500,
        merge_small_chunks: bool = True,
    ):
        request_params = {
            "name": collection_name,
            "project": self.config.project,
            "description": description,
            "version": version,
            "data_type": data_type,
            "preprocessing": {
                "chunking_strategy": chunking_strategy,
                "chunk_length": chunk_length,
                "merge_small_chunks": merge_small_chunks,
            },
        }

        create_collection_req = prepare_request(
            method="POST",
            path=create_collection_path,
            config=self.config,
            data=request_params,
        )
        resp = requests.request(
            method=create_collection_req.method,
            url="https://{}{}".format(
                g_knowledge_base_domain, create_collection_req.path
            ),
            headers=create_collection_req.headers,
            data=create_collection_req.body,
        )

        result = resp.json()
        if result["code"] != 0:
            logger.error(f"Error in create_collection: {result['message']}")
            raise ValueError(f"Error in create_collection: {result['message']}")
        return result

    def collection_exists(self, collection_name: str) -> bool:
        request_params = {
            "project": self.config.project,
        }
        list_collections_req = prepare_request(
            method="POST",
            path=list_collections_path,
            config=self.config,
            data=request_params,
        )
        resp = requests.request(
            method=list_collections_req.method,
            url="https://{}{}".format(
                g_knowledge_base_domain, list_collections_req.path
            ),
            headers=list_collections_req.headers,
            data=list_collections_req.body,
        )

        result = resp.json()
        if result["code"] != 0:
            logger.error(f"Error in list_collections: {result['message']}")
            raise ValueError(f"Error in list_collections: {result['message']}")

        collections = result["data"]["collection_list"]
        if not collections:
            raise ValueError(f"No collections found in project {self.config.project}.")

        collection_list = set()

        for collection in collections:
            collection_list.add(collection["collection_name"])
        # check the collection exist or not
        if collection_name in collection_list:
            return True
        else:
            return False