veadk-python 0.2.27__py3-none-any.whl

veadk/integrations/ve_viking_db_memory/__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/integrations/ve_viking_db_memory/ve_viking_db_memory.py

@@ -0,0 +1,308 @@
+# 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 json
+import threading
+from veadk.utils.misc import getenv
+from volcengine.ApiInfo import ApiInfo
+from volcengine.auth.SignerV4 import SignerV4
+from volcengine.base.Service import Service
+from volcengine.Credentials import Credentials
+from volcengine.ServiceInfo import ServiceInfo
+
+
+class VikingDBMemoryException(Exception):
+    def __init__(self, code, request_id, message=None):
+        self.code = code
+        self.request_id = request_id
+        self.message = "{}, code:{}，request_id:{}".format(
+            message, self.code, self.request_id
+        )
+
+    def __str__(self):
+        return self.message
+
+
+class VikingDBMemoryClient(Service):
+    _instance_lock = threading.Lock()
+
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(VikingDBMemoryClient, "_instance"):
+            with VikingDBMemoryClient._instance_lock:
+                if not hasattr(VikingDBMemoryClient, "_instance"):
+                    VikingDBMemoryClient._instance = object.__new__(cls)
+        return VikingDBMemoryClient._instance
+
+    def __init__(
+        self,
+        host="api-knowledgebase.mlp.cn-beijing.volces.com",
+        region="cn-beijing",
+        ak="",
+        sk="",
+        sts_token="",
+        scheme="https",
+        connection_timeout=30,
+        socket_timeout=30,
+    ):
+        env_host = getenv(
+            "DATABASE_VIKINGMEM_BASE_URL", default_value=None, allow_false_values=True
+        )
+        if env_host:
+            if env_host.startswith("http://"):
+                host = env_host.replace("http://", "")
+                scheme = "http"
+            elif env_host.startswith("https://"):
+                host = env_host.replace("https://", "")
+                scheme = "https"
+            else:
+                raise ValueError(
+                    "DATABASE_VIKINGMEM_BASE_URL must start with http:// or https://"
+                )
+
+        self.service_info = VikingDBMemoryClient.get_service_info(
+            host, region, scheme, connection_timeout, socket_timeout
+        )
+        self.api_info = VikingDBMemoryClient.get_api_info()
+        super(VikingDBMemoryClient, self).__init__(self.service_info, self.api_info)
+        if ak:
+            self.set_ak(ak)
+        if sk:
+            self.set_sk(sk)
+        if sts_token:
+            self.set_session_token(session_token=sts_token)
+        try:
+            self.get_body("Ping", {}, json.dumps({}))
+        except Exception as e:
+            raise VikingDBMemoryException(
+                1000028, "missed", "host or region is incorrect: {}".format(str(e))
+            ) from None
+
+    def setHeader(self, header):
+        api_info = VikingDBMemoryClient.get_api_info()
+        for key in api_info:
+            for item in header:
+                api_info[key].header[item] = header[item]
+        self.api_info = api_info
+
+    @staticmethod
+    def get_service_info(host, region, scheme, connection_timeout, socket_timeout):
+        service_info = ServiceInfo(
+            host,
+            {"Host": host},
+            Credentials("", "", "air", region),
+            connection_timeout,
+            socket_timeout,
+            scheme=scheme,
+        )
+        return service_info
+
+    @staticmethod
+    def get_api_info():
+        api_info = {
+            "CreateCollection": ApiInfo(
+                "POST",
+                "/api/memory/collection/create",
+                {},
+                {},
+                {"Accept": "application/json", "Content-Type": "application/json"},
+            ),
+            "GetCollection": ApiInfo(
+                "POST",
+                "/api/memory/collection/info",
+                {},
+                {},
+                {"Accept": "application/json", "Content-Type": "application/json"},
+            ),
+            "DropCollection": ApiInfo(
+                "POST",
+                "/api/memory/collection/delete",
+                {},
+                {},
+                {"Accept": "application/json", "Content-Type": "application/json"},
+            ),
+            "UpdateCollection": ApiInfo(
+                "POST",
+                "/api/memory/collection/update",
+                {},
+                {},
+                {"Accept": "application/json", "Content-Type": "application/json"},
+            ),
+            "SearchMemory": ApiInfo(
+                "POST",
+                "/api/memory/search",
+                {},
+                {},
+                {"Accept": "application/json", "Content-Type": "application/json"},
+            ),
+            "AddMessages": ApiInfo(
+                "POST",
+                "/api/memory/messages/add",
+                {},
+                {},
+                {"Accept": "application/json", "Content-Type": "application/json"},
+            ),
+            "Ping": ApiInfo(
+                "GET",
+                "/api/memory/ping",
+                {},
+                {},
+                {"Accept": "application/json", "Content-Type": "application/json"},
+            ),
+        }
+        return api_info
+
+    def get_body(self, api, params, body):
+        if api not in self.api_info:
+            raise Exception("no such api")
+        api_info = self.api_info[api]
+        r = self.prepare_request(api_info, params)
+        r.headers["Content-Type"] = "application/json"
+        r.headers["Traffic-Source"] = "SDK"
+        r.body = body
+
+        SignerV4.sign(r, self.service_info.credentials)
+
+        url = r.build()
+        resp = self.session.get(
+            url,
+            headers=r.headers,
+            data=r.body,
+            timeout=(
+                self.service_info.connection_timeout,
+                self.service_info.socket_timeout,
+            ),
+        )
+        if resp.status_code == 200:
+            return json.dumps(resp.json())
+        else:
+            raise Exception(resp.text.encode("utf-8"))
+
+    def get_body_exception(self, api, params, body):
+        try:
+            res = self.get_body(api, params, body)
+        except Exception as e:
+            try:
+                res_json = json.loads(e.args[0].decode("utf-8"))
+            except Exception as e:
+                raise VikingDBMemoryException(
+                    1000028, "missed", "json load res error, res:{}".format(str(e))
+                ) from None
+            code = res_json.get("code", 1000028)
+            request_id = res_json.get("request_id", 1000028)
+            message = res_json.get("message", None)
+
+            raise VikingDBMemoryException(code, request_id, message)
+
+        if res == "":
+            raise VikingDBMemoryException(
+                1000028,
+                "missed",
+                "empty response due to unknown error, please contact customer service",
+            ) from None
+        return res
+
+    def get_exception(self, api, params):
+        try:
+            res = self.get(api, params)
+        except Exception as e:
+            try:
+                res_json = json.loads(e.args[0].decode("utf-8"))
+            except Exception as e:
+                raise VikingDBMemoryException(
+                    1000028, "missed", "json load res error, res:{}".format(str(e))
+                ) from None
+            code = res_json.get("code", 1000028)
+            request_id = res_json.get("request_id", 1000028)
+            message = res_json.get("message", None)
+            raise VikingDBMemoryException(code, request_id, message)
+        if res == "":
+            raise VikingDBMemoryException(
+                1000028,
+                "missed",
+                "empty response due to unknown error, please contact customer service",
+            ) from None
+        return res
+
+    def create_collection(
+        self,
+        collection_name,
+        description="",
+        custom_event_type_schemas=[],
+        custom_entity_type_schemas=[],
+        builtin_event_types=[],
+        builtin_entity_types=[],
+    ):
+        params = {
+            "CollectionName": collection_name,
+            "Description": description,
+            "CustomEventTypeSchemas": custom_event_type_schemas,
+            "CustomEntityTypeSchemas": custom_entity_type_schemas,
+            "BuiltinEventTypes": builtin_event_types,
+            "BuiltinEntityTypes": builtin_entity_types,
+        }
+        res = self.json("CreateCollection", {}, json.dumps(params))
+        return json.loads(res)
+
+    def get_collection(self, collection_name):
+        params = {"CollectionName": collection_name}
+        res = self.json("GetCollection", {}, json.dumps(params))
+        return json.loads(res)
+
+    def drop_collection(self, collection_name):
+        params = {"CollectionName": collection_name}
+        res = self.json("DropCollection", {}, json.dumps(params))
+        return json.loads(res)
+
+    def update_collection(
+        self,
+        collection_name,
+        custom_event_type_schemas=[],
+        custom_entity_type_schemas=[],
+        builtin_event_types=[],
+        builtin_entity_types=[],
+    ):
+        params = {
+            "CollectionName": collection_name,
+            "CustomEventTypeSchemas": custom_event_type_schemas,
+            "CustomEntityTypeSchemas": custom_entity_type_schemas,
+            "BuiltinEventTypes": builtin_event_types,
+            "BuiltinEntityTypes": builtin_entity_types,
+        }
+        res = self.json("UpdateCollection", {}, json.dumps(params))
+        return json.loads(res)
+
+    def search_memory(self, collection_name, query, filter, limit=10):
+        params = {
+            "collection_name": collection_name,
+            "limit": limit,
+            "filter": filter,
+        }
+        if query:
+            params["query"] = query
+        res = self.json("SearchMemory", {}, json.dumps(params))
+        return json.loads(res)
+
+    def add_messages(
+        self, collection_name, session_id, messages, metadata, entities=None
+    ):
+        params = {
+            "collection_name": collection_name,
+            "session_id": session_id,
+            "messages": messages,
+            "metadata": metadata,
+        }
+        if entities is not None:
+            params["entities"] = entities
+        res = self.json("AddMessages", {}, json.dumps(params))
+        return json.loads(res)

