Refactor: Address code shortcomings and align with HF best practices

Phase 1 - Critical Fixes:
- Fix deprecated clear_gpu_memory() calls (remove model/tokenizer params)
- Register rate limiting middleware
- Add /v1/stats endpoint
- Improve thread safety with is_model_ready()
- Apply log_level from config dynamically

Phase 2 - Remove Redundancies:
- Simplify memory management (remove redundant cleanup in inference paths)
- Remove manual HF token env var setting (HF Hub handles it)
- Remove manual chat template loading (auto-loaded in transformers 4.45.0+)
- Remove manual device management (device_map='auto' handles it)

Phase 3 - Code Quality:
- Centralize version management in app/__init__.py
- Refactor long functions with helper methods
- Simplify memory cleanup to single pass

Phase 4 - Testing & Documentation:
- Rewrite unit tests to test actual provider logic
- Add test coverage for helper methods
- Update README with improvements and HF best practices alignment

README.md

@@ -136,6 +136,23 @@ response = client.chat.completions.create(
 - Development: L4x1 GPU (24GB VRAM)
 - Production: L40s GPU (48GB VRAM)
 
+## Recent Improvements
+
+### Code Quality & Hugging Face Best Practices Alignment
+
+This codebase has been optimized to align with Hugging Face inference best practices:
+
+- **Simplified Memory Management**: Removed redundant manual GPU memory cleanup - `device_map="auto"` handles this automatically
+- **Streamlined Token Management**: Hugging Face Hub now auto-detects tokens from environment variables
+- **Auto-Loading Chat Templates**: Leverages transformers 4.45.0+ automatic chat template loading
+- **Automatic Device Placement**: Removed manual device management - `device_map="auto"` handles GPU/CPU placement
+- **Improved Thread Safety**: Enhanced model access checks with thread-safe helpers
+- **Centralized Version Management**: Single source of truth for API version
+
+### Deprecated Functions
+
+- `clear_gpu_memory(model, tokenizer)` - Parameters deprecated, use `clear_gpu_memory()` without arguments
+
 ## Development
 
 ### Local Setup

app/__init__.py

@@ -1,2 +1,3 @@
-# empty package marker
+"""LLM Pro Finance API package."""
 
+__version__ = "1.0.0"

app/main.py

@@ -7,23 +7,29 @@ from typing import Dict
 from fastapi import FastAPI, status
 from fastapi.responses import JSONResponse
 
+from app import __version__
 from app.config import settings
 from app.middleware import api_key_guard
+from app.middleware.rate_limit import rate_limit_middleware
 from app.routers import openai_api
 
-# Configure logging
-logging.basicConfig(level=logging.INFO)
+# Configure logging with level from settings
+log_level = getattr(logging, settings.log_level.upper())
+logging.basicConfig(level=log_level)
 logger = logging.getLogger(__name__)
 
 app = FastAPI(
     title="LLM Pro Finance API (Transformers)",
     description="OpenAI-compatible API for financial LLM inference",
-    version="1.0.0"
+    version=__version__
 )
 
 # Mount routers
 app.include_router(openai_api.router, prefix="/v1")
 
+# Rate limiting middleware (applied first)
+app.middleware("http")(rate_limit_middleware)
+
 # Optional API key middleware
 app.middleware("http")(api_key_guard)
 
@@ -64,7 +70,7 @@ async def root() -> Dict[str, str]:
     return {
         "status": "ok", 
         "service": "Qwen Open Finance R 8B Inference", 
-        "version": "1.0.0",
+        "version": __version__,
         "model": settings.model,
         "backend": "Transformers"
     }

app/providers/transformers_provider.py

@@ -7,7 +7,7 @@ import re
 from typing import Dict, Any, AsyncIterator, Union, List, Optional
 import asyncio
 from threading import Thread, Lock
-from huggingface_hub import login, hf_hub_download
+from huggingface_hub import login
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, StoppingCriteria, StoppingCriteriaList
 
 from app.utils.constants import (
@@ -40,7 +40,6 @@ logger = logging.getLogger(__name__)
 # Global model state
 model = None
 tokenizer = None
-device = "cuda" if torch.cuda.is_available() else "cpu"
 _init_lock = Lock()
 _initializing = False
 _initialized = False
@@ -84,7 +83,7 @@ def initialize_model(force_reload: bool = False):
         # Clear previous model if force reloading
         if force_reload and model is not None:
             log_info("Force reload requested, clearing existing model...", print_output=True)
-            clear_gpu_memory(model, tokenizer)
+            clear_gpu_memory()
             model = None
             tokenizer = None
             _initialized = False
@@ -105,18 +104,12 @@ def initialize_model(force_reload: bool = False):
                 log_info(f"{token_source} found (length: {len(hf_token)})", print_output=True)
                 
                 # Authenticate with Hugging Face Hub
+                # login() automatically handles token precedence and environment variables
                 try:
                     login(token=hf_token, add_to_git_credential=False)
                     log_info("Successfully authenticated with Hugging Face Hub", print_output=True)
                 except Exception as e:
                     log_warning(f"Failed to authenticate with HF Hub: {e}", print_output=True)
-                
-                # Set token environment variables
-                os.environ.update({
-                    "HF_TOKEN": hf_token,
-                    "HUGGING_FACE_HUB_TOKEN": hf_token,
-                    "HF_API_TOKEN": hf_token,
-                })
             else:
                 log_warning(
                     "No HF token found! Model download may fail if model is gated.",
@@ -124,6 +117,7 @@ def initialize_model(force_reload: bool = False):
                 )
             
             # Load tokenizer
+            # Modern transformers (4.45.0+) auto-load chat templates from model repo
             log_info("Loading tokenizer...", print_output=True)
             tokenizer = AutoTokenizer.from_pretrained(
                 MODEL_NAME,
@@ -132,21 +126,9 @@ def initialize_model(force_reload: bool = False):
                 cache_dir=CACHE_DIR,
             )
             
-            # Load custom chat template if missing
+            # Verify chat template is available (should be auto-loaded)
             if not hasattr(tokenizer, 'chat_template') or tokenizer.chat_template is None:
-                try:
-                    template_path = hf_hub_download(
-                        repo_id=MODEL_NAME,
-                        filename="chat_template.jinja",
-                        repo_type="model",
-                        token=hf_token,
-                        cache_dir=CACHE_DIR,
-                    )
-                    with open(template_path, 'r', encoding='utf-8') as f:
-                        tokenizer.chat_template = f.read()
-                    log_info("Custom chat template applied", print_output=True)
-                except Exception as e:
-                    log_warning(f"Could not load custom template, using default: {e}")
+                log_warning("Chat template not found - will use fallback formatting")
             
             log_info("Tokenizer loaded", print_output=True)
             
@@ -178,7 +160,7 @@ def initialize_model(force_reload: bool = False):
             error_msg = f"Error initializing model: {e}"
             log_error(error_msg, exc_info=True, print_output=True)
             
-            clear_gpu_memory(model, tokenizer)
+            clear_gpu_memory()
             model = None
             tokenizer = None
             
@@ -222,8 +204,8 @@ class TransformersProvider:
     ) -> Union[Dict[str, Any], AsyncIterator[str]]:
         """Handle chat completion requests."""
         try:
-            # Initialize model on first use
-            if model is None:
+            # Initialize model on first use (thread-safe check)
+            if not is_model_ready():
                 log_info("Model not initialized, initializing now...")
                 initialize_model()
                 log_info("Model initialized successfully")
@@ -307,7 +289,8 @@ class TransformersProvider:
                 log_warning("No chat_template found, using fallback")
             
             # Tokenize
-            inputs = tokenizer(prompt, return_tensors="pt").to(device)
+            # device_map="auto" handles device placement automatically
+            inputs = tokenizer(prompt, return_tensors="pt")
             
             # Handle streaming vs non-streaming
             if stream:
@@ -323,110 +306,99 @@ class TransformersProvider:
         self, inputs, temperature: float, top_p: float, max_tokens: int, model_id: str, tools: Optional[List[Dict[str, Any]]] = None, json_output_required: bool = False
     ) -> Dict[str, Any]:
         """Generate non-streaming response."""
