small changes'

bespokelabsai · Dec 15, 2024 · 6a2cfd0 · 6a2cfd0
1 parent a25cb3c
commit 6a2cfd0
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 5 deletions.
diff --git a/src/bespokelabs/curator/request_processor/base_online_request_processor.py b/src/bespokelabs/curator/request_processor/base_online_request_processor.py
@@ -474,9 +474,6 @@ async def handle_single_request_with_retries(
                 status_tracker=status_tracker,
             )
 
-            if generic_response.response_message is None:
-                raise ValueError(f"Request {request.task_id} returned no response message")
-
             # Save response in the base class
             await self.append_generic_response(generic_response, save_filepath)
 

diff --git a/src/bespokelabs/curator/request_processor/litellm_online_request_processor.py b/src/bespokelabs/curator/request_processor/litellm_online_request_processor.py
@@ -18,6 +18,7 @@
 logger = logging.getLogger(__name__)
 
 litellm.suppress_debug_info = True
+REQUEST_TIMEOUT = 60.0
 
 
 class LiteLLMOnlineRequestProcessor(BaseOnlineRequestProcessor):
@@ -268,15 +269,15 @@ async def call_single_request(
                     await self.client.chat.completions.create_with_completion(
                         **request.api_specific_request,
                         response_model=request.prompt_formatter.response_format,
-                        timeout=60.0,
+                        timeout=REQUEST_TIMEOUT,
                     )
                 )
                 response_message = (
                     response.model_dump() if hasattr(response, "model_dump") else response
                 )
             else:
                 completion_obj = await litellm.acompletion(
-                    **request.api_specific_request, timeout=60.0
+                    **request.api_specific_request, timeout=REQUEST_TIMEOUT
                 )
                 response_message = completion_obj["choices"][0]["message"]["content"]
         except litellm.RateLimitError as e:
@@ -300,6 +301,11 @@ async def call_single_request(
         except litellm.NotFoundError as e:
             cost = 0
 
+        if response_message is None:
+            raise ValueError(
+                f"Request {request.task_id} returned no response message with raw response {completion_obj.model_dump()}"
+            )
+
         # Create and return response
         return GenericResponse(
             response_message=response_message,