veadk/knowledgebase/__init__.py

@@ -0,0 +1,17 @@
+# 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 .knowledgebase import KnowledgeBase
+
+__all__ = ["KnowledgeBase"]

veadk/knowledgebase/backends/__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/knowledgebase/backends/base_backend.py

@@ -0,0 +1,72 @@
+# 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 abc import ABC, abstractmethod
+
+from pydantic import BaseModel
+
+
+class BaseKnowledgebaseBackend(ABC, BaseModel):
+    """Base backend for knowledgebase.
+
+    Attributes:
+        index (str): Index or collection name of the vector storage.
+
+    Examples:
+        You can implement your own knowledgebase backend.
+
+        ```python
+        class CustomKnowledgebaseBackend(BaseKnowledgebaseBackend):
+            pass
+        ```
+    """
+
+    index: str
+
+    @abstractmethod
+    def precheck_index_naming(self) -> None:
+        """Check the index name is valid or not.
+
+        If index naming is not valid, raise an exception.
+        """
+
+    @abstractmethod
+    def add_from_directory(self, directory: str, *args, **kwargs) -> bool:
+        """Add knowledge from file path to knowledgebase"""
+
+    @abstractmethod
+    def add_from_files(self, files: list[str], *args, **kwargs) -> bool:
+        """Add knowledge (e.g, documents, strings, ...) to knowledgebase"""
+
+    @abstractmethod
+    def add_from_text(self, text: str | list[str], *args, **kwargs) -> bool:
+        """Add knowledge from text to knowledgebase"""
+
+    @abstractmethod
+    def search(self, *args, **kwargs) -> list:
+        """Search knowledge from knowledgebase"""
+
+    # Optional methods for future use:
+    # - `delete`: Delete collection or documents
+    # - `list_docs`: List original documents
+    # - `list_chunks`: List embedded document chunks
+
+    # def delete(self, **kwargs) -> bool:
+    #     """Delete knowledge from knowledgebase"""
+
+    # def list_docs(self, **kwargs) -> None:
+    #     """List original documents in knowledgebase"""
+
+    # def list_chunks(self, **kwargs) -> None:
+    #     """List embeded document chunks in knowledgebase"""