-        try:
-            # Prepare generation kwargs
-            generation_kwargs = {
-                "max_new_tokens": max_tokens,
-                "temperature": temperature,
-                "top_p": top_p,
-                "top_k": DEFAULT_TOP_K,
-                "do_sample": temperature > 0,
-                "pad_token_id": PAD_TOKEN_ID,
-                "eos_token_id": EOS_TOKENS,
-                "repetition_penalty": REPETITION_PENALTY,
-                "early_stopping": False,
-                "use_cache": True,
-            }
-            
-            # Note: Qwen reasoning models are designed to use reasoning tags
-            # We cannot completely disable reasoning, but we can:
-            # 1. Use strong prompts (already done above)
-            # 2. Post-process to extract desired output (done in _extract_json_from_text and _parse_tool_calls)
-            # 3. Set temperature to 0 for completely deterministic JSON output
-            #    Temperature=0 uses greedy decoding (always picks most likely token)
-            #    This maximizes consistency for structured outputs
-            if json_output_required:
-                # Set temperature to 0 for completely deterministic JSON output
-                # This uses greedy decoding which is ideal for structured formats
-                original_temp = generation_kwargs["temperature"]
-                generation_kwargs["temperature"] = 0.0
-                generation_kwargs["do_sample"] = False  # Explicitly set for temperature=0
-                log_info(f"Set temperature from {original_temp} to 0.0 (greedy decoding) for JSON output format")
-            
-            with torch.no_grad():
-                outputs = model.generate(
-                    **inputs,
-                    **generation_kwargs,
-                )
-            
-            # Extract token counts using tokenizer for accuracy
-            # Count prompt tokens (more accurate than shape[1] as it handles special tokens correctly)
-            prompt_tokens = len(inputs.input_ids[0])
-            generated_ids = outputs[0][inputs.input_ids.shape[1]:]
-            generated_text = tokenizer.decode(generated_ids, skip_special_tokens=True)
-            completion_tokens = len(generated_ids)
-            
-            # ✅ If JSON output is required, try to extract JSON from the response
-            if json_output_required:
-                generated_text = self._extract_json_from_text(generated_text)
-            
-            # ✅ Parse tool calls from generated text
-            tool_calls = None
-            if tools:
-                tool_calls = self._parse_tool_calls(generated_text, tools)
-                if tool_calls:
-                    log_info(f"Parsed {len(tool_calls)} tool calls from response")
-                    # Remove tool call markers from content if present
-                    generated_text = self._clean_tool_calls_from_text(generated_text)
-            
-            finish_reason = "tool_calls" if tool_calls else ("length" if completion_tokens >= max_tokens else "stop")
-            
-            log_info(f"Generated {completion_tokens} tokens (max: {max_tokens}), finish: {finish_reason}")
-            
-            # Record statistics
-            stats_tracker = get_stats_tracker()
-            stats_tracker.record_request(RequestStats(
-                timestamp=time.time(),
-                prompt_tokens=prompt_tokens,
-                completion_tokens=completion_tokens,
-                total_tokens=prompt_tokens + completion_tokens,
-                model=model_id,
-                finish_reason=finish_reason,
-            ))
-            
-            # Build message with optional tool_calls
-            message = {"role": "assistant", "content": generated_text if generated_text.strip() else None}
+        # Prepare generation kwargs
+        generation_kwargs = {
+            "max_new_tokens": max_tokens,
+            "temperature": temperature,
+            "top_p": top_p,
+            "top_k": DEFAULT_TOP_K,
+            "do_sample": temperature > 0,
+            "pad_token_id": PAD_TOKEN_ID,
+            "eos_token_id": EOS_TOKENS,
+            "repetition_penalty": REPETITION_PENALTY,
+            "early_stopping": False,
+            "use_cache": True,
+        }
+        
+        # Note: Qwen reasoning models are designed to use reasoning tags
+        # We cannot completely disable reasoning, but we can:
+        # 1. Use strong prompts (already done above)
+        # 2. Post-process to extract desired output (done in _extract_json_from_text and _parse_tool_calls)
+        # 3. Set temperature to 0 for completely deterministic JSON output
+        #    Temperature=0 uses greedy decoding (always picks most likely token)
+        #    This maximizes consistency for structured outputs
+        if json_output_required:
+            # Set temperature to 0 for completely deterministic JSON output
+            # This uses greedy decoding which is ideal for structured formats
+            original_temp = generation_kwargs["temperature"]
+            generation_kwargs["temperature"] = 0.0
+            generation_kwargs["do_sample"] = False  # Explicitly set for temperature=0
+            log_info(f"Set temperature from {original_temp} to 0.0 (greedy decoding) for JSON output format")
+        
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                **generation_kwargs,
+            )
+        
+        # Extract token counts using tokenizer for accuracy
+        # Count prompt tokens (more accurate than shape[1] as it handles special tokens correctly)
+        prompt_tokens = len(inputs.input_ids[0])
+        generated_ids = outputs[0][inputs.input_ids.shape[1]:]
+        generated_text = tokenizer.decode(generated_ids, skip_special_tokens=True)
+        completion_tokens = len(generated_ids)
+        
+        # ✅ If JSON output is required, try to extract JSON from the response
+        if json_output_required:
+            generated_text = self._extract_json_from_text(generated_text)
+        
+        # ✅ Parse tool calls from generated text
+        tool_calls = None
+        if tools:
+            tool_calls = self._parse_tool_calls(generated_text, tools)
             if tool_calls:
