diff --git a/resources_servers/workplace_assistant/app.py b/resources_servers/workplace_assistant/app.py index 953405e31cf..2eea41839e9 100644 --- a/resources_servers/workplace_assistant/app.py +++ b/resources_servers/workplace_assistant/app.py @@ -103,27 +103,31 @@ async def route_to_python_function(self, path: str, body: WorkbenchRequest, requ output=f"Error executing tool '{path}': {str(e)}" ) # return error to model so that it can correct itself - async def verify(self, body: WorkbenchVerifyRequest) -> WorkbenchVerifyResponse: - ground_truth = body.ground_truth - response = body.response.output + async def verify(self, request: Request, body: WorkbenchVerifyRequest) -> WorkbenchVerifyResponse: + session_id = request.session[SESSION_ID_KEY] + try: + ground_truth = body.ground_truth + response = body.response.output - total_score = 0.0 + total_score = 0.0 - # Convert list of ResponseFunctionToolCall objects into list of dictionaries - predicted_function_calls = [] + # Convert list of ResponseFunctionToolCall objects into list of dictionaries + predicted_function_calls = [] - for message in response: - if message.type == "function_call": - predicted_function_calls.append(message.model_dump()) + for message in response: + if message.type == "function_call": + predicted_function_calls.append(message.model_dump()) - predicted_chat_content = [] + predicted_chat_content = [] - for message in response: - if message.type == "output_text": - predicted_chat_content.append(message.model_dump()) + for message in response: + if message.type == "output_text": + predicted_chat_content.append(message.model_dump()) - total_score += is_correct(predicted_function_calls, ground_truth, None) * 1.0 - return WorkbenchVerifyResponse(**body.model_dump(), reward=total_score) + total_score += is_correct(predicted_function_calls, ground_truth, None) * 1.0 + return WorkbenchVerifyResponse(**body.model_dump(), reward=total_score) + finally: + self.session_id_to_tool_env.pop(session_id, None) if __name__ == "__main__": diff --git a/resources_servers/workplace_assistant/tests/test_app.py b/resources_servers/workplace_assistant/tests/test_app.py index 54be4996df5..3e58b585ca8 100644 --- a/resources_servers/workplace_assistant/tests/test_app.py +++ b/resources_servers/workplace_assistant/tests/test_app.py @@ -16,7 +16,7 @@ import pandas as pd from fastapi import Request -from pytest import fixture +from pytest import fixture, raises from nemo_gym.base_resources_server import BaseSeedSessionRequest from nemo_gym.openai_utils import ( @@ -1254,12 +1254,9 @@ async def test_verify(self, config: WorkbenchResourcesServerConfig) -> None: ] mock_email_df = pd.DataFrame(mock_email_data) - # Seed session - resources_server = self.init_server(config) session_id = "test_session_company_dir" mock_scope = {"type": "http", "session": {SESSION_ID_KEY: session_id}} mock_request_with_session = Request(scope=mock_scope) - await resources_server.seed_session(mock_request_with_session, BaseSeedSessionRequest()) with ( patch( @@ -1274,6 +1271,8 @@ async def test_verify(self, config: WorkbenchResourcesServerConfig) -> None: mock_email_csv.return_value = mock_email_df resources_server = self.init_server(config) + await resources_server.seed_session(mock_request_with_session, BaseSeedSessionRequest()) + assert session_id in resources_server.session_id_to_tool_env HARDCODED_CURRENT_TIME = pd.to_datetime("2023-11-30T23:59:00") SYS_PROMPT = ( @@ -1381,9 +1380,29 @@ async def test_verify(self, config: WorkbenchResourcesServerConfig) -> None: id="0", ) - verification_response = await resources_server.verify(verify_request) + verification_response = await resources_server.verify(mock_request_with_session, verify_request) assert verification_response.reward == 1.0 + assert session_id not in resources_server.session_id_to_tool_env + + async def test_verify_cleans_up_session_when_scoring_fails(self, config: WorkbenchResourcesServerConfig) -> None: + resources_server = self.init_server(config) + session_id = "test_session_scoring_failure" + mock_scope = {"type": "http", "session": {SESSION_ID_KEY: session_id}} + mock_request_with_session = Request(scope=mock_scope) + await resources_server.seed_session(mock_request_with_session, BaseSeedSessionRequest()) + + verify_request = MagicMock() + verify_request.ground_truth = [] + verify_request.response.output = [] + + with ( + patch("resources_servers.workplace_assistant.app.is_correct", side_effect=RuntimeError("scoring failed")), + raises(RuntimeError, match="scoring failed"), + ): + await resources_server.verify(mock_request_with_session, verify_request) + + assert session_id not in resources_server.session_id_to_tool_env async def test_stateful_email_deletion_and_fetch(self, config: WorkbenchResourcesServerConfig) -> None: """ @@ -1578,13 +1597,20 @@ async def test_stateful_email_deletion_and_fetch(self, config: WorkbenchResource id="1", ) - verification_response = await resources_server.verify(verify_request) + session_id = "test_session_stateful_email" + mock_scope = {"type": "http", "session": {SESSION_ID_KEY: session_id}} + mock_request_with_session = Request(scope=mock_scope) + await resources_server.seed_session(mock_request_with_session, BaseSeedSessionRequest()) + assert session_id in resources_server.session_id_to_tool_env + + verification_response = await resources_server.verify(mock_request_with_session, verify_request) # The reward should be 1.0 because the predicted function calls in our # crafted `response` object exactly match the `ground_truth`. assert verification_response.reward == 1.0, ( f"Verification failed with reward {verification_response.reward}" ) + assert session_id not in resources_server.session_id_to_tool_env async def test_extra_arguments_error_handling(self, config: WorkbenchResourcesServerConfig) -> None: """