veadk/knowledgebase/backends/in_memory_backend.py

@@ -0,0 +1,91 @@
+# 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 llama_index.core import Document, SimpleDirectoryReader, VectorStoreIndex
+from llama_index.core.schema import BaseNode
+from llama_index.embeddings.openai_like import OpenAILikeEmbedding
+from pydantic import Field
+from typing_extensions import Any, override
+
+from veadk.configs.model_configs import EmbeddingModelConfig, NormalEmbeddingModelConfig
+from veadk.knowledgebase.backends.base_backend import BaseKnowledgebaseBackend
+from veadk.knowledgebase.backends.utils import get_llama_index_splitter
+
+
+class InMemoryKnowledgeBackend(BaseKnowledgebaseBackend):
+    """A in-memory implementation backend for knowledgebase.
+
+    In-memory backend stores embedded text in a vector storage from Llama-index.
+
+    Attributes:
+        embedding_config (EmbeddingModelConfig):
+            Embedding config for text embedding and search.
+            Embedding config contains embedding model name and the corresponding dim.
+    """
+
+    embedding_config: NormalEmbeddingModelConfig | EmbeddingModelConfig = Field(
+        default_factory=EmbeddingModelConfig
+    )
+
+    def model_post_init(self, __context: Any) -> None:
+        self._embed_model = OpenAILikeEmbedding(
+            model_name=self.embedding_config.name,
+            api_key=self.embedding_config.api_key,
+            api_base=self.embedding_config.api_base,
+        )
+        self._vector_index = VectorStoreIndex([], embed_model=self._embed_model)
+
+    @override
+    def precheck_index_naming(self) -> None:
+        # Checking is not needed
+        pass
+
+    @override
+    def add_from_directory(self, directory: str) -> bool:
+        documents = SimpleDirectoryReader(input_dir=directory).load_data()
+        nodes = self._split_documents(documents)
+        self._vector_index.insert_nodes(nodes)
+        return True
+
+    @override
+    def add_from_files(self, files: list[str]) -> bool:
+        documents = SimpleDirectoryReader(input_files=files).load_data()
+        nodes = self._split_documents(documents)
+        self._vector_index.insert_nodes(nodes)
+        return True
+
+    @override
+    def add_from_text(self, text: str | list[str]) -> bool:
+        if isinstance(text, str):
+            documents = [Document(text=text)]
+        else:
+            documents = [Document(text=t) for t in text]
+        nodes = self._split_documents(documents)
+        self._vector_index.insert_nodes(nodes)
+        return True
+
+    @override
+    def search(self, query: str, top_k: int = 5) -> list[str]:
+        _retriever = self._vector_index.as_retriever(similarity_top_k=top_k)
+        retrieved_nodes = _retriever.retrieve(query)
+        return [node.text for node in retrieved_nodes]
+
+    def _split_documents(self, documents: list[Document]) -> list[BaseNode]:
+        """Split document into chunks"""
+        nodes = []
+        for document in documents:
+            splitter = get_llama_index_splitter(document.metadata.get("file_path", ""))
+            _nodes = splitter.get_nodes_from_documents([document])
+            nodes.extend(_nodes)
+        return nodes

