fix: adjustment of viking dependencies (#489)

Author: zakahan

pyproject.toml

@@ -34,9 +34,10 @@ dependencies = [
     "pymysql==1.1.1", # For MySQL database (short term memory)
     "aiomysql==0.3.2", # For async MySQL database (short term memory)
     "filetype==1.2.0",
-    "vikingdb-python-sdk>=0.1.3",
+    "vikingdb-python-sdk>=0.1.3", # For Viking DB
     "agentkit-sdk-python>=0.2.0",
     "python-frontmatter==1.1.0",
+    "tos>=2.8.4",   # For TOS storage and Viking DB
 ]
 
 [project.scripts]
@@ -45,7 +46,6 @@ veadk = "veadk.cli.cli:veadk"
 [project.optional-dependencies]
 extensions = [
     "redis>=5.0",                                           # For Redis database
-    "tos>=2.8.4",                                           # For TOS storage and Viking DB
     "cozeloop>=0.1.21",                                     # For Cozeloop Prompt manager
     "llama-index==0.14.0",                                  # For KnowledgeBase and LongTermMemory
     "llama-index-embeddings-openai-like==0.2.2",            # For Embeddings
@@ -58,7 +58,6 @@ database = [
     "redis>=5.0",                   # For Redis database
     "pymysql>=1.1.1",               # For MySQL database
     "volcengine>=1.0.193",          # For Viking DB
-    "tos>=2.8.4",                   # For TOS storage and Viking DB
     "mem0ai==0.1.118",              # For mem0
 ]
 speech = []

veadk/knowledgebase/backends/utils.py

@@ -12,19 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import json
 from pathlib import Path
-from typing import Literal
 
 from llama_index.core.node_parser import (
     CodeSplitter,
     HTMLNodeParser,
     MarkdownNodeParser,
     SentenceSplitter,
 )
-from volcengine.auth.SignerV4 import SignerV4
-from volcengine.base.Request import Request
-from volcengine.Credentials import Credentials
 
 
 def get_llama_index_splitter(
@@ -40,53 +35,3 @@ def get_llama_index_splitter(
         return HTMLNodeParser()
     else:
         return SentenceSplitter(chunk_size=512, chunk_overlap=50)
-
-
-def build_vikingdb_knowledgebase_request(
-    path: str,
-    volcengine_access_key: str,
-    volcengine_secret_key: str,
-    session_token: str = "",
-    method: Literal["GET", "POST", "PUT", "DELETE"] = "POST",
-    region: str = "cn-beijing",
-    params=None,
-    data=None,
-    doseq=0,
-) -> Request:
-    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(
-        volcengine_access_key, volcengine_secret_key, "air", region, session_token
-    )
-    SignerV4.sign(r, credentials)
-    return r

veadk/knowledgebase/backends/vikingdb_knowledge_backend.py

@@ -23,6 +23,9 @@
 from pydantic import Field
 from typing_extensions import override
 from volcengine.viking_knowledgebase import VikingKnowledgeBaseService
+from volcengine.auth.SignerV4 import SignerV4
+from volcengine.base.Request import Request
+from volcengine.Credentials import Credentials
 
 import veadk.config  # noqa E401
 from veadk.auth.veauth.utils import (
@@ -31,23 +34,65 @@
 )
 from veadk.configs.database_configs import NormalTOSConfig, TOSConfig
 from veadk.knowledgebase.backends.base_backend import BaseKnowledgebaseBackend
-from veadk.knowledgebase.backends.utils import (
-    build_vikingdb_knowledgebase_request,
-)
 from veadk.knowledgebase.entry import KnowledgebaseEntry
 from veadk.utils.logger import get_logger
 from veadk.utils.misc import formatted_timestamp, getenv
+from veadk.integrations.ve_tos.ve_tos import VeTOS
 
-try:
-    from veadk.integrations.ve_tos.ve_tos import VeTOS
-except ImportError:
-    raise ImportError(
-        "Please install VeADK extensions\npip install veadk-python[extensions]"
-    )
 
 logger = get_logger(__name__)
 
 
+def build_vikingdb_knowledgebase_request(
+    path: str,
+    volcengine_access_key: str,
+    volcengine_secret_key: str,
+    session_token: str = "",
+    method: Literal["GET", "POST", "PUT", "DELETE"] = "POST",
+    region: str = "cn-beijing",
+    params=None,
+    data=None,
+    doseq=0,
+) -> Request:
+    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(
+        volcengine_access_key, volcengine_secret_key, "air", region, session_token
+    )
+    SignerV4.sign(r, credentials)
+    return r
+
+
 def _read_file_to_bytes(file_path: str) -> tuple[bytes, str]:
     """Read file content to bytes, and file name"""
     with open(file_path, "rb") as f: