Minor type hint fix

VRSEN · Dec 9, 2024 · fa05d19 · fa05d19
1 parent 77add5d
commit fa05d19
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/agency_swarm/agency/agency.py b/agency_swarm/agency/agency.py
@@ -11,6 +11,7 @@
     Dict,
     List,
     Literal,
+    Optional,
     Type,
     TypedDict,
     TypeVar,
@@ -218,7 +219,7 @@ def get_completion(
     def get_completion_stream(
         self,
         message: str,
-        event_handler: type(AgencyEventHandler),
+        event_handler: Optional[Type[AgencyEventHandler]] = None,
         message_files: List[str] = None,
         recipient_agent: Agent = None,
         additional_instructions: str = None,
@@ -231,7 +232,7 @@ def get_completion_stream(
 
         Parameters:
             message (str): The message for which completion is to be retrieved.
-            event_handler (type(AgencyEventHandler)): The event handler class to handle the completion stream. https://github.com/openai/openai-python/blob/main/helpers.md
+            event_handler (Type[AgencyEventHandler], optional): The event handler class to handle the completion stream. https://github.com/openai/openai-python/blob/main/helpers.md Defaults to None.
             message_files (list, optional): A list of file ids to be sent as attachments with the message. When using this parameter, files will be assigned both to file_search and code_interpreter tools if available. It is recommended to assign files to the most sutiable tool manually, using the attachments parameter.  Defaults to None.
             recipient_agent (Agent, optional): The agent to which the message should be sent. Defaults to the first agent in the agency chart.
             additional_instructions (str, optional): Additional instructions to be sent with the message. Defaults to None.
@@ -280,7 +281,7 @@ def get_completion_parse(
 
         Parameters:
             message (str): The message for which completion is to be retrieved.
-            response_format (type(BaseModel)): The response format to use for the completion.
+            response_format (Type[BaseModel]): The response format to use for the completion.
             message_files (list, optional): A list of file ids to be sent as attachments with the message. When using this parameter, files will be assigned both to file_search and code_interpreter tools if available. It is recommended to assign files to the most sutiable tool manually, using the attachments parameter.  Defaults to None.
             recipient_agent (Agent, optional): The agent to which the message should be sent. Defaults to the first agent in the agency chart.
             additional_instructions (str, optional): Additional instructions to be sent with the message. Defaults to None.

diff --git a/agency_swarm/threads/thread.py b/agency_swarm/threads/thread.py
@@ -5,7 +5,7 @@
 import re
 import time
 from concurrent.futures import ThreadPoolExecutor, as_completed
-from typing import List, Optional, Union
+from typing import List, Optional, Type, Union
 
 from openai import APIError, BadRequestError
 from openai.types.beta import AssistantToolChoice
@@ -80,7 +80,7 @@ def init_thread(self):
     def get_completion_stream(
         self,
         message: Union[str, List[dict], None],
-        event_handler: type(AgencyEventHandler),
+        event_handler: Optional[Type[AgencyEventHandler]] = None,
         message_files: List[str] = None,
         attachments: Optional[List[Attachment]] = None,
         recipient_agent: Agent = None,
@@ -107,7 +107,7 @@ def get_completion(
         attachments: Optional[List[dict]] = None,
         recipient_agent: Union[Agent, None] = None,
         additional_instructions: str = None,
-        event_handler: type(AgencyEventHandler) = None,
+        event_handler: Optional[Type[AgencyEventHandler]] = None,
         tool_choice: AssistantToolChoice = None,
         yield_messages: bool = False,
         response_format: Optional[dict] = None,