Updated LocalLab v0.2.8

Author: UtkarshTheDev

CHANGELOG.md

@@ -2,13 +2,29 @@
 
 All notable changes to LocalLab will be documented in this file.
 
+## [0.2.8] - 2025-03-02
+
+### Fixed
+
+- Fixed parameter mismatch in text generation endpoints by properly handling `max_new_tokens` parameter
+- Resolved coroutine awaiting issues in streaming generation endpoints
+- Fixed async generator handling in `stream_chat` and `generate_stream` functions
+- Enhanced error handling in streaming responses to provide better error messages
+- Improved compatibility between route parameters and model manager methods
+
 ## [0.2.7] - 2025-03-02
 
+### Added
+
+- Added missing dependencies in `setup.py`: huggingface_hub, pynvml, and typing_extensions
+- Improved dependency management with dev extras for testing packages
+- Enhanced error handling for GPU memory detection
+
 ### Fixed
 
-- Added missing dependency `fastapi-cache2` that was causing server startup errors
-- Added missing dependency `nvidia-ml-py3` to properly monitor NVIDIA GPUs
-- Improved error handling for GPU monitoring when dependencies are missing
+- Fixed circular import issues between modules
+- Improved error handling in system utilities
+- Enhanced compatibility with Google Colab environments
 
 ## [0.2.6] - 2025-03-02
 

locallab/__init__.py

@@ -2,7 +2,7 @@
 LocalLab: Run LLMs locally with a friendly API similar to OpenAI
 """
 
-__version__ = "0.2.7" 
+__version__ = "0.2.8" 
 
 from typing import Dict, Any, Optional
 

locallab/model_manager.py

@@ -280,6 +280,7 @@ async def generate(
         prompt: str,
         stream: bool = False,
         max_length: Optional[int] = None,
+        max_new_tokens: Optional[int] = None,
         temperature: Optional[float] = None,
         top_p: Optional[float] = None,
         top_k: Optional[int] = None,
@@ -307,6 +308,10 @@ async def generate(
             from .config import get_model_generation_params
             gen_params = get_model_generation_params(self.current_model)
 
+            # Handle max_new_tokens parameter (map to max_length)
+            if max_new_tokens is not None:
+                max_length = max_new_tokens
+            
             # Override with user-provided parameters if specified
             if max_length is not None:
                 try:
@@ -423,8 +428,40 @@ def _stream_generate(
             logger.error(f"Streaming generation failed: {str(e)}")
             raise HTTPException(status_code=500, detail=f"Streaming generation failed: {str(e)}")
 
-    async def async_stream_generate(self, inputs: Dict[str, torch.Tensor], gen_params: Dict[str, Any]):
-        """Convert the synchronous stream generator to an async generator."""
+    async def async_stream_generate(self, inputs: Dict[str, torch.Tensor] = None, gen_params: Dict[str, Any] = None, prompt: str = None, system_prompt: Optional[str] = None, **kwargs):
+        """Convert the synchronous stream generator to an async generator.
+        
+        This can be called either with:
+        1. inputs and gen_params directly (internal use)
+        2. prompt, system_prompt and other kwargs (from generate_stream adapter)
+        """
+        # If called with prompt, prepare inputs and parameters
+        if prompt is not None:
+            # Get appropriate system instructions
+            from .config import system_instructions
+            instructions = str(system_instructions.get_instructions(self.current_model)) if not system_prompt else str(system_prompt)
+            
+            # Format prompt with system instructions
+            formatted_prompt = f"""<|system|>{instructions}</|system|>\n<|user|>{prompt}</|user|>\n<|assistant|>"""
+            
+            # Get model-specific generation parameters
+            from .config import get_model_generation_params
+            gen_params = get_model_generation_params(self.current_model)
+            
+            # Update with provided kwargs
+            for key, value in kwargs.items():
+                if key in ["max_length", "temperature", "top_p", "top_k", "repetition_penalty"]:
+                    gen_params[key] = value
+                elif key == "max_new_tokens":
+                    # Handle the max_new_tokens parameter by mapping to max_length
+                    gen_params["max_length"] = value
+            
+            # Tokenize the prompt
+            inputs = self.tokenizer(formatted_prompt, return_tensors="pt")
+            for key in inputs:
+                inputs[key] = inputs[key].to(self.device)
+        
+        # Now stream tokens using the prepared inputs and parameters
         for token in self._stream_generate(inputs, gen_params=gen_params):
             yield token
             await asyncio.sleep(0)
@@ -564,6 +601,11 @@ async def generate_text(self, prompt: str, system_prompt: Optional[str] = None,
         """
         # Make sure we're not streaming when generating text
         kwargs["stream"] = False
