Resources Server Implementation

1
# simplified
2
from typing import Any, Dict
3
4
from fastapi import FastAPI, HTTPException, Request
5
from pydantic import BaseModel, ConfigDict, Field
6
7
from nemo_gym.base_resources_server import (
8
    BaseResourcesServerConfig,
9
    BaseSeedSessionRequest,
10
    BaseSeedSessionResponse,
11
    BaseVerifyRequest,
12
    BaseVerifyResponse,
13
    SimpleResourcesServer,
14
)
15
from nemo_gym.server_utils import SESSION_ID_KEY
16
from resources_servers.workplace_assistant.utils import get_tools, is_correct
17
18
class WorkbenchResourcesServerConfig(BaseResourcesServerConfig):
19
    pass
20
21
class WorkbenchRequest(BaseModel):
22
    model_config = ConfigDict(extra="allow")
23
24
class WorkbenchResponse(BaseModel):
25
    model_config = ConfigDict(extra="allow")
26
27
class WorkbenchVerifyRequest(BaseVerifyRequest):
28
    ground_truth: list[Dict[str, str]] | str
29
    id: int
30
    category: str
31
    environment_name: str
32
33
class WorkbenchVerifyResponse(BaseVerifyResponse):
34
    pass
35
36
class WorkbenchResourcesServer(SimpleResourcesServer):
37
    config: WorkbenchResourcesServerConfig
38
    session_id_to_tool_env: Dict[str, Any] = Field(default_factory=dict)
39
40
    def setup_webserver(self) -> FastAPI:
41
        app = super().setup_webserver()
42
        # Dynamic routing: any path becomes a tool call
43
        app.post("/{path}")(self.route_to_python_function)
44
        return app
45
46
    async def seed_session(
47
        self,
48
        request: Request,
49
        body: BaseSeedSessionRequest
50
    ) -> BaseSeedSessionResponse:
51
        session_id = request.session[SESSION_ID_KEY]
52
53
        # Initialize multiple toolkits for this session
54
        toolkits = [
55
            "email",
56
            "calendar",
57
            "analytics",
58
            "project_management",
59
            "customer_relationship_manager",
60
        ]
61
        self.session_id_to_tool_env[session_id] = get_tools(toolkits)
62
        return BaseSeedSessionResponse()
63
64
    # Generic tool router - dispatches to Python functions dynamically
65
    async def route_to_python_function(
66
        self,
67
        path: str,
68
        body: WorkbenchRequest,
69
        request: Request
70
    ) -> WorkbenchResponse:
71
        session_id = request.session[SESSION_ID_KEY]
72
73
        if session_id not in self.session_id_to_tool_env:
74
            raise HTTPException(
75
                status_code=400,
76
                detail="Session not initialized. Please call seed_session first.",
77
            )
78
79
        tool_env = self.session_id_to_tool_env[session_id]
80
        args = {k: v for k, v in body.model_dump(exclude_unset=True).items() if v is not None}
81
82
        try:
83
            function = tool_env["functions"][path]
84
            result = function(**args)
85
            return WorkbenchResponse(output=result)
86
        except Exception as e:
87
            # Return error to model so it can self-correct
88
            return WorkbenchResponse(output=f"Error executing tool '{path}': {str(e)}")
89
90
    async def verify(self, body: WorkbenchVerifyRequest) -> WorkbenchVerifyResponse:
91
        ground_truth = body.ground_truth
92
        response = body.response.output
93
94
        # Extract function calls from response
95
        predicted_function_calls = [
96
            message.model_dump()
97
            for message in response
98
            if message.type == "function_call"
99
        ]
100
101
        # Compute reward using custom evaluation function
102
        total_score = is_correct(predicted_function_calls, ground_truth, None) * 1.0
103
        return WorkbenchVerifyResponse(**body.model_dump(), reward=total_score)
104
105
if __name__ == "__main__":
106
    WorkbenchResourcesServer.run_webserver()