11import requests
2- from typing import List , Dict
2+ from typing import List , Dict , Optional
3+
34
45class GAMEClient :
56 def __init__ (self , api_key : str ):
@@ -13,7 +14,7 @@ def _get_access_token(self) -> str:
1314 response = requests .post (
1415 "https://api.virtuals.io/api/accesses/tokens" ,
1516 json = {"data" : {}},
16- headers = {"x-api-key" : self .api_key }
17+ headers = {"x-api-key" : self .api_key },
1718 )
1819
1920 if response .status_code != 200 :
@@ -22,12 +23,21 @@ def _get_access_token(self) -> str:
2223 response_json = response .json ()
2324 return response_json ["data" ]["accessToken" ]
2425
25- def _post (self , endpoint : str , data : dict ) -> dict :
26+ def _post (
27+ self , endpoint : str , data : dict , extra_headers : Optional [Dict [str , str ]] = None
28+ ) -> dict :
2629 """
2730 Internal method to post data
2831 """
2932 access_token = self ._get_access_token ()
3033
34+ # Default headers with Authorization
35+ headers = {"Authorization" : f"Bearer { access_token } " }
36+
37+ # Merge additional headers if provided
38+ if extra_headers :
39+ headers .update (extra_headers )
40+
3141 response = requests .post (
3242 f"{ self .base_url } /prompts" ,
3343 json = {
@@ -40,7 +50,7 @@ def _post(self, endpoint: str, data: dict) -> dict:
4050 "data" : data ,
4151 },
4252 },
43- headers = { "Authorization" : f"Bearer { access_token } " } ,
53+ headers = headers ,
4454 )
4555
4656 if response .status_code != 200 :
@@ -89,20 +99,28 @@ def set_worker_task(self, agent_id: str, task: str) -> Dict:
8999 data = {"task" : task },
90100 )
91101
92- def get_worker_action (self , agent_id : str , submission_id : str , data : dict ) -> Dict :
102+ def get_worker_action (
103+ self ,
104+ agent_id : str ,
105+ submission_id : str ,
106+ data : dict ,
107+ model_name : str ,
108+ ) -> Dict :
93109 """
94110 Get worker actions (for standalone worker)
95111 """
96112 return self ._post (
97113 endpoint = f"/v2/agents/{ agent_id } /tasks/{ submission_id } /next" ,
98114 data = data ,
115+ extra_headers = {"model_name" : model_name },
99116 )
100117
101- def get_agent_action (self , agent_id : str , data : dict ) -> Dict :
118+ def get_agent_action (self , agent_id : str , data : dict , model_name : str ) -> Dict :
102119 """
103120 Get agent actions/next step (for agent)
104121 """
105122 return self ._post (
106123 endpoint = f"/v2/agents/{ agent_id } /actions" ,
107124 data = data ,
125+ extra_headers = {"model_name" : model_name },
108126 )
0 commit comments