-                message["tool_calls"] = tool_calls
-            
-            return {
-                "id": f"chatcmpl-{os.urandom(12).hex()}",
-                "object": "chat.completion",
-                "created": int(time.time()),
-                "model": model_id,
-                "choices": [
-                    {
-                        "index": 0,
-                        "message": message,
-                        "finish_reason": finish_reason,
-                    }
-                ],
-                "usage": {
-                    "prompt_tokens": prompt_tokens,
-                    "completion_tokens": completion_tokens,
-                    "total_tokens": prompt_tokens + completion_tokens,
-                },
-            }
-        finally:
-            # Clean up GPU memory
-            if 'inputs' in locals():
-                del inputs
-            if 'outputs' in locals():
-                del outputs
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
-            import gc
-            gc.collect()
+                log_info(f"Parsed {len(tool_calls)} tool calls from response")
+                # Remove tool call markers from content if present
+                generated_text = self._clean_tool_calls_from_text(generated_text)
+        
+        finish_reason = "tool_calls" if tool_calls else ("length" if completion_tokens >= max_tokens else "stop")
+        
+        log_info(f"Generated {completion_tokens} tokens (max: {max_tokens}), finish: {finish_reason}")
+        
+        # Record statistics
+        stats_tracker = get_stats_tracker()
+        stats_tracker.record_request(RequestStats(
+            timestamp=time.time(),
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+            total_tokens=prompt_tokens + completion_tokens,
+            model=model_id,
+            finish_reason=finish_reason,
+        ))
+        
+        # Build message with optional tool_calls
+        message = {"role": "assistant", "content": generated_text if generated_text.strip() else None}
+        if tool_calls:
+            message["tool_calls"] = tool_calls
+        
+        return {
+            "id": f"chatcmpl-{os.urandom(12).hex()}",
+            "object": "chat.completion",
+            "created": int(time.time()),
+            "model": model_id,
+            "choices": [
+                {
+                    "index": 0,
+                    "message": message,
+                    "finish_reason": finish_reason,
+                }
+            ],
+            "usage": {
+                "prompt_tokens": prompt_tokens,
+                "completion_tokens": completion_tokens,
+                "total_tokens": prompt_tokens + completion_tokens,
+            },
+        }
     
     async def _chat_stream(
         self, inputs, temperature: float, top_p: float, max_tokens: int, model_id: str, tools: Optional[List[Dict[str, Any]]] = None, json_output_required: bool = False
@@ -455,12 +427,8 @@ class TransformersProvider:
         }
         
         def generate():
