diff --git a/pyproject.toml b/pyproject.toml index 760ed086..d245c51c 100644 --- a/pyproject.toml +++ b/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 = [] diff --git a/veadk/knowledgebase/backends/utils.py b/veadk/knowledgebase/backends/utils.py index abb38643..fd9ee78f 100644 --- a/veadk/knowledgebase/backends/utils.py +++ b/veadk/knowledgebase/backends/utils.py @@ -12,9 +12,7 @@ # 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, @@ -22,9 +20,6 @@ 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 diff --git a/veadk/knowledgebase/backends/vikingdb_knowledge_backend.py b/veadk/knowledgebase/backends/vikingdb_knowledge_backend.py index cab84b87..986d2b96 100644 --- a/veadk/knowledgebase/backends/vikingdb_knowledge_backend.py +++ b/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: