Merge pull request #628 from AnguseZhang/fix/project-error

feat: add async capability to project ID retrieval and logging enhanc…

Author: Rasic2

agents/matmaster_agent/base_callbacks/private_callback.py

@@ -450,8 +450,10 @@ def _inject_ak(ctx: Union[InvocationContext, ToolContext], executor, storage):
     return access_key, executor, storage
 
 
-def _inject_projectId(ctx: Union[InvocationContext, ToolContext], executor, storage):
-    project_id = _get_projectId(ctx)
+async def _inject_projectId(
+    ctx: Union[InvocationContext, ToolContext], executor, storage
+):
+    project_id = await _get_projectId(ctx)
     if executor is not None:
         if executor['type'] == 'dispatcher':  # BohriumExecutor
             executor['machine']['remote_profile']['project_id'] = int(project_id)
@@ -571,7 +573,7 @@ async def wrapper(
 
         # 获取 project_id
         try:
-            project_id, tool.executor, tool.storage = _inject_projectId(
+            project_id, tool.executor, tool.storage = await _inject_projectId(
                 tool_context, tool.executor, tool.storage
             )
         except ValueError as e:

agents/matmaster_agent/core_agents/public_agents/job_agents/result_core_agent/agent.py

@@ -73,7 +73,7 @@ async def _run_events(self, ctx: InvocationContext) -> AsyncGenerator[Event, Non
             access_key, Executor, BohriumStorge = _inject_ak(
                 ctx, get_BohriumExecutor(), get_BohriumStorage()
             )
-            project_id, Executor, BohriumStorge = _inject_projectId(
+            project_id, Executor, BohriumStorge = await _inject_projectId(
                 ctx, Executor, BohriumStorge
             )
         else:

agents/matmaster_agent/services/project.py

@@ -0,0 +1,25 @@
+import asyncio
+import json
+
+import aiohttp
+
+from agents.matmaster_agent.constant import OPENAPI_HOST
+
+
+async def get_project_list(access_key: str):
+    user_project_list_url = f"{OPENAPI_HOST}/openapi/v1/open/user/project/list"
+    params = {'accessKey': 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())
+            project_list = res.get('data', {}).get('items', [])
+
+            if project_list:
+                return [item['project_id'] for item in project_list]
+            else:
+                return project_list
+
+
+if __name__ == '__main__':
+    result = asyncio.run(get_project_list('ba5cdc51b4934d1689e38332b27bb271'))

agents/matmaster_agent/utils/callback_utils.py

@@ -1,3 +1,4 @@
+import logging
 import os
 from typing import Union
 
@@ -10,20 +11,35 @@
     FRONTEND_STATE_KEY,
     MATERIALS_ACCESS_KEY,
     MATERIALS_PROJECT_ID,
+    MATMASTER_AGENT_NAME,
 )
+from agents.matmaster_agent.logger import PrefixFilter
+from agents.matmaster_agent.services.project import get_project_list
 from agents.matmaster_agent.utils.helper_func import (
     check_None_wrapper,
     get_session_state,
 )
 
+logger = logging.getLogger(__name__)
+logger.addFilter(PrefixFilter(MATMASTER_AGENT_NAME))
+logger.setLevel(logging.INFO)
+
 
 @check_None_wrapper
-def _get_projectId(ctx: Union[InvocationContext, ToolContext]):
+async def _get_projectId(ctx: Union[InvocationContext, ToolContext]):
     if USER_DIRECT_CONSUME:
         session_state = get_session_state(ctx)
-        return session_state[FRONTEND_STATE_KEY]['biz'].get('projectId') or os.getenv(
-            'BOHRIUM_PROJECT_ID'
-        )
+        project_id = session_state[FRONTEND_STATE_KEY]['biz'].get(
+            'projectId'
+        ) or os.getenv('BOHRIUM_PROJECT_ID')
+        project_list = await get_project_list(access_key=_get_ak(ctx))
+        if project_id in project_list:
+            return project_id
+        else:
+            logger.warning(
+                f'{ctx.session.id} project_id <{project_id}> is not exist, use project_list[0] <{project_list[0]}>'
+            )
+            return project_list[0]
     else:
         return MATERIALS_PROJECT_ID