feat: add user phonon balance check callback to ensure sufficient balance before job creation

Author: Rasic2

agents/matmaster_agent/base_agents/callback.py

@@ -4,7 +4,7 @@
 import os
 import traceback
 from functools import wraps
-from typing import Optional, Union
+from typing import Callable, Optional, Union
 
 import aiohttp
 import litellm
@@ -27,6 +27,7 @@
 )
 from agents.matmaster_agent.prompt import get_params_check_info_prompt
 from agents.matmaster_agent.utils.auth import ak_to_ticket, ak_to_username
+from agents.matmaster_agent.utils.finance import get_user_photon_balance
 from agents.matmaster_agent.utils.helper_func import (
     check_None_wrapper,
     function_calls_to_str,
@@ -164,6 +165,100 @@ async def default_before_tool_callback(tool, args, tool_context):
     return
 
 
+def default_cost_func() -> int:
+    return 300
+
+
+def check_user_phonon_balance(
+    func: BeforeToolCallback, cost_func: Callable[[], int] = default_cost_func
+) -> BeforeToolCallback:
+    @wraps(func)
+    async def wrapper(
+        tool: BaseTool, args: dict, tool_context: ToolContext
+    ) -> Optional[dict]:
+        # 两步操作：
+        # 1. 调用被装饰的 before_tool_callback；
+        # 2. 如果调用的 before_tool_callback 有返回值，以这个为准
+        if (before_tool_result := await func(tool, args, tool_context)) is not None:
+            return before_tool_result
+
+        # 如果 tool 不是 CalculationMCPTool，不应该调用这个 callback
+        if not isinstance(tool, CalculationMCPTool):
+            raise TypeError(
+                "Not CalculationMCPTool type, current tool can't create job!"
+            )
+
+        user_id = _get_user_id(tool_context)
+        cost = cost_func()
+        balance = await get_user_photon_balance(user_id)
+
+        logger.info(
+            f"[check_user_phonon_balance] user_id={user_id}, cost={cost}, balance={balance}"
+        )
+        if balance < cost:
+            raise RuntimeError('Phonon is not enough, Please recharge.')
+
+    return wrapper
+
+
+def check_job_create(func: BeforeToolCallback) -> BeforeToolCallback:
+    @wraps(func)
+    async def wrapper(
+        tool: BaseTool, args: dict, tool_context: ToolContext
+    ) -> Optional[dict]:
+        # 两步操作：
+        # 1. 调用被装饰的 before_tool_callback；
+        # 2. 如果调用的 before_tool_callback 有返回值，以这个为准
+        if (before_tool_result := await func(tool, args, tool_context)) is not None:
+            return before_tool_result
+
+        # 如果 tool 不是 CalculationMCPTool，不应该调用这个 callback
+        if not isinstance(tool, CalculationMCPTool):
+            raise TypeError(
+                "Not CalculationMCPTool type, current tool can't create job!"
+            )
+
+        if tool.executor is not None:
+            job_create_url = f"{OPENAPI_HOST}/openapi/v1/job/create"
+            user_project_list_url = f"{OPENAPI_HOST}/openapi/v1/open/user/project/list"
+            payload = {
+                'projectId': MATERIALS_PROJECT_ID,
+                'name': 'check_job_create',
+            }
+            params = {'accessKey': MATERIALS_ACCESS_KEY}
+
+            logger.info(
+                f"[check_job_create] project_id = {MATERIALS_PROJECT_ID}, "
+                f"ak = {MATERIALS_ACCESS_KEY}"
+            )
+
+            async with aiohttp.ClientSession() as session:
+                async with session.get(
+                    user_project_list_url, params=params
+                ) as response:
+                    res = json.loads(await response.text())
+                    logger.info(f"[check_job_create] res = {res}")
+                    project_name = [
+                        item['project_name']
+                        for item in res['data']['items']
+                        if item['project_id'] == MATERIALS_PROJECT_ID
+                    ][0]
+
+            async with aiohttp.ClientSession() as session:
+                async with session.post(
+                    job_create_url, json=payload, params=params
+                ) as response:
+                    res = json.loads(await response.text())
+                    if res['code'] != 0:
+                        if res['code'] == 2000:
+                            res['error'][
+                                'msg'
+                            ] = f"您所用项目为 `{project_name}`，该项目余额不足，请充值或更换项目后重试。"
+                        return res
+
+    return wrapper
+
+
 @check_None_wrapper
 def _get_ak(ctx: Union[InvocationContext, ToolContext, CallbackContext]):
     session_state = get_session_state(ctx)
@@ -180,6 +275,14 @@ def _get_projectId(ctx: Union[InvocationContext, ToolContext]):
     )
 
 
+@check_None_wrapper
+def _get_user_id(ctx: Union[InvocationContext, ToolContext]):
+    session_state = get_session_state(ctx)
+    return session_state[FRONTEND_STATE_KEY].get('adk_user_id') or os.getenv(
+        'MATERIALS_USER_ID'
+    )
+
+
 def _inject_ak(ctx: Union[InvocationContext, ToolContext], executor, storage):
     access_key = _get_ak(ctx)
     if executor is not None:
