Added Max Time Params Handling

Author: UtkarshTheDev

CHANGELOG.md

@@ -2,6 +2,18 @@
 
 All notable changes to LocalLab will be documented in this file.
 
+## [0.7.1] - 2025-05-18
+
+### Fixed
+
+- Fixed critical error: "ModelManager.generate() got an unexpected keyword argument 'max_time'"
+- Added proper handling of the `max_time` parameter in all generation endpoints
+- Updated `ModelManager.generate()` method to accept the `max_time` parameter
+- Added `max_time` parameter to all request models (GenerationRequest, BatchGenerationRequest, ChatRequest)
+- Ensured consistent parameter passing between client and server
+- Set default max_time to 180 seconds (3 minutes) when not specified
+- Improved error handling for generation timeouts
+
 ## Client Package [1.1.0] - 2025-05-17
 
 ### Added

locallab/__init__.py

@@ -6,7 +6,7 @@
 # This ensures Hugging Face's progress bars are displayed correctly
 from .utils.early_config import configure_hf_logging
 
-__version__ = "0.7.0"  # Improved stream generation and non-streaming generation quality
+__version__ = "0.7.1"  # Fixed max_time parameter handling in generation endpoints
 
 # Only import what's necessary initially, lazy-load the rest
 from .logger import get_logger

locallab/model_manager.py

@@ -399,7 +399,8 @@ async def generate(
         top_k: Optional[int] = None,
         repetition_penalty: Optional[float] = None,
         system_instructions: Optional[str] = None,
-        do_sample: bool = True
+        do_sample: bool = True,
+        max_time: Optional[float] = None
     ) -> str:
         """Generate text from the model"""
         # Check model timeout
@@ -527,10 +528,13 @@ async def generate(
                     }
 
                     # Set a reasonable max time for generation to prevent hanging
-                    # Use the DEFAULT_MAX_TIME from config (increased to 180 seconds)
-                    if "max_time" not in generate_params and not stream:
-                        from .config import DEFAULT_MAX_TIME
-                        generate_params["max_time"] = DEFAULT_MAX_TIME  # Use the default max time from config
+                    # Use the provided max_time or a default value of 180 seconds
+                    if not stream:
+                        if max_time is not None:
+                            generate_params["max_time"] = max_time
+                        elif "max_time" not in generate_params:
+                            # Default to 180 seconds (3 minutes) if not specified
+                            generate_params["max_time"] = 180.0  # Default max time in seconds
 
                     # Define comprehensive stop sequences for proper termination
                     stop_sequences = [
@@ -918,7 +922,7 @@ async def async_stream_generate(self, inputs: Dict[str, torch.Tensor] = None, ge
 
             # Update with provided kwargs
             for key, value in kwargs.items():
-                if key in ["max_length", "temperature", "top_p", "top_k", "repetition_penalty"]:
+                if key in ["max_length", "temperature", "top_p", "top_k", "repetition_penalty", "max_time"]:
                     gen_params[key] = value
                 elif key == "max_new_tokens":
                     # Handle the max_new_tokens parameter by mapping to max_length

locallab/routes/generate.py

@@ -35,6 +35,7 @@ class GenerationRequest(BaseModel):
     top_p: float = Field(default=DEFAULT_TOP_P, ge=0.0, le=1.0)
     top_k: int = Field(default=80, ge=1, le=1000)  # Added top_k parameter
     repetition_penalty: float = Field(default=1.15, ge=1.0, le=2.0)  # Added repetition_penalty parameter
+    max_time: Optional[float] = Field(default=None, ge=0.0, description="Maximum time in seconds for generation")
     system_prompt: Optional[str] = Field(default=DEFAULT_SYSTEM_INSTRUCTIONS)
     stream: bool = Field(default=False)
 
@@ -47,6 +48,7 @@ class BatchGenerationRequest(BaseModel):
     top_p: float = Field(default=DEFAULT_TOP_P, ge=0.0, le=1.0)
     top_k: int = Field(default=80, ge=1, le=1000)  # Added top_k parameter
     repetition_penalty: float = Field(default=1.15, ge=1.0, le=2.0)  # Added repetition_penalty parameter
+    max_time: Optional[float] = Field(default=None, ge=0.0, description="Maximum time in seconds for generation")
     system_prompt: Optional[str] = Field(default=DEFAULT_SYSTEM_INSTRUCTIONS)
 
 
@@ -64,6 +66,7 @@ class ChatRequest(BaseModel):
     top_p: float = Field(default=DEFAULT_TOP_P, ge=0.0, le=1.0)
     top_k: int = Field(default=80, ge=1, le=1000)  # Added top_k parameter
     repetition_penalty: float = Field(default=1.15, ge=1.0, le=2.0)  # Added repetition_penalty parameter
+    max_time: Optional[float] = Field(default=None, ge=0.0, description="Maximum time in seconds for generation")
     stream: bool = Field(default=False)
 
 
@@ -129,7 +132,7 @@ async def generate_text(request: GenerationRequest) -> GenerationResponse:
         # Return a streaming response
         return StreamingResponse(
             generate_stream(request.prompt, request.max_tokens, request.temperature,
-                           request.top_p, request.system_prompt),
+                           request.top_p, request.system_prompt, request.max_time),
             media_type="text/event-stream"
         )
 
@@ -144,7 +147,8 @@ async def generate_text(request: GenerationRequest) -> GenerationResponse:
             "top_p": request.top_p if request.top_p is not None else 0.92,  # Optimized default
             "top_k": request.top_k if request.top_k is not None else 80,  # Optimized default
             "repetition_penalty": request.repetition_penalty if request.repetition_penalty is not None else 1.15,  # Optimized default
-            "do_sample": model_params.get("do_sample", True)  # Pass do_sample from model params
+            "do_sample": model_params.get("do_sample", True),  # Pass do_sample from model params
+            "max_time": request.max_time  # Pass max_time parameter
         }
 
         # Merge model-specific params with request params
