From 61635b2f3bf46efdd70e4c38b91efffcf249a752 Mon Sep 17 00:00:00 2001 From: Rasic2 <1051987201@qq.com> Date: Tue, 23 Sep 2025 14:50:56 +0800 Subject: [PATCH 1/3] refactor: streamline access key and project ID retrieval across callbacks and remove unused functions --- .../matmaster_agent/base_agents/callback.py | 111 ++---------------- .../matmaster_agent/base_agents/job_agent.py | 29 ++--- agents/matmaster_agent/callback.py | 6 +- agents/matmaster_agent/constant.py | 64 +++++----- agents/matmaster_agent/prompt.py | 2 +- 5 files changed, 61 insertions(+), 151 deletions(-) diff --git a/agents/matmaster_agent/base_agents/callback.py b/agents/matmaster_agent/base_agents/callback.py index e6009e7e..de8779c1 100644 --- a/agents/matmaster_agent/base_agents/callback.py +++ b/agents/matmaster_agent/base_agents/callback.py @@ -20,6 +20,8 @@ CURRENT_ENV, FRONTEND_STATE_KEY, LOCAL_EXECUTOR, + MATMASTER_ACCESS_KEY, + MATMASTER_PROJECT_ID, OPENAPI_HOST, Transfer2Agent, ) @@ -178,14 +180,6 @@ def _get_projectId(ctx: Union[InvocationContext, ToolContext]): ) -@check_None_wrapper -def _get_machineType(ctx: Union[InvocationContext, ToolContext]): - session_state = get_session_state(ctx) - return session_state[FRONTEND_STATE_KEY]['biz'].get('machineType') or os.getenv( - 'MACHINE_TYPE', 'c32_m64_cpu' - ) - - def _inject_ak(ctx: Union[InvocationContext, ToolContext], executor, storage): access_key = _get_ak(ctx) if executor is not None: @@ -219,8 +213,7 @@ def _inject_projectId(ctx: Union[InvocationContext, ToolContext], executor, stor def _inject_username(ctx: Union[InvocationContext, ToolContext], executor): - access_key = _get_ak(ctx) - username = ak_to_username(access_key=access_key) + username = ak_to_username(access_key=MATMASTER_ACCESS_KEY) if username: if executor is not None: if executor['type'] == 'dispatcher': # BohriumExecutor @@ -238,8 +231,7 @@ def _inject_username(ctx: Union[InvocationContext, ToolContext], executor): def _inject_ticket(ctx: Union[InvocationContext, ToolContext], executor): - access_key = _get_ak(ctx) - ticket = ak_to_ticket(access_key=access_key) + ticket = ak_to_ticket(access_key=MATMASTER_ACCESS_KEY) if ticket: if executor is not None: if executor['type'] == 'dispatcher': # BohriumExecutor @@ -270,65 +262,6 @@ def _inject_current_env(executor): return executor -def _inject_machine_type(ctx: Union[InvocationContext, ToolContext], executor): - machine_type = _get_machineType(ctx) - session_state = get_session_state(ctx) - logger.info( - f"biz = {session_state[FRONTEND_STATE_KEY]['biz']}; " - f"machineType = {machine_type}" - ) - if executor is not None: - if executor['type'] == 'dispatcher': # BohriumExecutor - current_machine_type = executor['machine']['remote_profile']['machine_type'] - if not current_machine_type: - executor['machine']['remote_profile']['machine_type'] = str( - machine_type - ) - logger.info(f"After inject_machine_type, executor = {executor}") - - return executor - - -def inject_ak_projectId(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 为 Transfer2Agent，不做 ak 和 project_id 设置/校验 - if tool.name == Transfer2Agent: - return None - - # 如果 tool 不是 CalculationMCPTool，不应该调用这个 callback - if not isinstance(tool, CalculationMCPTool): - raise TypeError( - 'Not CalculationMCPTool type, current tool does not have ' - ) - - # 获取 access_key - access_key, tool.executor, tool.storage = _inject_ak( - tool_context, tool.executor, tool.storage - ) - - # 获取 project_id - try: - project_id, tool.executor, tool.storage = _inject_projectId( - tool_context, tool.executor, tool.storage - ) - except ValueError as e: - raise ValueError('ProjectId is invalid') from e - - tool_context.state['ak'] = access_key - tool_context.state['project_id'] = project_id - - return wrapper - - def inject_username_ticket(func: BeforeToolCallback) -> BeforeToolCallback: @wraps(func) async def wrapper( @@ -383,14 +316,14 @@ async def wrapper( 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': int(tool_context.state['project_id']), + 'projectId': MATMASTER_PROJECT_ID, 'name': 'check_job_create', } - params = {'accessKey': tool_context.state['ak']} + params = {'accessKey': MATMASTER_ACCESS_KEY} logger.info( - f"[check_job_create] project_id = {tool_context.state['project_id']}, " - f"ak = {tool_context.state['ak']}" + f"[check_job_create] project_id = {MATMASTER_PROJECT_ID}, " + f"ak = {MATMASTER_ACCESS_KEY}" ) async with aiohttp.ClientSession() as session: @@ -402,7 +335,7 @@ async def wrapper( project_name = [ item['project_name'] for item in res['data']['items'] - if item['project_id'] == int(tool_context.state['project_id']) + if item['project_id'] == MATMASTER_PROJECT_ID ][0] async with aiohttp.ClientSession() as session: @@ -420,32 +353,6 @@ async def wrapper( return wrapper -def inject_machineType(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 为 Transfer2Agent，不做 ak 和 project_id 设置/校验 - if tool.name == Transfer2Agent: - return None - - # 如果 tool 不是 CalculationMCPTool，不应该调用这个 callback - if not isinstance(tool, CalculationMCPTool): - raise TypeError( - 'Not CalculationMCPTool type, current tool does not have ' - ) - - _inject_machine_type(tool_context, tool.executor) - - return wrapper - - # 总应该在最后 def catch_before_tool_callback_error(func: BeforeToolCallback) -> BeforeToolCallback: @wraps(func) diff --git a/agents/matmaster_agent/base_agents/job_agent.py b/agents/matmaster_agent/base_agents/job_agent.py index 8a494c74..00cb62ad 100644 --- a/agents/matmaster_agent/base_agents/job_agent.py +++ b/agents/matmaster_agent/base_agents/job_agent.py @@ -20,9 +20,7 @@ default_after_model_callback, default_after_tool_callback, default_before_tool_callback, - inject_ak_projectId, inject_current_env, - inject_machineType, inject_username_ticket, remove_function_call, tgz_oss_to_oss_list, @@ -37,6 +35,8 @@ LOADING_START, LOADING_STATE_KEY, LOADING_TITLE, + MATMASTER_ACCESS_KEY, + MATMASTER_PROJECT_ID, TMP_FRONTEND_STATE_KEY, ModelRole, OpenAPIJobAPI, @@ -160,14 +160,8 @@ def __init__( # Todo: support List[before_tool_callback] before_tool_callback = catch_before_tool_callback_error( - inject_machineType( - check_job_create( - inject_current_env( - inject_username_ticket( - inject_ak_projectId(before_tool_callback) - ) - ) - ) + check_job_create( + inject_current_env(inject_username_ticket(before_tool_callback)) ) ) after_tool_callback = check_before_tool_callback_effect( @@ -554,11 +548,15 @@ async def _run_async_impl( try: await self.tools[0].get_tools() if not ctx.session.state['dflow']: - access_key, Executor, BohriumStorge = _inject_ak( - ctx, get_BohriumExecutor(), get_BohriumStorage() + access_key, Executor, BohriumStorge = ( + MATMASTER_ACCESS_KEY, + get_BohriumExecutor(), + get_BohriumStorage(), ) - project_id, Executor, BohriumStorge = _inject_projectId( - ctx, Executor, BohriumStorge + project_id, Executor, BohriumStorge = ( + MATMASTER_PROJECT_ID, + Executor, + BohriumStorge, ) else: access_key, Executor, BohriumStorge = _inject_ak( @@ -850,6 +848,9 @@ async def _run_async_impl( params_check_completed_json: dict = json.loads( response.choices[0].message.content ) + logger.info( + f"[BaseAsyncJobAgent] params_check_completed_json = {params_check_completed_json}" + ) params_check_completed = params_check_completed_json['flag'] params_check_reason = params_check_completed_json['reason'] params_check_msg = params_check_completed_json['analyzed_messages'] diff --git a/agents/matmaster_agent/callback.py b/agents/matmaster_agent/callback.py index 02d7279d..73337a99 100644 --- a/agents/matmaster_agent/callback.py +++ b/agents/matmaster_agent/callback.py @@ -11,8 +11,7 @@ from google.genai import types from google.genai.types import FunctionCall, Part -from agents.matmaster_agent.base_agents.callback import _get_ak -from agents.matmaster_agent.constant import FRONTEND_STATE_KEY +from agents.matmaster_agent.constant import FRONTEND_STATE_KEY, MATMASTER_ACCESS_KEY from agents.matmaster_agent.locales import i18n from agents.matmaster_agent.model import UserContent from agents.matmaster_agent.prompt import get_user_content_lang @@ -112,7 +111,6 @@ async def matmaster_check_job_status( jobs_dict ): # 确认当前有在运行中的任务 running_job_ids = get_running_jobs_detail(jobs_dict) # 从 state 里面拿 - access_key = _get_ak(callback_context) # 从 state 或环境变量里面拿 if callback_context.state['target_language'] in [ 'Chinese', 'zh-CN', @@ -130,7 +128,7 @@ async def matmaster_check_job_status( '[matmaster_check_job_status] new LlmResponse, prepare call API' ) job_status = await get_job_status( - job_query_url, access_key=access_key + job_query_url, access_key=MATMASTER_ACCESS_KEY ) # 查询任务的最新状态 callback_context.state['new_query_job_status'][ 'origin_job_id' diff --git a/agents/matmaster_agent/constant.py b/agents/matmaster_agent/constant.py index 8c376e2f..0d133569 100644 --- a/agents/matmaster_agent/constant.py +++ b/agents/matmaster_agent/constant.py @@ -16,13 +16,38 @@ # DB DBUrl = os.getenv('SESSION_API_URL') +OPENAPI_HOST = '' +DFLOW_HOST = '' +DFLOW_K8S_API_SERVER = '' +BOHRIUM_API_URL = '' + +CURRENT_ENV = os.getenv('OPIK_PROJECT_NAME', 'prod') +if CURRENT_ENV == 'test': + OPENAPI_HOST = 'https://openapi.test.dp.tech' + DFLOW_HOST = 'https://lbg-workflow-mlops.test.dp.tech' + DFLOW_K8S_API_SERVER = 'https://lbg-workflow-mlops.test.dp.tech' + BOHRIUM_API_URL = 'https://bohrium-api.test.dp.tech' +elif CURRENT_ENV == 'uat': + OPENAPI_HOST = 'https://openapi.uat.dp.tech' + BOHRIUM_API_URL = 'https://bohrium-api.uat.dp.tech' +elif CURRENT_ENV == 'prod': + OPENAPI_HOST = 'https://openapi.dp.tech' + DFLOW_HOST = 'https://workflows.deepmodeling.com' + DFLOW_K8S_API_SERVER = 'https://workflows.deepmodeling.com' + BOHRIUM_API_URL = 'https://bohrium-api.dp.tech' + +OpenAPIJobAPI = f"{OPENAPI_HOST}/openapi/v1/job" + +MATMASTER_ACCESS_KEY = str(os.getenv('MATMASTER_ACCESS_KEY')) +MATMASTER_PROJECT_ID = int(os.getenv('MATMASTER_PROJECT_ID')) + # Bohrium Constant BohriumStorge = { 'type': 'https', 'plugin': { 'type': 'bohrium', - 'access_key': '', - 'project_id': -1, + 'access_key': MATMASTER_ACCESS_KEY, + 'project_id': MATMASTER_PROJECT_ID, 'app_key': 'agent', }, } @@ -34,38 +59,17 @@ 'batch_type': 'OpenAPI', 'context_type': 'OpenAPI', 'remote_profile': { - 'access_key': '', - 'project_id': -1, + 'access_key': MATMASTER_ACCESS_KEY, + 'project_id': MATMASTER_PROJECT_ID, 'app_key': 'agent', - 'image_address': 'registry.dp.tech/dptech/dp/native/prod-19853/dpa-mcp:0.0.0', + 'image_address': '', 'platform': 'ali', - 'machine_type': '', + 'machine_type': 'c2_m8_cpu', }, }, - 'resources': {'envs': {}}, + 'resources': {'envs': {'BOHRIUM_PROJECT_ID': MATMASTER_PROJECT_ID}}, } -OPENAPI_HOST = '' -DFLOW_HOST = '' -DFLOW_K8S_API_SERVER = '' -BOHRIUM_API_URL = '' - -CURRENT_ENV = os.getenv('OPIK_PROJECT_NAME', 'prod') -if CURRENT_ENV == 'test': - OPENAPI_HOST = 'https://openapi.test.dp.tech' - DFLOW_HOST = 'https://lbg-workflow-mlops.test.dp.tech' - DFLOW_K8S_API_SERVER = 'https://lbg-workflow-mlops.test.dp.tech' - BOHRIUM_API_URL = 'https://bohrium-api.test.dp.tech' -elif CURRENT_ENV == 'uat': - OPENAPI_HOST = 'https://openapi.uat.dp.tech' - BOHRIUM_API_URL = 'https://bohrium-api.uat.dp.tech' -elif CURRENT_ENV == 'prod': - OPENAPI_HOST = 'https://openapi.dp.tech' - DFLOW_HOST = 'https://workflows.deepmodeling.com' - DFLOW_K8S_API_SERVER = 'https://workflows.deepmodeling.com' - BOHRIUM_API_URL = 'https://bohrium-api.dp.tech' -OpenAPIJobAPI = f"{OPENAPI_HOST}/openapi/v1/job" - DFlowExecutor = { 'type': 'local', 'dflow': True, @@ -74,8 +78,8 @@ 'DFLOW_K8S_API_SERVER': DFLOW_K8S_API_SERVER, 'DFLOW_S3_REPO_KEY': 'oss-bohrium', 'DFLOW_S3_STORAGE_CLIENT': 'dflow.plugins.bohrium.TiefblueClient', - 'BOHRIUM_ACCESS_KEY': '', - 'BOHRIUM_PROJECT_ID': '', + 'BOHRIUM_ACCESS_KEY': MATMASTER_ACCESS_KEY, + 'BOHRIUM_PROJECT_ID': str(MATMASTER_PROJECT_ID), 'BOHRIUM_APP_KEY': 'agent', }, } diff --git a/agents/matmaster_agent/prompt.py b/agents/matmaster_agent/prompt.py index 88a8b757..9f8fa2ed 100644 --- a/agents/matmaster_agent/prompt.py +++ b/agents/matmaster_agent/prompt.py @@ -684,7 +684,7 @@ def gen_params_check_completed_agent_instruction(): {{ "flag": , "reason": , // *A concise explanation of the reasoning behind the judgment, covering both positive and negative evidence found in the context messages. Return empty string only if there is absolutely no relevant content to analyze.* - "analyzed_message": List[] // *Quote the key messages that were analyzed to make this determination.* + "analyzed_messages": List[] // *Quote the key messages that were analyzed to make this determination.* }} Return `flag: true` ONLY IF ALL of the following conditions are met: From 1a1db09b17586df7621f1624ba54b9cce6b7bc1b Mon Sep 17 00:00:00 2001 From: Rasic2 <1051987201@qq.com> Date: Tue, 23 Sep 2025 14:58:20 +0800 Subject: [PATCH 2/3] refactor: rename constants MATMASTER_ACCESS_KEY and MATMASTER_PROJECT_ID to MATERIALS_ACCESS_KEY and MATERIALS_PROJECT_ID for clarity and consistency --- agents/matmaster_agent/base_agents/callback.py | 18 +++++++++--------- .../matmaster_agent/base_agents/job_agent.py | 8 ++++---- agents/matmaster_agent/callback.py | 4 ++-- agents/matmaster_agent/constant.py | 18 +++++++++--------- 4 files changed, 24 insertions(+), 24 deletions(-) diff --git a/agents/matmaster_agent/base_agents/callback.py b/agents/matmaster_agent/base_agents/callback.py index de8779c1..eb47047a 100644 --- a/agents/matmaster_agent/base_agents/callback.py +++ b/agents/matmaster_agent/base_agents/callback.py @@ -20,8 +20,8 @@ CURRENT_ENV, FRONTEND_STATE_KEY, LOCAL_EXECUTOR, - MATMASTER_ACCESS_KEY, - MATMASTER_PROJECT_ID, + MATERIALS_ACCESS_KEY, + MATERIALS_PROJECT_ID, OPENAPI_HOST, Transfer2Agent, ) @@ -213,7 +213,7 @@ def _inject_projectId(ctx: Union[InvocationContext, ToolContext], executor, stor def _inject_username(ctx: Union[InvocationContext, ToolContext], executor): - username = ak_to_username(access_key=MATMASTER_ACCESS_KEY) + username = ak_to_username(access_key=MATERIALS_ACCESS_KEY) if username: if executor is not None: if executor['type'] == 'dispatcher': # BohriumExecutor @@ -231,7 +231,7 @@ def _inject_username(ctx: Union[InvocationContext, ToolContext], executor): def _inject_ticket(ctx: Union[InvocationContext, ToolContext], executor): - ticket = ak_to_ticket(access_key=MATMASTER_ACCESS_KEY) + ticket = ak_to_ticket(access_key=MATERIALS_ACCESS_KEY) if ticket: if executor is not None: if executor['type'] == 'dispatcher': # BohriumExecutor @@ -316,14 +316,14 @@ async def wrapper( 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': MATMASTER_PROJECT_ID, + 'projectId': MATERIALS_PROJECT_ID, 'name': 'check_job_create', } - params = {'accessKey': MATMASTER_ACCESS_KEY} + params = {'accessKey': MATERIALS_ACCESS_KEY} logger.info( - f"[check_job_create] project_id = {MATMASTER_PROJECT_ID}, " - f"ak = {MATMASTER_ACCESS_KEY}" + f"[check_job_create] project_id = {MATERIALS_PROJECT_ID}, " + f"ak = {MATERIALS_ACCESS_KEY}" ) async with aiohttp.ClientSession() as session: @@ -335,7 +335,7 @@ async def wrapper( project_name = [ item['project_name'] for item in res['data']['items'] - if item['project_id'] == MATMASTER_PROJECT_ID + if item['project_id'] == MATERIALS_PROJECT_ID ][0] async with aiohttp.ClientSession() as session: diff --git a/agents/matmaster_agent/base_agents/job_agent.py b/agents/matmaster_agent/base_agents/job_agent.py index 00cb62ad..c57a7bb1 100644 --- a/agents/matmaster_agent/base_agents/job_agent.py +++ b/agents/matmaster_agent/base_agents/job_agent.py @@ -35,8 +35,8 @@ LOADING_START, LOADING_STATE_KEY, LOADING_TITLE, - MATMASTER_ACCESS_KEY, - MATMASTER_PROJECT_ID, + MATERIALS_ACCESS_KEY, + MATERIALS_PROJECT_ID, TMP_FRONTEND_STATE_KEY, ModelRole, OpenAPIJobAPI, @@ -549,12 +549,12 @@ async def _run_async_impl( await self.tools[0].get_tools() if not ctx.session.state['dflow']: access_key, Executor, BohriumStorge = ( - MATMASTER_ACCESS_KEY, + MATERIALS_ACCESS_KEY, get_BohriumExecutor(), get_BohriumStorage(), ) project_id, Executor, BohriumStorge = ( - MATMASTER_PROJECT_ID, + MATERIALS_PROJECT_ID, Executor, BohriumStorge, ) diff --git a/agents/matmaster_agent/callback.py b/agents/matmaster_agent/callback.py index 73337a99..405e0df8 100644 --- a/agents/matmaster_agent/callback.py +++ b/agents/matmaster_agent/callback.py @@ -11,7 +11,7 @@ from google.genai import types from google.genai.types import FunctionCall, Part -from agents.matmaster_agent.constant import FRONTEND_STATE_KEY, MATMASTER_ACCESS_KEY +from agents.matmaster_agent.constant import FRONTEND_STATE_KEY, MATERIALS_ACCESS_KEY from agents.matmaster_agent.locales import i18n from agents.matmaster_agent.model import UserContent from agents.matmaster_agent.prompt import get_user_content_lang @@ -128,7 +128,7 @@ async def matmaster_check_job_status( '[matmaster_check_job_status] new LlmResponse, prepare call API' ) job_status = await get_job_status( - job_query_url, access_key=MATMASTER_ACCESS_KEY + job_query_url, access_key=MATERIALS_ACCESS_KEY ) # 查询任务的最新状态 callback_context.state['new_query_job_status'][ 'origin_job_id' diff --git a/agents/matmaster_agent/constant.py b/agents/matmaster_agent/constant.py index 0d133569..8dede87c 100644 --- a/agents/matmaster_agent/constant.py +++ b/agents/matmaster_agent/constant.py @@ -38,16 +38,16 @@ OpenAPIJobAPI = f"{OPENAPI_HOST}/openapi/v1/job" -MATMASTER_ACCESS_KEY = str(os.getenv('MATMASTER_ACCESS_KEY')) -MATMASTER_PROJECT_ID = int(os.getenv('MATMASTER_PROJECT_ID')) +MATERIALS_ACCESS_KEY = str(os.getenv('MATERIALS_ACCESS_KEY')) +MATERIALS_PROJECT_ID = int(os.getenv('MATERIALS_PROJECT_ID')) # Bohrium Constant BohriumStorge = { 'type': 'https', 'plugin': { 'type': 'bohrium', - 'access_key': MATMASTER_ACCESS_KEY, - 'project_id': MATMASTER_PROJECT_ID, + 'access_key': MATERIALS_ACCESS_KEY, + 'project_id': MATERIALS_PROJECT_ID, 'app_key': 'agent', }, } @@ -59,15 +59,15 @@ 'batch_type': 'OpenAPI', 'context_type': 'OpenAPI', 'remote_profile': { - 'access_key': MATMASTER_ACCESS_KEY, - 'project_id': MATMASTER_PROJECT_ID, + 'access_key': MATERIALS_ACCESS_KEY, + 'project_id': MATERIALS_PROJECT_ID, 'app_key': 'agent', 'image_address': '', 'platform': 'ali', 'machine_type': 'c2_m8_cpu', }, }, - 'resources': {'envs': {'BOHRIUM_PROJECT_ID': MATMASTER_PROJECT_ID}}, + 'resources': {'envs': {'BOHRIUM_PROJECT_ID': MATERIALS_PROJECT_ID}}, } DFlowExecutor = { @@ -78,8 +78,8 @@ 'DFLOW_K8S_API_SERVER': DFLOW_K8S_API_SERVER, 'DFLOW_S3_REPO_KEY': 'oss-bohrium', 'DFLOW_S3_STORAGE_CLIENT': 'dflow.plugins.bohrium.TiefblueClient', - 'BOHRIUM_ACCESS_KEY': MATMASTER_ACCESS_KEY, - 'BOHRIUM_PROJECT_ID': str(MATMASTER_PROJECT_ID), + 'BOHRIUM_ACCESS_KEY': MATERIALS_ACCESS_KEY, + 'BOHRIUM_PROJECT_ID': str(MATERIALS_PROJECT_ID), 'BOHRIUM_APP_KEY': 'agent', }, } From 9e5a71dedb69cdb6aae61fb9a58f0b2a4e678ee8 Mon Sep 17 00:00:00 2001 From: Rasic2 <1051987201@qq.com> Date: Tue, 23 Sep 2025 22:27:42 +0800 Subject: [PATCH 3/3] fix: update API endpoints, logging, and data types for enhanced compatibility and clarity --- agents/matmaster_agent/base_agents/job_agent.py | 3 +++ agents/matmaster_agent/constant.py | 2 +- agents/matmaster_agent/model.py | 2 +- agents/matmaster_agent/structure_generate_agent/agent.py | 2 +- agents/matmaster_agent/structure_generate_agent/constant.py | 2 +- 5 files changed, 7 insertions(+), 4 deletions(-) diff --git a/agents/matmaster_agent/base_agents/job_agent.py b/agents/matmaster_agent/base_agents/job_agent.py index c57a7bb1..c6ac2d9f 100644 --- a/agents/matmaster_agent/base_agents/job_agent.py +++ b/agents/matmaster_agent/base_agents/job_agent.py @@ -404,6 +404,9 @@ async def _run_async_impl( ): raw_result = part.function_response.response['result'] results = json.loads(raw_result.content[0].text) + logger.info( + f"[SubmitCoreCalculationMCPLlmAgent] results = {results}" + ) origin_job_id = results['job_id'] job_name = part.function_response.name job_status = results['status'] diff --git a/agents/matmaster_agent/constant.py b/agents/matmaster_agent/constant.py index 8dede87c..93d3a334 100644 --- a/agents/matmaster_agent/constant.py +++ b/agents/matmaster_agent/constant.py @@ -36,7 +36,7 @@ DFLOW_K8S_API_SERVER = 'https://workflows.deepmodeling.com' BOHRIUM_API_URL = 'https://bohrium-api.dp.tech' -OpenAPIJobAPI = f"{OPENAPI_HOST}/openapi/v1/job" +OpenAPIJobAPI = f"{OPENAPI_HOST}/openapi/v1/sandbox/job" MATERIALS_ACCESS_KEY = str(os.getenv('MATERIALS_ACCESS_KEY')) MATERIALS_PROJECT_ID = int(os.getenv('MATERIALS_PROJECT_ID')) diff --git a/agents/matmaster_agent/model.py b/agents/matmaster_agent/model.py index ffa5e0f4..94efb981 100644 --- a/agents/matmaster_agent/model.py +++ b/agents/matmaster_agent/model.py @@ -51,7 +51,7 @@ class JobResult(BaseModel): class BohrJobInfo(BaseModel): origin_job_id: str - job_id: int + job_id: Union[int, str] job_query_url: str job_detail_url: str job_status: JobStatus diff --git a/agents/matmaster_agent/structure_generate_agent/agent.py b/agents/matmaster_agent/structure_generate_agent/agent.py index 26d23be4..a42906ae 100644 --- a/agents/matmaster_agent/structure_generate_agent/agent.py +++ b/agents/matmaster_agent/structure_generate_agent/agent.py @@ -26,7 +26,7 @@ ] = 'registry.dp.tech/dptech/dp/native/prod-788025/structure-generate-agent:small' StructureGenerateBohriumExecutor['machine']['remote_profile'][ 'machine_type' -] = 'c8_m31_1 * NVIDIA T4' +] = 'c8_m32_1 * NVIDIA 4090' sse_params = SseServerParams(url=StructureGenerateServerUrl) diff --git a/agents/matmaster_agent/structure_generate_agent/constant.py b/agents/matmaster_agent/structure_generate_agent/constant.py index 7600f88a..d306aeb5 100644 --- a/agents/matmaster_agent/structure_generate_agent/constant.py +++ b/agents/matmaster_agent/structure_generate_agent/constant.py @@ -3,6 +3,6 @@ StructureGenerateAgentName = 'structure_generate_agent' if CURRENT_ENV in ['test', 'uat']: - StructureGenerateServerUrl = 'http://pfmx1355864.bohrium.tech:50003/sse' + StructureGenerateServerUrl = 'http://hftp1387064.bohrium.tech:50001/sse' else: StructureGenerateServerUrl = 'https://cystalformer-uuid1754551471.app-space.dplink.cc/sse?token=1750cd294e6c4270946ae37107a725ff'