@@ -16,8 +16,8 @@ def __init__(self) -> None:
1616
1717 """
1818 self .max_requests_per_batch = BatchRequestContent .MAX_REQUESTS
19- self .batches : list [BatchRequestContent ] = []
2019 self .current_batch : BatchRequestContent = BatchRequestContent ()
20+ self .batches : list [BatchRequestContent ] = [self .current_batch ]
2121
2222 def add_batch_request_item (self , request : BatchRequestItem ) -> None :
2323 """
@@ -28,8 +28,8 @@ def add_batch_request_item(self, request: BatchRequestItem) -> None:
2828 if len (self .current_batch .requests ) >= self .max_requests_per_batch :
2929 self .batches .append (self .current_batch .finalize ())
3030 self .current_batch = BatchRequestContent ()
31+ self .batches .append (self .current_batch )
3132 self .current_batch .add_request (request .id , request )
32- self .batches .append (self .current_batch )
3333
3434 def remove_batch_request_item (self , request_id : str ) -> None :
3535 """
@@ -49,16 +49,30 @@ def new_batch_with_failed_requests(self) -> Optional[BatchRequestContent]:
4949 Optional[BatchRequestContent]: A new batch with failed requests.
5050 """
5151 # Use IDs to get response status codes, generate new batch with failed requests
52- batch_with_failed_responses : Optional [BatchRequestContent ] = BatchRequestContent ()
52+ batch_with_failed_responses : Optional [BatchRequestContent ] = (
53+ BatchRequestContent ()
54+ )
5355 for batch in self .batches :
5456 for request in batch .requests :
55- if request .status_code not in [ # type: ignore # Method should be deprecated
56- 200 , 201 , 202 , 203 , 204 , 205 , 206 , 207 , 208 , 226
57- ]:
57+ if (
58+ request .status_code
59+ not in [ # type: ignore # Method should be deprecated
60+ 200 ,
61+ 201 ,
62+ 202 ,
63+ 203 ,
64+ 204 ,
65+ 205 ,
66+ 206 ,
67+ 207 ,
68+ 208 ,
69+ 226 ,
70+ ]
71+ ):
5872 if batch_with_failed_responses is not None :
5973 batch_with_failed_responses .add_request (
6074 request .id , # type: ignore # Bug. Method should be deprecated
61- request # type: ignore
75+ request , # type: ignore
6276 )
6377 else :
6478 raise ValueError ("batch_with_failed_responses is None" )
0 commit comments