vikingdb-python-sdk 0.1.9__py3-none-any.whl

vikingdb/__init__.py

@@ -0,0 +1,36 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from .auth import APIKey, IAM
+from .request_options import RequestOptions
+from . import vector
+from . import memory
+from .vector import (
+    CollectionClient,
+    EmbeddingClient,
+    IndexClient,
+    RerankClient,
+    VikingVector,
+)
+from .memory import (
+    VikingMem,
+    Collection,
+)
+from .version import __version__
+__all__ = [
+    "IAM",
+    "APIKey",
+    "CollectionClient",
+    "EmbeddingClient",
+    "RerankClient",
+    "IndexClient",
+    "RequestOptions",
+    "VikingVector",
+    "vector",
+    "memory",
+    "VikingMem",
+    "Collection",
+    "__version__",
+]

vikingdb/_client.py

@@ -0,0 +1,262 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Shared base client logic for Viking services."""
+
+from __future__ import annotations
+
+import json
+from abc import ABC, abstractmethod
+from json import JSONDecodeError
+from typing import Any, Mapping, Optional
+
+import aiohttp
+
+from volcengine.ApiInfo import ApiInfo
+from volcengine.ServiceInfo import ServiceInfo
+from volcengine.base.Request import Request
+from volcengine.base.Service import Service
+import requests
+
+from .auth import Auth, IAM, APIKey
+from .exceptions import (
+    DEFAULT_UNKNOWN_ERROR_CODE,
+    VikingAPIException,
+)
+
+
+_REQUEST_ID_HEADER = "X-Tt-Logid"
+
+
+class Client(Service, ABC):
+    """Reusable base client built on top of volcengine Service."""
+
+    def __init__(
+        self,
+        *,
+        host: str,
+        region: str,
+        service: str,
+        auth: Auth,
+        sts_token: str = "",
+        scheme: str = "http",
+        timeout: int = 30,
+    ):
+        self.region = region
+        self.service = service
+        self.auth_provider = auth
+        credentials = auth.initialize(service=service, region=region)
+        self.service_info = self._build_service_info(
+            host=host,
+            credentials=credentials,
+            scheme=scheme,
+            timeout=timeout,
+        )
+        self.api_info = self._build_api_info()
+        # 判断auth是不是IAM 还是 APIKey类型
+        if isinstance(auth, IAM):
+            super().__init__(self.service_info, self.api_info)
+        elif isinstance(auth, APIKey):
+            self.session = requests.session()
+        else:
+            raise ValueError("auth must be IAM or APIKey type")
+
+        if sts_token:
+            self.set_session_token(session_token=sts_token)
+
+    @abstractmethod
+    def _build_api_info(self) -> Mapping[str, ApiInfo]:
+        """Return the API metadata mapping used by this client."""
+
+    @staticmethod
+    def _build_service_info(
+        *,
+        host: str,
+        credentials,
+        scheme: str,
+        timeout: int,
+    ) -> ServiceInfo:
+        return ServiceInfo(
+            host,
+            {},
+            credentials,
+            timeout,
+            timeout,
+            scheme=scheme,
+        )
+
+    def prepare_request(self, api_info: ApiInfo, params: Optional[Mapping[str, Any]], doseq: int = 0):
+        """Prepare a volcengine request without adding implicit headers."""
+        request = Request()
+        request.set_shema(self.service_info.scheme)
+        request.set_method(api_info.method)
+        request.set_host(self.service_info.host)
+        request.set_path(api_info.path)
+        request.set_connection_timeout(self.service_info.connection_timeout)
+        request.set_socket_timeout(self.service_info.socket_timeout)
+        request.set_headers(dict(api_info.header))
+        if params:
+            request.set_query(params)
+        return request
+
+    def _json(
+        self,
+        api: str,
+        params: Optional[Mapping[str, Any]],
+        body: Any,
+        headers: Optional[Mapping[str, str]] = None,
+        timeout: Optional[int] = None,
+    ) -> Any:
+        """Send a JSON request synchronously.
+        
+        Args:
+            api: API name
+            params: Query parameters
+            body: Request body
+            headers: Additional headers
+            timeout: Timeout in seconds (optional). If not provided, uses default connection_timeout and socket_timeout.
+        """
+        if api not in self.api_info:
+            raise Exception("no such api")
+        api_info = self.api_info[api]
+        request = self.prepare_request(api_info, params)
+        if headers:
+            for key, value in headers.items():
+                request.headers[key] = value
+        request.headers["Content-Type"] = "application/json"
+        request.body = body
+        self.auth_provider.sign_request(request)
+        url = request.build()
+        
+        request_id_value = request.headers.get(_REQUEST_ID_HEADER)
+        request_id = str(request_id_value) if request_id_value else "unknown"
+        
+        # Use custom timeout if provided, otherwise use default
+        if timeout is not None:
+            request_timeout = (timeout, timeout)
+        else:
+            request_timeout = (
+                self.service_info.connection_timeout,
+                self.service_info.socket_timeout,
+            )
+        
+        try:
+            response = self.session.post(
+                url,
+                headers=request.headers,
+                data=request.body,
+                timeout=request_timeout,
+            )
+        except Exception as exc:
+            raise VikingAPIException(
+                    DEFAULT_UNKNOWN_ERROR_CODE,
+                    request_id=request_id,
+                    message=f"failed to run session.post {api}: {exc}",
+                ) from exc
+
+        payload_text_attr = getattr(response, "text", "")
+        payload_text = payload_text_attr if isinstance(payload_text_attr, str) else ""
+        payload_text = payload_text or ""
+
+        if response.status_code != 200:
+            error = VikingAPIException.from_response(
+                payload_text,
+                request_id=request_id,
+                status_code=response.status_code,
+            )
+            raise error
+
+        try:
+            return response.json()
+        except (ValueError, JSONDecodeError) as exc:
+            raise VikingAPIException(
+                DEFAULT_UNKNOWN_ERROR_CODE,
+                request_id=request_id,
+                message=f"failed to decode JSON response for {api}: {exc}",
+                status_code=response.status_code,
+            ) from exc
+
+        except Exception as exc:  # pragma: no cover - defensive fallback
+            raise VikingAPIException(
+                DEFAULT_UNKNOWN_ERROR_CODE,
+                request_id=request_id,
+                message=f"unexpected error parsing response for {api}: {exc}",
+                status_code=response.status_code,
+            ) from exc
+
+    async def async_json(
+        self,
+        api: str,
+        params: Optional[Mapping[str, Any]],
+        body: Any,
+        headers: Optional[Mapping[str, str]] = None,
+        timeout: Optional[int] = None,
+    ) -> Any:
+        """Send a JSON request asynchronously.
+        
+        Args:
+            api: API name
+            params: Query parameters
+            body: Request body
+            headers: Additional headers
+            timeout: Timeout in seconds (optional). If not provided, uses default connection_timeout and socket_timeout.
+        """
+        if api not in self.api_info:
+            raise Exception("no such api")
+        api_info = self.api_info[api]
+        request = self.prepare_request(api_info, params)
+        if headers:
+            for key, value in headers.items():
+                request.headers[key] = value
+        request.headers["Content-Type"] = "application/json"
+        request.body = body
+
+        self.auth_provider.sign_request(request)
+        
+        # Use custom timeout if provided, otherwise use default
+        if timeout is not None:
+            client_timeout = aiohttp.ClientTimeout(
+                connect=timeout,
+                sock_connect=timeout,
+                sock_read=timeout,
+            )
+        else:
+            client_timeout = aiohttp.ClientTimeout(
+                connect=self.service_info.connection_timeout,
+                sock_connect=self.service_info.socket_timeout,
+            )
+        
+        url = request.build()
+        try:
+            async with aiohttp.request(
+                "POST",
+                url,
+                headers=request.headers,
+                data=request.body,
+                timeout=client_timeout,
+            ) as response:
+                request_id_value = response.headers.get(_REQUEST_ID_HEADER)
+                request_id = str(request_id_value) if request_id_value else "unknown"
+                payload = await response.text(encoding="utf-8")
+                if response.status != 200:
+                    error = VikingAPIException.from_response(
+                        payload,
+                        request_id=request_id,
+                        status_code=response.status,
+                    )
+                    raise error
+                try:
+                    return json.loads(payload)
+                except JSONDecodeError as exc:
+                    raise VikingAPIException(
+                        DEFAULT_UNKNOWN_ERROR_CODE,
+                        request_id=request_id,
+                        message=f"failed to decode JSON response for {api}: {exc}",
+                        status_code=response.status,
+                    ) from exc
+        except Exception as exc:
+            raise VikingAPIException(
+                    DEFAULT_UNKNOWN_ERROR_CODE,
+                    request_id=request_id,
+                    message=f"failed to run aiohttp {api}: {exc}",
+                ) from exc

vikingdb/auth.py

@@ -0,0 +1,68 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Common authentication helpers shared by Viking services."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from volcengine.Credentials import Credentials
+from volcengine.auth.SignerV4 import SignerV4
+
+
+class Auth(ABC):
+    """Base class for authentication providers compatible with volcengine."""
+
+    @abstractmethod
+    def initialize(self, *, service: str, region: str):
+        """Prepare the provider for the given service/region. Returns credentials understood by ServiceInfo."""
+
+    @abstractmethod
+    def sign_request(self, request) -> None:
+        """Sign or otherwise authorise the outgoing request."""
+
+
+class IAM(Auth):
+    """IAM-style AK/SK signature authentication provider."""
+
+    def __init__(self, *, ak: str, sk: str):
+        if not ak or not sk:
+            raise ValueError("ak and sk must be provided for IAM authentication")
+        self._ak = ak
+        self._sk = sk
+        self._credentials = None
+        self._service = None
+        self._region = None
+
+    def initialize(self, *, service: str, region: str):
+        if not service or not region:
+            raise ValueError("service and region must be provided for IAM credentials")
+        if self._credentials is not None and (service != self._service or region != self._region):
+            raise ValueError(
+                f"IAM credentials already initialised for service={self._service!r}, region={self._region!r}"
+            )
+        self._service = service
+        self._region = region
+        if self._credentials is None:
+            self._credentials = Credentials(self._ak, self._sk, service, region)
+        return self._credentials
+
+    def sign_request(self, request) -> None:
+        if self._credentials is None:
+            raise ValueError("IAM provider must be initialised before signing requests")
+        SignerV4.sign(request, self._credentials)
+
+
+class APIKey(Auth):
+    """API Key authentication provider (placeholder for future support)."""
+
+    def __init__(self, *, api_key: str):
+        if not api_key:
+            raise ValueError("api_key must be provided for API key authentication")
+        self.api_key = api_key
+
+    def initialize(self, *, service: str, region: str):
+        return None
+
+    def sign_request(self, request) -> None:
+        request.headers["Authorization"] = f"Bearer {self.api_key}"

vikingdb/exceptions.py

@@ -0,0 +1,160 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from typing import Any, Mapping, Optional, Type, TypeVar, Union
+
+DEFAULT_UNKNOWN_ERROR_CODE: Union[int, str] = 1000028
+NETWORK_ERROR_CODE = 1001
+
+
+@dataclass
+class ParsedError:
+    code: Union[int, str]
+    request_id: str
+    message: Optional[str]
+    payload: Any
+    raw: Optional[str] = None
+
+
+def _normalize_code(value: Any, default: Union[int, str]) -> Union[int, str]:
+    if value is None:
+        return default
+    if isinstance(value, int):
+        return value
+    if isinstance(value, str):
+        stripped = value.strip()
+        return stripped or default
+    try:
+        return int(value)
+    except (TypeError, ValueError):
+        return default
+
+
+def parse_error_payload(payload: Any) -> ParsedError:
+    """
+    Parse a server error payload into a consistent structure.
+    
+    Supports both the legacy {"ResponseMetadata": {"Error": {...}}} format and the
+    newer flat {"code": ..., "request_id": ..., "message": ...} shape. When JSON
+    decoding fails, the raw text is preserved as the message.
+    """
+    code = DEFAULT_UNKNOWN_ERROR_CODE
+    request_id = "unknown"
+    message: Optional[str] = None
+    raw_text: Optional[str] = None
+    parsed_payload: Any = payload
+
+    if isinstance(payload, bytes):
+        raw_text = payload.decode("utf-8", errors="replace")
+    elif isinstance(payload, str):
+        raw_text = payload
+
+    if raw_text is not None:
+        try:
+            parsed_payload = json.loads(raw_text)
+        except json.JSONDecodeError:
+            message = raw_text.strip() or None
+            return ParsedError(code, request_id, message, raw_text, raw_text)
+
+    if isinstance(parsed_payload, Mapping):
+        metadata = parsed_payload.get("ResponseMetadata")
+        if isinstance(metadata, Mapping):
+            error = metadata.get("Error")
+            if isinstance(error, Mapping):
+                code_value = error.get("Code") if "Code" in error else error.get("code")
+                code = _normalize_code(code_value, code)
+                request_id = str(error.get("RequestId") or request_id)
+                if error.get("Message"):
+                    message = str(error["Message"])
+            if metadata.get("RequestId"):
+                request_id = str(metadata["RequestId"])
+        else:
+            code_value = parsed_payload.get("code")
+            if code_value is None and "Code" in parsed_payload:
+                code_value = parsed_payload.get("Code")
+            code = _normalize_code(code_value, code)
+            if parsed_payload.get("request_id"):
+                request_id = str(parsed_payload["request_id"])
+            if parsed_payload.get("message"):
+                message = str(parsed_payload["message"])
+    else:
+        if message is None and raw_text is None:
+            message = str(parsed_payload)
+
+    if raw_text is None and isinstance(parsed_payload, (dict, list)):
+        try:
+            raw_text = json.dumps(parsed_payload, ensure_ascii=False)
+        except Exception:
+            raw_text = None
+
+    return ParsedError(code, request_id, message, parsed_payload, raw_text)
+
+
+T_VikingException = TypeVar("T_VikingException", bound="VikingException")
+
+
+class VikingException(Exception):
+    """
+    Base exception for all Viking SDK errors.
+
+    Captures standard fields returned by the service for consistent error handling.
+    """
+
+    def __init__(
+        self,
+        code: Union[int, str],
+        request_id: str = "unknown",
+        message: Optional[str] = None,
+        *,
+        status_code: Optional[int] = None,
+    ) -> None:
+        self.code = code
+        self.request_id = request_id or "unknown"
+        self.status_code = status_code
+        self.message = message or f"request failed (code={self.code})"
+        super().__init__(self.message)
+
+    def __str__(self) -> str:
+        status = f", http_status={self.status_code}" if self.status_code is not None else ""
+        return f"{self.message} (code={self.code}, request_id={self.request_id}{status})"
+
+    def promote(self, target_cls: Type[T_VikingException]) -> T_VikingException:
+        """
+        Promote this exception to a more specific subclass, reusing captured context.
+        """
+        if isinstance(self, target_cls):
+            return self  # type: ignore[return-value]
+        return target_cls(
+            self.code,
+            self.request_id,
+            self.message,
+            status_code=self.status_code,
+        )
+
+
+class VikingAPIException(VikingException):
+    """Raised when the remote API returns an error payload."""
+
+    @classmethod
+    def from_response(cls, payload: Any, *, request_id: Optional[str] = "unknown", status_code: Optional[int] = None) -> "VikingAPIException":
+        parsed = parse_error_payload(payload)
+        return cls(
+            parsed.code,
+            parsed.request_id or request_id,
+            parsed.message or "unknown api error",
+            status_code=status_code,
+        )
+
+
+def promote_exception(
+    exc: VikingException,
+    *,
+    exception_map: Optional[Mapping[int, Type[T_VikingException]]] = None,
+    default_cls: Type[T_VikingException],
+) -> T_VikingException:
+    """
+    Promote a VikingException to a service-specific subclass using the provided map.
+    """
+    target_cls = (exception_map or {}).get(exc.code, default_cls)
+    return exc.promote(target_cls)

vikingdb/memory/__init__.py

@@ -0,0 +1,70 @@
+# coding:utf-8
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Viking Memory SDK
+
+Provides memory management features including:
+- Collection management (create, delete, update, query)
+- Profile management (add, delete, update)
+- Event management (add, update, delete, batch delete)
+- Session management (add messages)
+- Memory search (semantic search)
+
+All APIs return plain dictionaries (dict) or lists of dictionaries (list[dict]) without object encapsulation.
+
+Authentication support:
+- AK/SK signature authentication (standard volcengine SignerV4)
+- API Key authentication (reserved interface)
+"""
+
+from .client import VikingMem
+from .collection import Collection
+from .exceptions import (
+    VikingMemException,
+    UnauthorizedException,
+    InvalidRequestException,
+    CollectionExistException,
+    CollectionNotExistException,
+    IndexExistException,
+    IndexNotExistException,
+    DataNotFoundException,
+    DelOpFailedException,
+    UpsertOpFailedException,
+    InvalidVectorException,
+    InvalidPrimaryKeyException,
+    InvalidFilterException,
+    IndexSearchException,
+    IndexFetchException,
+    IndexInitializingException,
+    EmbeddingException,
+    InternalServerException,
+    QuotaLimiterException,
+)
+
+__all__ = [
+    # Main client classes
+    "VikingMem",
+    "Collection",
+    # Exception classes
+    "VikingMemException",
+    "UnauthorizedException",
+    "InvalidRequestException",
+    "CollectionExistException",
+    "CollectionNotExistException",
+    "IndexExistException",
+    "IndexNotExistException",
+    "DataNotFoundException",
+    "DelOpFailedException",
+    "UpsertOpFailedException",
+    "InvalidVectorException",
+    "InvalidPrimaryKeyException",
+    "InvalidFilterException",
+    "IndexSearchException",
+    "IndexFetchException",
+    "IndexInitializingException",
+    "EmbeddingException",
+    "InternalServerException",
+    "QuotaLimiterException",
+]