From 65814a464433644e01d01200847c5d6030dd06c2 Mon Sep 17 00:00:00 2001 From: Junyan Qin Date: Sun, 13 Jul 2025 21:39:33 +0800 Subject: [PATCH] feat: binary storage api --- pkg/entity/persistence/bstorage.py | 22 +++++++ pkg/plugin/handler.py | 102 +++++++++++++++++++++++++++++ 2 files changed, 124 insertions(+) create mode 100644 pkg/entity/persistence/bstorage.py diff --git a/pkg/entity/persistence/bstorage.py b/pkg/entity/persistence/bstorage.py new file mode 100644 index 00000000..674dee29 --- /dev/null +++ b/pkg/entity/persistence/bstorage.py @@ -0,0 +1,22 @@ +import sqlalchemy + +from .base import Base + + +class BinaryStorage(Base): + """Current for plugin use only""" + + __tablename__ = 'binary_storages' + + unique_key = sqlalchemy.Column(sqlalchemy.String(255), primary_key=True) + key = sqlalchemy.Column(sqlalchemy.String(255), nullable=False) + owner_type = sqlalchemy.Column(sqlalchemy.String(255), nullable=False) + owner = sqlalchemy.Column(sqlalchemy.String(255), nullable=False) + value = sqlalchemy.Column(sqlalchemy.LargeBinary, nullable=False) + created_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=False, server_default=sqlalchemy.func.now()) + updated_at = sqlalchemy.Column( + sqlalchemy.DateTime, + nullable=False, + server_default=sqlalchemy.func.now(), + onupdate=sqlalchemy.func.now(), + ) diff --git a/pkg/plugin/handler.py b/pkg/plugin/handler.py index 8da1bc7f..d0905f1f 100644 --- a/pkg/plugin/handler.py +++ b/pkg/plugin/handler.py @@ -2,6 +2,7 @@ from __future__ import annotations import typing from typing import Any +import base64 import sqlalchemy @@ -18,6 +19,7 @@ import langbot_plugin.api.entities.builtin.provider.message as provider_message import langbot_plugin.api.entities.builtin.resource.tool as resource_tool from ..entity.persistence import plugin as persistence_plugin +from ..entity.persistence import bstorage as persistence_bstorage from ..core import app from ..utils import constants @@ -252,6 +254,106 @@ class RuntimeConnectionHandler(handler.Handler): }, ) + @self.action(RuntimeToLangBotAction.SET_BINARY_STORAGE) + async def set_binary_storage(data: dict[str, Any]) -> handler.ActionResponse: + """Set binary storage""" + key = data['key'] + owner_type = data['owner_type'] + owner = data['owner'] + value = base64.b64decode(data['value_base64']) + + result = await self.ap.persistence_mgr.execute_async( + sqlalchemy.select(persistence_bstorage.BinaryStorage) + .where(persistence_bstorage.BinaryStorage.key == key) + .where(persistence_bstorage.BinaryStorage.owner_type == owner_type) + .where(persistence_bstorage.BinaryStorage.owner == owner) + ) + + if result.first() is not None: + await self.ap.persistence_mgr.execute_async( + sqlalchemy.update(persistence_bstorage.BinaryStorage) + .where(persistence_bstorage.BinaryStorage.key == key) + .where(persistence_bstorage.BinaryStorage.owner_type == owner_type) + .where(persistence_bstorage.BinaryStorage.owner == owner) + .values(value=value) + ) + else: + await self.ap.persistence_mgr.execute_async( + sqlalchemy.insert(persistence_bstorage.BinaryStorage).values( + unique_key=f'{owner_type}:{owner}:{key}', + key=key, + owner_type=owner_type, + owner=owner, + value=value, + ) + ) + + return handler.ActionResponse.success( + data={}, + ) + + @self.action(RuntimeToLangBotAction.GET_BINARY_STORAGE) + async def get_binary_storage(data: dict[str, Any]) -> handler.ActionResponse: + """Get binary storage""" + key = data['key'] + owner_type = data['owner_type'] + owner = data['owner'] + + result = await self.ap.persistence_mgr.execute_async( + sqlalchemy.select(persistence_bstorage.BinaryStorage) + .where(persistence_bstorage.BinaryStorage.key == key) + .where(persistence_bstorage.BinaryStorage.owner_type == owner_type) + .where(persistence_bstorage.BinaryStorage.owner == owner) + ) + + storage = result.first() + if storage is None: + return handler.ActionResponse.error( + message=f'Storage with key {key} not found', + ) + + return handler.ActionResponse.success( + data={ + 'value_base64': base64.b64encode(storage.value).decode('utf-8'), + }, + ) + + @self.action(RuntimeToLangBotAction.DELETE_BINARY_STORAGE) + async def delete_binary_storage(data: dict[str, Any]) -> handler.ActionResponse: + """Delete binary storage""" + key = data['key'] + owner_type = data['owner_type'] + owner = data['owner'] + + await self.ap.persistence_mgr.execute_async( + sqlalchemy.delete(persistence_bstorage.BinaryStorage) + .where(persistence_bstorage.BinaryStorage.key == key) + .where(persistence_bstorage.BinaryStorage.owner_type == owner_type) + .where(persistence_bstorage.BinaryStorage.owner == owner) + ) + + return handler.ActionResponse.success( + data={}, + ) + + @self.action(RuntimeToLangBotAction.GET_BINARY_STORAGE_KEYS) + async def get_binary_storage_keys(data: dict[str, Any]) -> handler.ActionResponse: + """Get binary storage keys""" + owner_type = data['owner_type'] + owner = data['owner'] + + result = await self.ap.persistence_mgr.execute_async( + sqlalchemy.select(persistence_bstorage.BinaryStorage.key) + .where(persistence_bstorage.BinaryStorage.owner_type == owner_type) + .where(persistence_bstorage.BinaryStorage.owner == owner) + ) + + return handler.ActionResponse.success( + data={ + 'keys': result.scalars().all(), + }, + ) + async def ping(self) -> dict[str, Any]: """Ping the runtime""" return await self.call_action(