+        
+        # Handle max_new_tokens parameter by mapping to max_length if needed
+        if "max_new_tokens" in kwargs and "max_length" not in kwargs:
+            kwargs["max_length"] = kwargs.pop("max_new_tokens")
+            
         # Directly await the generate method to return the string result
         return await self.generate(prompt=prompt, system_instructions=system_prompt, **kwargs)
 
@@ -572,7 +614,14 @@ async def generate_stream(self, prompt: str, system_prompt: Optional[str] = None
         Calls the async_stream_generate method with proper parameters."""
         # Ensure streaming is enabled
         kwargs["stream"] = True
-        return self.async_stream_generate(prompt=prompt, system_prompt=system_prompt, **kwargs)
+        
+        # Handle max_new_tokens parameter by mapping to max_length
+        if "max_new_tokens" in kwargs and "max_length" not in kwargs:
+            kwargs["max_length"] = kwargs.pop("max_new_tokens")
+            
+        # Call async_stream_generate with the prompt and parameters
+        async for token in self.async_stream_generate(prompt=prompt, system_prompt=system_prompt, **kwargs):
+            yield token
 
     def is_model_loaded(self, model_id: str) -> bool:
         """Check if a specific model is loaded.

locallab/routes/generate.py

@@ -5,7 +5,7 @@
 from fastapi import APIRouter, HTTPException, BackgroundTasks, Request
 from fastapi.responses import JSONResponse, StreamingResponse
 from pydantic import BaseModel, Field
-from typing import Dict, List, Any, Optional, Generator, Tuple
+from typing import Dict, List, Any, Optional, Generator, Tuple, AsyncGenerator
 import json
 
 from ..logger import get_logger
@@ -212,9 +212,9 @@ async def generate_stream(
     temperature: float, 
     top_p: float, 
     system_prompt: Optional[str]
-) -> Generator[str, None, None]:
+) -> AsyncGenerator[str, None]:
     """
-    Generate text in a streaming fashion
+    Generate text in a streaming fashion and return as server-sent events
     """
     try:
         # Get model-specific generation parameters
@@ -230,12 +230,15 @@ async def generate_stream(
         # Merge model-specific params with request params
         generation_params.update(model_params)
 
-        # Stream tokens
-        async for token in model_manager.generate_stream(
+        # Get the stream generator
+        stream_generator = model_manager.generate_stream(
             prompt=prompt,
             system_prompt=system_prompt,
             **generation_params
-        ):
+        )
+        
+        # Stream tokens
+        async for token in stream_generator:
             # Format as server-sent event
             data = token.replace("\n", "\\n")
             yield f"data: {data}\n\n"
@@ -252,9 +255,9 @@ async def stream_chat(
     max_tokens: int,
     temperature: float,
     top_p: float
-) -> Generator[str, None, None]:
+) -> AsyncGenerator[str, None]:
     """
-    Stream chat completion
+    Stream chat completion responses as server-sent events
     """
     try:
         # Get model-specific generation parameters
@@ -270,19 +273,21 @@ async def stream_chat(
         # Merge model-specific params with request params
         generation_params.update(model_params)
 
-        # Generate streaming tokens
-        async for token in model_manager.generate_stream(
+        # Generate streaming tokens - properly await the async generator
+        stream_generator = model_manager.generate_stream(
             prompt=formatted_prompt,
             **generation_params
-        ):
+        )
+        
+        async for token in stream_generator:
             # Format as a server-sent event with the structure expected by chat clients
             data = json.dumps({"role": "assistant", "content": token})
             yield f"data: {data}\n\n"
 
         # End of stream marker
         yield "data: [DONE]\n\n"
     except Exception as e:
-        logger.error(f"Chat streaming failed: {str(e)}")
+        logger.error(f"Streaming generation failed: {str(e)}")
         error_data = json.dumps({"error": str(e)})
         yield f"data: {error_data}\n\n"
         yield "data: [DONE]\n\n"

setup.py

@@ -5,7 +5,7 @@
 
 setup(
     name="locallab",
-    version="0.2.7",
+    version="0.2.8",
     packages=find_packages(include=["locallab", "locallab.*"]),
     install_requires=[
         "fastapi>=0.95.0,<1.0.0",