@@ -212,7 +216,7 @@ async def chat_completion(request: ChatRequest) -> ChatResponse:
     # If streaming is requested, return a streaming response
     if request.stream:
         return StreamingResponse(
-            stream_chat(formatted_prompt, request.max_tokens, request.temperature, request.top_p),
+            stream_chat(formatted_prompt, request.max_tokens, request.temperature, request.top_p, request.max_time),
             media_type="text/event-stream"
         )
 
@@ -227,7 +231,8 @@ async def chat_completion(request: ChatRequest) -> ChatResponse:
             "top_p": request.top_p if request.top_p is not None else 0.92,  # Optimized default
             "top_k": request.top_k if request.top_k is not None else 80,  # Optimized default
             "repetition_penalty": request.repetition_penalty if request.repetition_penalty is not None else 1.15,  # Optimized default
-            "do_sample": model_params.get("do_sample", True)  # Pass do_sample from model params
+            "do_sample": model_params.get("do_sample", True),  # Pass do_sample from model params
+            "max_time": request.max_time  # Pass max_time parameter
         }
 
         # Merge model-specific params with request params
@@ -292,7 +297,8 @@ async def generate_stream(
     max_tokens: int,
     temperature: float,
     top_p: float,
-    system_prompt: Optional[str]
+    system_prompt: Optional[str],
+    max_time: Optional[float] = None
 ) -> AsyncGenerator[str, None]:
     """
     Generate text in a streaming fashion and return as server-sent events
@@ -309,7 +315,8 @@ async def generate_stream(
             "top_p": top_p,
             "top_k": 80,  # Optimized top_k for high-quality streaming
             "repetition_penalty": 1.15,  # Optimized repetition_penalty for high-quality streaming
-            "do_sample": model_params.get("do_sample", True)  # Pass do_sample from model params
+            "do_sample": model_params.get("do_sample", True),  # Pass do_sample from model params
+            "max_time": max_time  # Pass max_time parameter
         }
 
         # Merge model-specific params with request params
@@ -361,7 +368,8 @@ async def stream_chat(
     formatted_prompt: str,
     max_tokens: int,
     temperature: float,
-    top_p: float
+    top_p: float,
+    max_time: Optional[float] = None
 ) -> AsyncGenerator[str, None]:
     """
     Stream chat completion responses as server-sent events
@@ -378,7 +386,8 @@ async def stream_chat(
             "top_p": top_p,
             "top_k": 80,  # Optimized top_k for high-quality streaming
             "repetition_penalty": 1.15,  # Optimized repetition_penalty for high-quality streaming
-            "do_sample": model_params.get("do_sample", True)  # Pass do_sample from model params
+            "do_sample": model_params.get("do_sample", True),  # Pass do_sample from model params
+            "max_time": max_time  # Pass max_time parameter
         }
 
         # Merge model-specific params with request params
@@ -438,7 +447,8 @@ async def batch_generate(request: BatchGenerationRequest) -> BatchGenerationResp
             "top_p": request.top_p if request.top_p is not None else 0.92,  # Optimized default
             "top_k": request.top_k if request.top_k is not None else 80,  # Optimized default
             "repetition_penalty": request.repetition_penalty if request.repetition_penalty is not None else 1.15,  # Optimized default
-            "do_sample": model_params.get("do_sample", True)  # Pass do_sample from model params
+            "do_sample": model_params.get("do_sample", True),  # Pass do_sample from model params
+            "max_time": request.max_time  # Pass max_time parameter
         }
 
         # Merge model-specific params with request params

setup.py

@@ -47,7 +47,7 @@
 
 setup(
     name="locallab",
-    version="0.7.0",
+    version="0.7.1",
     packages=find_packages(include=["locallab", "locallab.*"]),
     install_requires=install_requires,
     extras_require={