veadk/knowledgebase/backends/opensearch_backend.py

@@ -0,0 +1,162 @@
+# 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 re
+
+from llama_index.core import (
+    Document,
+    SimpleDirectoryReader,
+    StorageContext,
+    VectorStoreIndex,
+)
+from llama_index.core.schema import BaseNode
+from llama_index.embeddings.openai_like import OpenAILikeEmbedding
+from pydantic import Field
+from typing_extensions import Any, override
+
+import veadk.config  # noqa E401
+from veadk.configs.database_configs import OpensearchConfig
+from veadk.configs.model_configs import EmbeddingModelConfig, NormalEmbeddingModelConfig
+from veadk.knowledgebase.backends.base_backend import BaseKnowledgebaseBackend
+from veadk.knowledgebase.backends.utils import get_llama_index_splitter
+
+try:
+    from llama_index.vector_stores.opensearch import (
+        OpensearchVectorClient,
+        OpensearchVectorStore,
+    )
+except ImportError:
+    raise ImportError(
+        "Please install VeADK extensions\npip install veadk-python[extensions]"
+    )
+
+
+class OpensearchKnowledgeBackend(BaseKnowledgebaseBackend):
+    """Opensearch-based backend for knowledgebase.
+
+    Opensearch backend stores embedded text in a opensearch database by Llama-index.
+
+    Attributes:
+        opensearch_config (OpensearchConfig):
+            Opensearch database configurations.
+            Mainly contains opensearch host, port, username, password, etc.
+        embedding_config (EmbeddingModelConfig):
+            Embedding config for text embedding and search.
+            Embedding config contains embedding model name and the corresponding dim.
+
+    Examples:
+        Init a knowledgebase based on opensearch backend.
+
+        ```python
+        knowledgebase = Knowledgebase(backend="opensearch")
+        ```
+
+        With more configurations:
+
+        ```python
+        ...
+        ```
+    """
+
+    opensearch_config: OpensearchConfig = Field(default_factory=OpensearchConfig)
+    """Opensearch client configs"""
+
+    embedding_config: EmbeddingModelConfig | NormalEmbeddingModelConfig = Field(
+        default_factory=EmbeddingModelConfig
+    )
+    """Embedding model configs"""
+
+    def model_post_init(self, __context: Any) -> None:
+        self.precheck_index_naming()
+        self._opensearch_client = OpensearchVectorClient(
+            endpoint=self.opensearch_config.host,
+            port=self.opensearch_config.port,
+            http_auth=(
+                self.opensearch_config.username,
+                self.opensearch_config.password,
+            ),
+            use_ssl=True,
+            verify_certs=False,
+            dim=self.embedding_config.dim,
+            index=self.index,  # collection name
+        )
+
+        self._vector_store = OpensearchVectorStore(client=self._opensearch_client)
+
+        self._storage_context = StorageContext.from_defaults(
+            vector_store=self._vector_store
+        )
+
+        self._embed_model = OpenAILikeEmbedding(
+            model_name=self.embedding_config.name,
+            api_key=self.embedding_config.api_key,
+            api_base=self.embedding_config.api_base,
+        )
+
+        self._vector_index = VectorStoreIndex.from_documents(
+            documents=[],
+            storage_context=self._storage_context,
+            embed_model=self._embed_model,
+        )
+
+    @override
+    def precheck_index_naming(self) -> None:
+        if not (
+            isinstance(self.index, str)
+            and not self.index.startswith(("_", "-"))
+            and self.index.islower()
+            and re.match(r"^[a-z0-9_\-.]+$", self.index)
+        ):
+            raise ValueError(
+                "The index name does not conform to the naming rules of OpenSearch"
+            )
+
+    @override
+    def add_from_directory(self, directory: str) -> bool:
+        documents = SimpleDirectoryReader(input_dir=directory).load_data()
+        nodes = self._split_documents(documents)
+        self._vector_index.insert_nodes(nodes)
+        return True
+
+    @override
+    def add_from_files(self, files: list[str]) -> bool:
+        documents = SimpleDirectoryReader(input_files=files).load_data()
+        nodes = self._split_documents(documents)
+        self._vector_index.insert_nodes(nodes)
+        return True
+
+    @override
+    def add_from_text(self, text: str | list[str]) -> bool:
+        if isinstance(text, str):
+            documents = [Document(text=text)]
+        else:
+            documents = [Document(text=t) for t in text]
+        nodes = self._split_documents(documents)
+        self._vector_index.insert_nodes(nodes)
+        return True
+
+    @override
+    def search(self, query: str, top_k: int = 5) -> list[str]:
+        _retriever = self._vector_index.as_retriever(similarity_top_k=top_k)
+        retrieved_nodes = _retriever.retrieve(query)
+        return [node.text for node in retrieved_nodes]
+
+    def _split_documents(self, documents: list[Document]) -> list[BaseNode]:
+        """Split document into chunks"""
+        nodes = []
+        for document in documents:
+            splitter = get_llama_index_splitter(document.metadata.get("file_path", ""))
+            _nodes = splitter.get_nodes_from_documents([document])
+            nodes.extend(_nodes)
+        return nodes