@@ -295,64 +398,6 @@ async def wrapper(
     return wrapper
 
 
-def check_job_create(func: BeforeToolCallback) -> BeforeToolCallback:
-    @wraps(func)
-    async def wrapper(
-        tool: BaseTool, args: dict, tool_context: ToolContext
-    ) -> Optional[dict]:
-        # 两步操作：
-        # 1. 调用被装饰的 before_tool_callback；
-        # 2. 如果调用的 before_tool_callback 有返回值，以这个为准
-        if (before_tool_result := await func(tool, args, tool_context)) is not None:
-            return before_tool_result
-
-        # 如果 tool 不是 CalculationMCPTool，不应该调用这个 callback
-        if not isinstance(tool, CalculationMCPTool):
-            raise TypeError(
-                "Not CalculationMCPTool type, current tool can't create job!"
-            )
-
-        if tool.executor is not None:
-            job_create_url = f"{OPENAPI_HOST}/openapi/v1/job/create"
-            user_project_list_url = f"{OPENAPI_HOST}/openapi/v1/open/user/project/list"
-            payload = {
-                'projectId': MATERIALS_PROJECT_ID,
-                'name': 'check_job_create',
-            }
-            params = {'accessKey': MATERIALS_ACCESS_KEY}
-
-            logger.info(
-                f"[check_job_create] project_id = {MATERIALS_PROJECT_ID}, "
-                f"ak = {MATERIALS_ACCESS_KEY}"
-            )
-
-            async with aiohttp.ClientSession() as session:
-                async with session.get(
-                    user_project_list_url, params=params
-                ) as response:
-                    res = json.loads(await response.text())
-                    logger.info(f"[check_job_create] res = {res}")
-                    project_name = [
-                        item['project_name']
-                        for item in res['data']['items']
-                        if item['project_id'] == MATERIALS_PROJECT_ID
-                    ][0]
-
-            async with aiohttp.ClientSession() as session:
-                async with session.post(
-                    job_create_url, json=payload, params=params
-                ) as response:
-                    res = json.loads(await response.text())
-                    if res['code'] != 0:
-                        if res['code'] == 2000:
-                            res['error'][
-                                'msg'
-                            ] = f"您所用项目为 `{project_name}`，该项目余额不足，请充值或更换项目后重试。"
-                        return res
-
-    return wrapper
-
-
 # 总应该在最后
 def catch_before_tool_callback_error(func: BeforeToolCallback) -> BeforeToolCallback:
     @wraps(func)

agents/matmaster_agent/base_agents/job_agent.py

@@ -17,6 +17,7 @@
     catch_before_tool_callback_error,
     check_before_tool_callback_effect,
     check_job_create,
+    check_user_phonon_balance,
     default_after_model_callback,
     default_after_tool_callback,
     default_before_tool_callback,
@@ -161,7 +162,9 @@ def __init__(
         # Todo: support List[before_tool_callback]
         before_tool_callback = catch_before_tool_callback_error(
             inject_current_env(
-                inject_username_ticket(check_job_create(before_tool_callback))
+                inject_username_ticket(
+                    check_job_create(check_user_phonon_balance(before_tool_callback))
+                )
             )
         )
         after_tool_callback = check_before_tool_callback_effect(
@@ -359,6 +362,8 @@ async def _run_async_impl(
                 ):
                     try:
                         dict_result = load_tool_response(event)
+                        if dict_result.get('status', None) == 'error':
+                            raise eval(dict_result['error_type'])(dict_result['error'])
                     except BaseException:
                         yield event
                         raise

agents/matmaster_agent/constant.py

@@ -47,8 +47,8 @@
 FINANCE_INFO_API = f"{FINANCE_API_BASE}/api/integral/info"
 FINANCE_CONSUME_API = f"{FINANCE_API_BASE}/api/integral/consume"
 
-MATERIALS_USER_ID = int(os.getenv('MATERIALS_USER_ID'), -1)
-MATERIALS_ORG_ID = int(os.getenv('MATERIALS_ORG_ID'), -1)
+MATERIALS_USER_ID = int(os.getenv('MATERIALS_USER_ID', -1))
+MATERIALS_ORG_ID = int(os.getenv('MATERIALS_ORG_ID', -1))
 MATERIALS_ACCESS_KEY = str(os.getenv('MATERIALS_ACCESS_KEY'))
 MATERIALS_PROJECT_ID = int(os.getenv('MATERIALS_PROJECT_ID'))
 

agents/matmaster_agent/utils/finance.py

@@ -0,0 +1,22 @@
+import logging
+
+import aiohttp
+
+from agents.matmaster_agent.constant import (
+    FINANCE_INFO_API,
+)
+
+logger = logging.getLogger(__name__)
+
+
+async def get_user_photon_balance(user_id):
+    payload = {'userId': int(user_id)}
+    async with aiohttp.ClientSession() as session:
+        async with session.post(
+            FINANCE_INFO_API, json=payload, headers={'Content-Type': 'application/json'}
+        ) as response:
+            response.raise_for_status()
+            res = await response.json()
+
+            logger.info(f'[get_user_photon_balance] balance = {res['data']['balance']}')
+            return res['data']['balance']