Merge pull request #22 from bettersg/agent/uat/openai-assistants-api-1

Agent/uat/openai assistants api 1
bettersg · Mar 24, 2024 · e6a83a4 · e6a83a4
2 parents cbf32ec + 63afe90
commit e6a83a4
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 16 deletions.
diff --git a/base/agent_base_class.py b/base/agent_base_class.py
@@ -22,27 +22,18 @@ def get_website_content(self, link):
 
     def message_handler(self, message: MessagePayload):
         messageId = message.messageId
-        try:
-            vote = self.check_message(message)
-        except UnsupportedMessageTypeException as e:
-            logging.info("Message type not supported")
-            return
-        except Exception as e:
-            raise ValueError(f"Error in check_message implementation: {e}")
-        # check if valid Vote 
-        if not isinstance(vote, Vote):
-            raise TypeError(f"Expected Vote, got {type(vote)} from check_message implementation")
-
         api_host = os.getenv("API_HOST")
         agent_name = os.getenv("AGENT_NAME")
         if not api_host:
             raise ValueError("API_HOST environment variable not set")
         if not agent_name:
             raise ValueError("AGENT_NAME environment variable not set")
-
-        auth_req = google.auth.transport.requests.Request()
-        id_token = google.oauth2.id_token.fetch_id_token(auth_req, api_host)
-        headers = {"Authorization": f"Bearer {id_token}"}
+        try:
+            auth_req = google.auth.transport.requests.Request()
+            id_token = google.oauth2.id_token.fetch_id_token(auth_req, api_host)
+            headers = {"Authorization": f"Bearer {id_token}"}
+        except Exception as e:
+            raise ValueError(f"Error fetching ID token: {e}")
         try:
             vote_initialisation = VoteInitialisation(factCheckerName=agent_name)
             res = requests.post(f"{api_host}/messages/{messageId}/voteRequests", json=vote_initialisation.model_dump(mode="json"), headers=headers)
@@ -58,6 +49,16 @@ def message_handler(self, message: MessagePayload):
 
         if messageId not in vote_request_path:
             raise ValueError(f"Vote request path does not contain messageId")
+        try:
+            vote = self.check_message(message)
+        except UnsupportedMessageTypeException as e:
+            logging.info("Message type not supported")
+            return
+        except Exception as e:
+            raise ValueError(f"Error in check_message implementation: {e}")
+        # check if valid Vote 
+        if not isinstance(vote, Vote):
+            raise TypeError(f"Expected Vote, got {type(vote)} from check_message implementation")
 
         try:
             res = requests.patch(f"{api_host}/{vote_request_path}", json=vote.model_dump(mode="json"), headers=headers)

diff --git a/implementation/agent.py b/implementation/agent.py
@@ -74,7 +74,7 @@ def process_call(tool_call):
 
         start = time()
 
-        while not agent_complete and ((time() - start) < 60):
+        while not agent_complete and ((time() - start) < 120):
             sleep(1)
             run = client.beta.threads.runs.retrieve(
                 thread_id=thread.id,