-            try:
-                with torch.no_grad():
-                    model.generate(**inputs, **generation_kwargs)
-            finally:
-                if torch.cuda.is_available():
-                    torch.cuda.empty_cache()
+            with torch.no_grad():
+                model.generate(**inputs, **generation_kwargs)
         
         generation_thread = Thread(target=generate)
         generation_thread.start()
@@ -504,11 +472,6 @@ class TransformersProvider:
                 model=model_id,
                 finish_reason=finish_reason,
             ))
-            
-            if 'inputs' in locals():
-                del inputs
-            import gc
-            gc.collect()
         
         # Send final chunk
         final_chunk = {
@@ -536,6 +499,51 @@ class TransformersProvider:
         prompt += "Assistant: "
         return prompt
     
+    def _remove_reasoning_tags(self, text: str) -> str:
+        """Remove Qwen reasoning tags from text."""
+        # Remove reasoning tags - matches <think>...</think>
+        cleaned_text = re.sub(
+            r'<think>.*?</think>',
+            '',
+            text,
+            flags=re.DOTALL | re.IGNORECASE
+        )
+        
+        # Handle unclosed reasoning tags (split on closing tag)
+        if "</think>" in cleaned_text:
+            parts = cleaned_text.split("</think>", 1)
+            if len(parts) > 1:
+                cleaned_text = parts[1].strip()
+        
+        # If still has opening tag but no closing, remove everything before first {
+        if "<think>" in cleaned_text.lower() and "{" in cleaned_text:
+            brace_pos = cleaned_text.find('{')
+            if brace_pos != -1:
+                cleaned_text = cleaned_text[brace_pos:]
+        
+        return cleaned_text
+    
+    def _extract_json_by_brace_matching(self, text: str, start_pos: int = 0) -> Optional[str]:
+        """Extract JSON object by matching braces starting at given position."""
+        brace_start = text.find('{', start_pos)
+        if brace_start == -1:
+            return None
+        
+        brace_count = 0
+        for i in range(brace_start, len(text)):
+            if text[i] == '{':
+                brace_count += 1
+            elif text[i] == '}':
+                brace_count -= 1
+                if brace_count == 0:
+                    json_candidate = text[brace_start:i+1]
+                    try:
+                        json.loads(json_candidate)
+                        return json_candidate
+                    except json.JSONDecodeError:
+                        return None
+        return None
+    
     def _format_tools_for_prompt(self, tools: List[Dict[str, Any]]) -> str:
         """Format tools for inclusion in system prompt."""
         tools_text = (
@@ -580,18 +588,8 @@ class TransformersProvider:
         """Parse tool calls from generated text."""
         tool_calls = []
         
-        # First, remove reasoning tags to get clean text
-        cleaned_text = generated_text
-        cleaned_text = re.sub(
-            r'<think>.*?</think>',
-            '',
-            cleaned_text,
-            flags=re.DOTALL | re.IGNORECASE
-        )
-        if "</think>" in cleaned_text:
-            parts = cleaned_text.split("</think>", 1)
-            if len(parts) > 1:
-                cleaned_text = parts[1].strip()
+        # Remove reasoning tags to get clean text
+        cleaned_text = self._remove_reasoning_tags(generated_text)
         
         # Pattern to match <tool_call>...</tool_call> blocks
         pattern = r'<tool_call>\s*({.*?})\s*</tool_call>'
@@ -608,27 +606,22 @@ class TransformersProvider:
         if not matches:
             tool_names = [t.get("function", {}).get("name", "") for t in tools]
             # Look for JSON objects that might be tool calls
-            brace_start = cleaned_text.find('{')
-            while brace_start != -1:
-                # Try to extract JSON object starting at this position
-                brace_count = 0
-                for i in range(brace_start, len(cleaned_text)):
-                    if cleaned_text[i] == '{':
-                        brace_count += 1
-                    elif cleaned_text[i] == '}':
-                        brace_count -= 1
-                        if brace_count == 0:
-                            json_candidate = cleaned_text[brace_start:i+1]
-                            try:
-                                candidate_data = json.loads(json_candidate)
-                                if "name" in candidate_data and candidate_data["name"] in tool_names:
-                                    matches.append(json_candidate)
-                                    break
-                            except json.JSONDecodeError:
-                                pass
-                            break
+            brace_start = 0
+            while True:
+                json_candidate = self._extract_json_by_brace_matching(cleaned_text, brace_start)
+                if json_candidate is None:
+                    break
+                try:
+                    candidate_data = json.loads(json_candidate)
+                    if "name" in candidate_data and candidate_data["name"] in tool_names:
+                        matches.append(json_candidate)
+                        break
+                except json.JSONDecodeError:
+                    pass
                 # Find next {
-                brace_start = cleaned_text.find('{', brace_start + 1)
+                brace_start = cleaned_text.find('{', cleaned_text.find(json_candidate) + len(json_candidate))
+                if brace_start == -1:
+                    break
         
         for i, match in enumerate(matches):
             try:
@@ -676,30 +669,7 @@ class TransformersProvider:
     def _extract_json_from_text(self, text: str) -> str:
         """Extract JSON from text, handling cases where JSON is wrapped in markdown, reasoning tags, or other text."""
         # Step 1: Remove reasoning tags first (Qwen reasoning models)
-        # Handle <think> tags (Qwen reasoning format - actual tag is <think>)
-        cleaned_text = text
-        
-        # Remove reasoning tags - matches <think>...</think>
-        cleaned_text = re.sub(
-            r'<think>.*?</think>',
-            '',
-            cleaned_text,
-            flags=re.DOTALL | re.IGNORECASE
-        )
-        
-        # Also handle unclosed reasoning tags (split on closing tag)
-        if "</think>" in cleaned_text:
-            parts = cleaned_text.split("</think>", 1)
-            if len(parts) > 1:
-                cleaned_text = parts[1].strip()
-        
-        # If still has opening tag but no closing, remove everything before first {
-        # This handles cases where reasoning tag is not closed but JSON follows
-        if "<think>" in cleaned_text.lower() and "{" in cleaned_text:
-            # Find first { and take everything from there
-            brace_pos = cleaned_text.find('{')
-            if brace_pos != -1:
-                cleaned_text = cleaned_text[brace_pos:]
+        cleaned_text = self._remove_reasoning_tags(text)
         
         # Step 2: Try to find JSON wrapped in markdown code blocks
         json_code_block = re.search(r'```(?:json)?\s*(\{.*?\})\s*```', cleaned_text, re.DOTALL)
@@ -733,24 +703,10 @@ class TransformersProvider:
         if best_match:
             return best_match.strip()
         
-        # Step 4: Fallback - try to find any JSON-like structure
-        # Look for { ... } and try to extract it, even if nested
-        brace_start = cleaned_text.find('{')
-        if brace_start != -1:
-            # Find matching closing brace
-            brace_count = 0
-            for i in range(brace_start, len(cleaned_text)):
-                if cleaned_text[i] == '{':
-                    brace_count += 1
-                elif cleaned_text[i] == '}':
-                    brace_count -= 1
-                    if brace_count == 0:
-                        json_candidate = cleaned_text[brace_start:i+1]
-                        try:
-                            json.loads(json_candidate)
-                            return json_candidate.strip()
-                        except json.JSONDecodeError:
-                            break
+        # Step 4: Fallback - try to find any JSON-like structure using brace matching
+        json_candidate = self._extract_json_by_brace_matching(cleaned_text)
+        if json_candidate:
+            return json_candidate.strip()
         
         # Step 5: If no JSON found, return cleaned text (without reasoning tags)
         # This allows the caller to handle it or show an error

app/routers/openai_api.py

@@ -19,6 +19,17 @@ async def list_models_endpoint():
     return await list_models()
 
 
+@router.get("/stats")
+async def get_stats():
+    """Get API usage statistics.
+    
+    Returns:
+        Dictionary containing request counts, token usage, and performance metrics.
+    """
+    from app.utils.stats import get_stats_tracker
+    return get_stats_tracker().get_stats()
+
+
 @router.post("/models/reload")
 async def reload_model(force: bool = Query(False, description="Force reload from Hugging Face Hub")):
     """

app/utils/memory.py

@@ -41,14 +41,9 @@ def clear_gpu_memory(model: Optional[Any] = None, tokenizer: Optional[Any] = Non
     if not torch.cuda.is_available():
         return
     
-    # Clear CUDA cache
+    # Clear CUDA cache and run garbage collection
+    # Single pass is sufficient with modern PyTorch and device_map="auto"
     torch.cuda.empty_cache()
     torch.cuda.synchronize()
     gc.collect()
-    
-    # Force multiple garbage collection passes
-    for _ in range(3):
-        gc.collect()
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
 

tests/test_providers.py

@@ -1,51 +1,164 @@
+"""Tests for Transformers provider."""
+
 import pytest
-from unittest.mock import patch, AsyncMock
-import httpx
+from unittest.mock import patch, MagicMock, AsyncMock
+import torch
 
-from app.providers.transformers_provider import list_models, chat
+from app.providers.transformers_provider import list_models, chat, is_model_ready, TransformersProvider
 
 
 @pytest.mark.asyncio
 async def test_list_models_success():
     """Test successful model listing."""
-    mock_response = {"data": [{"id": "test-model"}]}
+    result = await list_models()
     
-    with patch('httpx.AsyncClient') as mock_client:
-        mock_response_obj = AsyncMock()
-        mock_response_obj.json.return_value = mock_response
-        mock_response_obj.raise_for_status.return_value = None
-        
-        mock_client.return_value.__aenter__.return_value.get.return_value = mock_response_obj
-        
-        result = await list_models()
-        assert result == mock_response
+    assert "object" in result
+    assert result["object"] == "list"
+    assert "data" in result
+    assert len(result["data"]) > 0
+    assert result["data"][0]["object"] == "model"
 
 
 @pytest.mark.asyncio
-async def test_chat_success():
-    """Test successful chat completion."""
-    payload = {"model": "test", "messages": [{"role": "user", "content": "hello"}]}
-    mock_response = {"choices": [{"message": {"content": "hi"}}]}
-    
-    with patch('httpx.AsyncClient') as mock_client:
-        mock_response_obj = AsyncMock()
-        mock_response_obj.json.return_value = mock_response
-        mock_response_obj.raise_for_status.return_value = None
-        
-        mock_client.return_value.__aenter__.return_value.post.return_value = mock_response_obj
+async def test_list_models_structure():
+    """Test model listing returns correct structure."""
+    result = await list_models()
+    
+    model = result["data"][0]
+    assert "id" in model
+    assert "object" in model
+    assert "owned_by" in model
+    assert model["object"] == "model"
+
+
+@pytest.mark.asyncio
+async def test_chat_with_mock_model():
+    """Test chat completion with mocked model."""
+    payload = {
+        "model": "test-model",
+        "messages": [{"role": "user", "content": "hello"}],
+        "temperature": 0.7,
+        "max_tokens": 100
+    }
+    
+    # Mock the global model and tokenizer
+    mock_tokenizer = MagicMock()
+    mock_tokenizer.apply_chat_template.return_value = "formatted prompt"
+    mock_tokenizer.encode.return_value = [1, 2, 3]
+    mock_tokenizer.decode.return_value = "test response"
+    mock_tokenizer.__call__.return_value = {
+        "input_ids": torch.tensor([[1, 2, 3]]),
+        "attention_mask": torch.tensor([[1, 1, 1]])
+    }
+    
+    mock_model = MagicMock()
+    mock_outputs = MagicMock()
+    mock_outputs[0] = torch.tensor([[1, 2, 3, 4, 5]])
+    mock_model.generate.return_value = mock_outputs
+    mock_model.get_input_embeddings.return_value.num_embeddings = 1000
+    
+    with patch('app.providers.transformers_provider.model', mock_model), \
+         patch('app.providers.transformers_provider.tokenizer', mock_tokenizer), \
+         patch('app.providers.transformers_provider.is_model_ready', return_value=True), \
+         patch('app.providers.transformers_provider._initialized', True):
         
         result = await chat(payload, stream=False)
-        assert result == mock_response
+        
+        assert "id" in result
+        assert "object" in result
+        assert result["object"] == "chat.completion"
+        assert "choices" in result
+        assert len(result["choices"]) > 0
+        assert "usage" in result
 
 
 @pytest.mark.asyncio
-async def test_chat_stream():
+async def test_chat_streaming():
     """Test chat completion with streaming."""
-    payload = {"model": "test", "messages": [{"role": "user", "content": "hello"}]}
-    mock_stream = AsyncMock()
+    payload = {
+        "model": "test-model",
+        "messages": [{"role": "user", "content": "hello"}],
+        "stream": True
+    }
+    
+    # Mock for streaming
+    mock_tokenizer = MagicMock()
+    mock_tokenizer.apply_chat_template.return_value = "formatted prompt"
+    mock_tokenizer.__call__.return_value = {
+        "input_ids": torch.tensor([[1, 2, 3]]),
+        "attention_mask": torch.tensor([[1, 1, 1]])
+    }
     
-    with patch('httpx.AsyncClient') as mock_client:
-        mock_client.return_value.__aenter__.return_value.stream.return_value = mock_stream
+    with patch('app.providers.transformers_provider.model', MagicMock()), \
+         patch('app.providers.transformers_provider.tokenizer', mock_tokenizer), \
+         patch('app.providers.transformers_provider.is_model_ready', return_value=True), \
+         patch('app.providers.transformers_provider._initialized', True):
         
         result = await chat(payload, stream=True)
-        assert result == mock_stream
+        
+        # Should return an async iterator
+        assert hasattr(result, '__aiter__')
+
+
+def test_is_model_ready_false_when_not_initialized():
+    """Test is_model_ready returns False when model not initialized."""
+    with patch('app.providers.transformers_provider._initialized', False), \
+         patch('app.providers.transformers_provider.model', None), \
+         patch('app.providers.transformers_provider.tokenizer', None):
+        
+        assert is_model_ready() is False
+
+
+def test_is_model_ready_true_when_initialized():
+    """Test is_model_ready returns True when model is initialized."""
+    mock_model = MagicMock()
+    mock_tokenizer = MagicMock()
+    
+    with patch('app.providers.transformers_provider._initialized', True), \
+         patch('app.providers.transformers_provider.model', mock_model), \
+         patch('app.providers.transformers_provider.tokenizer', mock_tokenizer):
+        
+        assert is_model_ready() is True
+
+
+def test_provider_format_tools_for_prompt():
+    """Test tool formatting for prompt."""
+    provider = TransformersProvider()
+    tools = [
+        {
+            "function": {
+                "name": "test_tool",
+                "description": "A test tool",
+                "parameters": {"type": "object", "properties": {}}
+            }
+        }
+    ]
+    
+    result = provider._format_tools_for_prompt(tools)
+    
+    assert "test_tool" in result
+    assert "CRITICAL" in result
+    assert "<tool_call>" in result
+
+
+def test_provider_remove_reasoning_tags():
+    """Test reasoning tag removal."""
+    provider = TransformersProvider()
+    
+    text_with_tags = "<think>Some reasoning</think>Actual answer"
+    result = provider._remove_reasoning_tags(text_with_tags)
+    
+    assert "<think>" not in result
+    assert "Actual answer" in result
+
+
+def test_provider_extract_json_by_brace_matching():
+    """Test JSON extraction by brace matching."""
+    provider = TransformersProvider()
+    
+    text = "Some text {\"key\": \"value\"} more text"
+    result = provider._extract_json_by_brace_matching(text)
+    
+    assert result is not None
+    assert "key" in result
+    assert "value" in result