Spaces:

AshokGakr
/

tiny_model

Sleeping

Update app.py
db1e4cb verified 9 months ago
45.1 kB
"""
FastAPI backend with optimized streaming for LFM2.5-1.2B-Instruct
"""

import os, json, re, traceback, asyncio, time, sys
from typing import Optional, List, Dict, Any, AsyncGenerator
from datetime import datetime, timezone
from contextlib import asynccontextmanager

import torch
import numpy as np
from fastapi.staticfiles import StaticFiles
from fastapi import (
    FastAPI,
    HTTPException,
    Form,
    UploadFile,
    File,
    Request,
    status,
)
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse
from fastapi.templating import Jinja2Templates
from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModelForCausalLM, TextIteratorStreamer
from threading import Thread

# Import the tools module
try:
    from tools import web_search
    print("✅ Tools module imported successfully")
except ImportError as e:
    print(f"⚠️ Could not import tools module: {e}")
    def web_search(query: str) -> str:
        return f"Search results for: {query} (tools module not available)"

# ================== CONFIG =================
ENVIRONMENT = os.getenv("ENVIRONMENT", "development")
DEBUG = ENVIRONMENT != "production"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"🔧 Using device: {DEVICE}")
print(f"🔧 PyTorch version: {torch.__version__}")
print(f"🔧 CUDA available: {torch.cuda.is_available()}")

# Updated models from provided URLs
LFM_MODEL_NAME = "AshokGakr/model-tiny"

MAX_FILE_SIZE = 50 * 1024 * 1024
MAX_TOKENS_DEFAULT = 4500
MAX_TOKENS_HARD_LIMIT = 24000  # Conservative limit for stability
DEFAULT_TEMPERATURE = 0.7  # LFM2.5 default is 0.1, but keep adjustable
STREAM_DELAY = 0.01
MAX_CONTEXT_LENGTH = 32768  # Safe context length for LFM2.5 (max 32768)

# Global model variables
lfm_tokenizer = None
lfm_model = None

# Setup directories
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
TEMPLATES_DIR = os.path.join(BASE_DIR, "templates")
STATIC_DIR = os.path.join(BASE_DIR, "static")

os.makedirs(TEMPLATES_DIR, exist_ok=True)
os.makedirs(STATIC_DIR, exist_ok=True)
templates = Jinja2Templates(directory=TEMPLATES_DIR)

# ================== MODEL LOADING =================

async def load_models():
    """Load and initialize all ML models with optimized settings for LFM2.5"""
    global lfm_tokenizer, lfm_model
    
    print("\n" + "="*50)
    print("🔄 LOADING AI MODELS - LFM2.5-1.2B-INSTRUCT")
    print("="*50)
    
    # ==== OPTIMIZED LFM2.5 LOADING ====
    lfm_success = False
    try:
        print(f"📥 Loading GAKR tiny model model")
        print(f"model loading from {LFM_MODEL_NAME}")
        # Load tokenizer with correct settings
        lfm_tokenizer = AutoTokenizer.from_pretrained(
            LFM_MODEL_NAME,
            padding_side="left",  # Essential for streaming
            trust_remote_code=False,
            use_fast=True
        )
        
        # IMPORTANT: LFM2.5 uses special tokens <|im_start|> and <|im_end|>
        # No need to set pad token separately as the tokenizer handles it
        
        print(f"   Tokenizer loaded with vocab size: {lfm_tokenizer.vocab_size}")
        print(f"   Chat template available: {lfm_tokenizer.chat_template is not None}")
        
        # Configure model loading based on available hardware
        load_kwargs = {
            "torch_dtype": torch.bfloat16 if DEVICE == "cuda" else torch.float32,
            "trust_remote_code": False,
            "device_map": "auto" if DEVICE == "cuda" else None,
        }
        
        lfm_model = AutoModelForCausalLM.from_pretrained(
            LFM_MODEL_NAME,
            **load_kwargs
        )
        
        # If device_map not used, manually move to device
        if DEVICE != "cuda" or load_kwargs.get("device_map") is None:
            lfm_model = lfm_model.to(DEVICE)
        
        lfm_model.eval()
        
        # Enable eval mode for inference optimizations
        lfm_model.config.use_cache = True  # Enable KV cache for faster generation
        
        print(f"✅ GAKR tiny model loaded successfully")
        print(f"   Model device: {lfm_model.device}")
        print(f"   Model dtype: {lfm_model.dtype}")
        print(f"   Model parameters: ~1.2B")
        print(f"   Context window: 32,768 tokens")
        lfm_success = True
        
    except Exception as e:
        print(f"❌ Failed to load LFM2.5 model: {e}")
        traceback.print_exc()
        
        # Fallback: Try with simpler settings
        try:
            print("🔄 Trying simpler loading method...")
            lfm_tokenizer = AutoTokenizer.from_pretrained(LFM_MODEL_NAME)
            
            lfm_model = AutoModelForCausalLM.from_pretrained(
                LFM_MODEL_NAME,
                torch_dtype=torch.float32
            )
            lfm_model.to(DEVICE).eval()
            print(f"✅ LFM2.5 model loaded (fallback method)")
            lfm_success = True
        except Exception as e2:
            print(f"❌ Fallback loading failed: {e2}")
            lfm_tokenizer = None
            lfm_model = None
    
    print("="*50)
    print(f"📊 MODEL LOADING SUMMARY:")
    print(f"  • GAKR tiny Model: {'✅' if lfm_success else '❌'}")
    print(f"  • Device: {DEVICE}")
    print(f"  • Max Context: {MAX_CONTEXT_LENGTH} tokens")
    print("="*50 + "\n")
    
    return lfm_success

async def cleanup_models():
    """Cleanup models on shutdown"""
    print("🧹 Cleaning up models...")
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.synchronize()
    print("✅ Cleanup complete")

# ================== LIFESPAN MANAGER =================

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Modern lifespan context manager for FastAPI"""
    print("\n" + "="*50)
    print("🚀 STARTING GAKR AI ASSISTANT")
    print("="*50)
    print(f"📁 Working directory: {BASE_DIR}")
    print(f"🐍 Python version: {sys.version}")
    print(f"🔥 Torch version: {torch.__version__}")
    print("="*50 + "\n")
    
    model_loaded = await load_models()
    if not model_loaded:
        print("⚠️ CRITICAL: tiny model not loaded. Assistant will not function.")
    
    yield
    
    await cleanup_models()
    print("\n" + "="*50)
    print("👋 SHUTTING DOWN GAKR AI ASSISTANT")
    print("="*50 + "\n")

# ================== APP SETUP =================

app = FastAPI(
    title="GAKR AI Assistant",
    description="FastAPI backend with streaming for LFM2.5-1.2B-Instruct",
    version="5.0.0",
    lifespan=lifespan
)

app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")

app.add_middleware(
    CORSMiddleware, 
    allow_origins=["*"], 
    allow_credentials=True, 
    allow_methods=["*"], 
    allow_headers=["*"]
)

# ================== OPTIMIZED HELPERS =================

def validate_prompt_length(prompt: str, max_length: int = MAX_CONTEXT_LENGTH) -> bool:
    """Validate prompt doesn't exceed context window"""
    if not lfm_tokenizer:
        return True
    
    tokens = lfm_tokenizer.encode(prompt, truncation=False)
    return len(tokens) <= max_length

def get_truncated_prompt(prompt: str, max_tokens: int = 3072) -> str:
    """Truncate prompt to fit within context window"""
    if not lfm_tokenizer:
        return prompt[:1000]
    
    tokens = lfm_tokenizer.encode(prompt, truncation=False)
    if len(tokens) <= max_tokens:
        return prompt
    
    # Truncate and add ellipsis
    truncated_tokens = tokens[:max_tokens]
    truncated_text = lfm_tokenizer.decode(truncated_tokens, skip_special_tokens=True)
    return truncated_text + " [Content truncated...]"

# ================== OPTIMIZED STREAMING GENERATION =================
async def generate_reply_stream(
    prompt: str, 
    files: List[UploadFile], 
    max_tokens: int, 
    temperature: float
) -> AsyncGenerator[str, None]:
    """Stream the response with enhanced streaming for LFM2.5"""
    
    print(f"\n🎯 Processing: '{prompt[:50]}...'")
    
    # Validate models
    if not lfm_model or not lfm_tokenizer:
        error_msg = "LFM2.5 model not available. Please check server health."
        print(f"❌ {error_msg}")
        yield json.dumps({
            "type": "error",
            "data": {"error": error_msg}
        }, ensure_ascii=False) + "\n\n"
        return
    
    # Validate prompt length
    if not validate_prompt_length(prompt, MAX_CONTEXT_LENGTH - max_tokens):
        prompt = get_truncated_prompt(prompt, MAX_CONTEXT_LENGTH - max_tokens)
        print("⚠️ Prompt truncated to fit context window")
    
    # --- File Processing ---
    file_text = ""
    if files:
        for f in files:
            try:
                content = await f.read(min(MAX_FILE_SIZE, 1024 * 1024))  # Max 1MB per file
                file_text += f"\n--- File: {f.filename} ---\n"
                file_text += content.decode('utf-8', errors='ignore')[:2000] + "\n"
                print(f"✅ Processed file: {f.filename}")
            except Exception as e:
                print(f"⚠️ File error {f.filename}: {e}")
    
    # ==== ENHANCED LFM2.5 GENERATION WITH TOOL CALL LOOP ====
    max_tool_calls = 5
    tool_calls = 0
    tool_results: List[str] = []
    last_tool_name: Optional[str] = None
    prompt_lower = prompt.lower()
    
    # ==== CONVERSATION HISTORY ====
    conversation_history = []  # Store all messages
    assistant_tool_calls = []  # Store actual tool call texts
    tool_required = any(kw in prompt_lower for kw in [
        "latest", "today", "current", "price", "news", "up to date", "up-to-date",
        "search the web", "search web", "search online", "look up", "lookup",
        "use tool", "use the tool", "web_search", "web search"
    ])

    # Define tools JSON (ONCE)
    tools_json = [
        {
            "name": "web_search",
            "description": "Search the web for current information, facts, or recent updates.",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {"type": "string", "description": "Search query"}
                },
                "required": ["query"]
            }
        }
    ]
    tools_json_str = json.dumps(tools_json, ensure_ascii=False)
    
    print(f"🧰 Tools JSON defined: {tools_json_str[:100]}...")
    
    # ==== INITIALIZE CONVERSATION HISTORY (ONCE) ====
    # Build initial system message (ONCE)
    system_message_content = f"""
You are **GAKR AI**, an expert assistant specializing in depth-calibrated responses. Your core function is to match and anticipate user needs through adaptive response strategies.
## **⚙️ STRICT TOOL USAGE PROTOCOL**
**MANDATORY TOOL CALL REQUIREMENTS:**
1. **Tool Call Trigger Conditions:**
   - ANY request for current/latest/today/tomorrow/now/recent information
   - ANY request for news, prices, stock quotes, cryptocurrency values
   - ANY request requiring factual verification or data retrieval
   - ANY explicit user request for web search
**Critical Requirements:**
- the tool call must start with the <|tool_call_start|> and must end with the <|tool_call_end|> and after the end tool call no extra text
 **Tool Call Format (EXACT AND UNCHANGEABLE):**
<|tool_call_start|>web_search(query="your optimized search query here")<|tool_call_end|>
**Critical Requirements:**
- the tool call start with the <|tool_call_start|> and end with the <|tool_call_end|> and after the end tool call no extra text
- When a tool is required, your first response MUST contain ONLY the tool call (no extra text before or after)
- After tool results return, answer using standard response formatting
- If no tool is needed, answer immediately using standard formatting
**Available Tools:**
<|tool_list_start|> {tools_json_str} <|tool_list_end|>
## **Core Principles**
### 1. **Intent-First Response Design**
- **Foundation Requests** (definitions, basics): Start with clear, complete answers. Add structure only when it clarifies.
- **Deep Dive Requests** ("more details", "explain deeply"): Escalate complexity significantly. Build upon previous information with advanced concepts, trade-offs, and implementation nuances.
- **Comparative Analysis**: Use tables/matrices to highlight distinctions. Include practical selection criteria.
- **Implementation Requests**: Provide executable, well-commented code with configuration considerations.
- **Narrative Requests**: Maintain immersive flow without structural interruptions.
### 2. **Progressive Depth Scaling**
- **Foundation** - Core definition + essential characteristics + 1-2 examples
- **Comprehensive** - Full conceptual breakdown + multiple examples + use cases
- **Expert** - Architecture patterns + optimization strategies + trade-off analysis
- **Mastery** - Industry case studies + advanced implementation + future trends
*Escalate depth based on explicit user requests or implied need for complexity.*
### 3. **Predictive Enhancement** *(When Naturally Relevant)*
- Suggest logical next learning steps
- Warn of common implementation pitfalls
- Recommend high-value resources
- *Apply judiciously—don't force these into simple answers*
### 4. **Output Excellence**
- **Never** include meta-commentary ("Okay, let's tackle this...")
- **Never** explain your structure or reasoning in responses
- **Always** begin with substantive content
- Use complete tables, executable code blocks, and clear formatting
- Maintain professional, accessible tone that matches the query type
## **Response Protocol**
1. Recognize primary intent and depth requirement
2. Deliver comprehensive response matching that level
3. Add predictive elements only when they provide clear value
4. Ensure all information is accurate, logically organized, and practically useful
*Prioritize accuracy and relevance over rigid formatting rules. Adapt structure to content needs.*
"""

    conversation_history.append({"role": "system", "content": system_message_content})
    
    # Original user query (with file content if any)
    user_content = prompt
    if file_text:
        user_content += f"\n\nADDITIONAL FILES:\n{file_text}"
    conversation_history.append({"role": "user", "content": user_content})
    
    print(f"📝 Initialized conversation history with {len(conversation_history)} messages")

    while True:
        # ==== ENHANCED LFM2.5 GENERATION ====
        try:
            # Start with fresh copy for this iteration
            current_messages = conversation_history.copy()
            
            # ==== DEBUG: SHOW EXACT MESSAGE FLOW ====
            print(f"\n🔍 DEBUG: Tool calls history = {tool_calls}")
            print(f"🔍 DEBUG: Conversation history has {len(current_messages)} messages:")
            for i, msg in enumerate(current_messages):
                role = msg["role"]
                content_preview = msg["content"][:150].replace('\n', ' ').strip()
                print(f"  [{i}] {role.upper()}: {content_preview}")
            
            # Check token count BEFORE formatting
            if lfm_tokenizer:
                test_tokens = lfm_tokenizer.apply_chat_template(
                    current_messages,
                    tools=tools_json,
                    tokenize=True,
                    add_generation_prompt=True,
                )
                print(f"🔍 DEBUG: Token count = {len(test_tokens[0])}")
            
            # Apply chat template to format the messages
            try:
                formatted_prompt = lfm_tokenizer.apply_chat_template(
                    current_messages,  # Use the cumulative conversation history
                    tools=tools_json,
                    tokenize=False,
                    add_generation_prompt=True,
                )
            except TypeError:
                formatted_prompt = lfm_tokenizer.apply_chat_template(
                    current_messages,  # Use the cumulative conversation history
                    tokenize=False,
                    add_generation_prompt=True,
                )
            
            print(f"🔍 Formatted prompt length: {len(formatted_prompt)} characters")
            print(f"🔍 Sample formatted prompt: {formatted_prompt[:200]}...")
            
            # Tokenize the formatted prompt
            inputs = lfm_tokenizer(
                formatted_prompt,
                return_tensors="pt",
                truncation=True,
                max_length=MAX_CONTEXT_LENGTH - max_tokens,
            ).to(DEVICE)
            
            input_length = inputs['input_ids'].shape[1]
            print(f"🔍 Tokenized input: {input_length} tokens")
            
            # Create streamer
            streamer = TextIteratorStreamer(
                lfm_tokenizer,
                skip_prompt=True,
                skip_special_tokens=False,
                timeout=300.0
            )
            
            # LFM2.5 OPTIMIZED GENERATION PARAMETERS
            # Based on official recommendations from model card
            generation_config = {
                **inputs,
                "max_new_tokens": min(max_tokens, MAX_TOKENS_HARD_LIMIT),
                "do_sample": temperature > 0,
                "temperature": max(0.1, min(temperature, 1.0)),
                "top_p": 0.1,  # LFM2.5 default
                "top_k": 50,   # LFM2.5 default
                "repetition_penalty": 1.05,  # LFM2.5 default
                "no_repeat_ngram_size": 3,
                "pad_token_id": lfm_tokenizer.pad_token_id,
                "eos_token_id": lfm_tokenizer.eos_token_id,
                "streamer": streamer,
                "use_cache": True,  # Enable KV cache
            }
            
            # For deterministic output when temperature is 0
            if temperature <= 0.1:
                generation_config["do_sample"] = False
                generation_config.pop("temperature", None)
                generation_config.pop("top_p", None)
                generation_config.pop("top_k", None)
            
            # Start generation thread
            thread = Thread(target=lfm_model.generate, kwargs=generation_config)
            thread.start()
            
            print("🔄 Started streaming generation")
            
            # Stream tokens with tool-call detection
            buffer = ""
            response_text = ""
            token_count = 0
            start_time = time.time()
            streaming_started = False
            tool_call = None
            debug_buffer_logged = False
            special_tokens = [
                "<|im_end|>",
                "<|im_start|>",
                "<|endoftext|>",
                "<|startoftext|>",
                "<|tool_call_start|>",
                "<|tool_call_end|>",
                "<|tool_list_start|>",
                "<|tool_list_end|>",
            ]

            def strip_special_tokens(text: str) -> str:
                cleaned = text
                for tok in special_tokens:
                    if tok in cleaned:
                        cleaned = cleaned.replace(tok, "")
                return cleaned

            def normalize_tool_name(name: str) -> str:
                if not name:
                    return ""
                # Remove spaces, convert to lowercase, normalize underscores
                normalized = name.strip().lower().replace(" ", "_").replace("-", "_")
                # Handle variations: "web search", "websearch", "web_search" -> "web_search"
                if "web" in normalized and "search" in normalized:
                    return "web_search"
                if normalized == "websearch":
                    return "web_search"
                return normalized

            def parse_tool_call(text: str) -> Optional[Dict[str, Any]]:
                # ==== ENHANCED DETECTOR: Handle incomplete tool calls & variations ====
                # When <|tool_call_start|> is detected, it's 100% confirmation of a tool call
                # The model may not always generate <|tool_call_end|>, so we extract what we can
                # Also handles variations like "web_ Search" (with space/case variations)
                
                if "<|tool_call_start|>" in text:
                    start_pos = text.find("<|tool_call_start|>") + len("<|tool_call_start|>")
                    
                    # Case 1: Complete token pair - has both start and end
                    if "<|tool_call_end|>" in text[start_pos:]:
                        end_pos = text.find("<|tool_call_end|>", start_pos)
                        payload = text[start_pos:end_pos].strip()
                        print(f"🔍 [PARSE] Complete token pair detected. Payload: {payload[:80]}")
                    else:
                        # Case 2: Incomplete - only start token, extract what we have
                        remaining = text[start_pos:].strip()
                        print(f"🔍 [PARSE] Incomplete token pair. Remaining buffer: {remaining[:80]}")
                        payload = None
                        
                        # Try to extract from brackets first: [web_search(...)] or [web_ Search(...)]
                        # Pattern: [word(args)] - allowing for spaces and case variations
                        bracket_match = re.search(r'\[([a-zA-Z_\s]+?)\s*\((.*?)\)\s*\]', remaining)
                        if bracket_match:
                            func_base = bracket_match.group(1).replace(" ", "_").lower()  # Normalize spaces
                            func_args = bracket_match.group(2)
                            payload = f"{func_base}({func_args})"
                            print(f"🔍 [PARSE] Extracted from brackets (normalized): {payload}")
                        else:
                            # Try bare function call: web_search(...) or web_ Search(...)
                            # Allow for spaces: web_ search, web_search, etc.
                            func_match = re.search(r'([a-zA-Z_]+(?:\s+[a-zA-Z_]+)?)\s*\((.*?)\)', remaining)
                            if func_match:
                                func_base = func_match.group(1).replace(" ", "_").lower()  # Normalize spaces
                                func_args = func_match.group(2)
                                payload = f"{func_base}({func_args})"
                                print(f"🔍 [PARSE] Extracted bare function (normalized): {payload}")
                            else:
                                print(f"⚠️ [PARSE] Could not extract tool call from: {remaining[:50]}")
                                return None
                    
                    # Now parse the extracted payload
                    if payload:
                        # JSON payload
                        if payload.startswith("{") and payload.endswith("}"):
                            try:
                                obj = json.loads(payload)
                                if isinstance(obj, dict) and "tool_call" in obj:
                                    tc = obj.get("tool_call")
                                    if isinstance(tc, dict) and "name" in tc and "args" in tc:
                                        return tc
                            except Exception as e:
                                print(f"⚠️ [PARSE] JSON parsing failed: {e}")

                        # Pythonic function call payload (optionally wrapped in [])
                        clean_payload = payload
                        if clean_payload.startswith("[") and clean_payload.endswith("]"):
                            clean_payload = clean_payload[1:-1].strip()

                        m = re.match(r"^([a-zA-Z_]+(?:\s+[a-zA-Z_]+)?)\s*\((.*)\)\s*$", clean_payload)
                        if m:
                            # Normalize function name: remove spaces, convert to lowercase
                            name_raw = m.group(1)
                            name = normalize_tool_name(name_raw)
                            args_str = m.group(2).strip()
                            args: Dict[str, Any] = {}

                            # Parse query="..." or query='...'
                            mq = re.search(r'query\s*=\s*\"(.*?)\"|query\s*=\s*\'(.*?)\'', args_str)
                            if mq:
                                args["query"] = mq.group(1) or mq.group(2) or ""
                            else:
                                # Positional single string: "..." or '...'
                                mpos = re.match(r'^\"(.*)\"$|^\'(.*)\'$', args_str)
                                if mpos:
                                    pos = mpos.group(1) or mpos.group(2)
                                    if pos is not None:
                                        args["query"] = pos

                            print(f"✅ [PARSE] Tool call extracted: name={name}, query={args.get('query', 'N/A')[:50]}")
                            return {"name": name, "args": args}

                # 2) JSON-only tool call format (legacy - without special tokens)
                stripped = text.strip()
                if stripped.startswith("{") and stripped.endswith("}"):
                    try:
                        obj = json.loads(stripped)
                    except Exception:
                        return None
                    if isinstance(obj, dict) and "tool_call" in obj:
                        tc = obj.get("tool_call")
                        if isinstance(tc, dict) and "name" in tc and "args" in tc:
                            return {"name": normalize_tool_name(tc.get("name", "")), "args": tc.get("args", {})}

                # 3) Bare Pythonic call without special tokens (fallback) - also handle spacing
                stripped = text.strip()
                raw = stripped
                if raw.startswith("[") and raw.endswith("]"):
                    raw = raw[1:-1].strip()
                
                # Handle: web_search(...), web_Search(...), web_ search(...), etc.
                bare_match = re.match(r'([a-zA-Z_]+(?:\s+[a-zA-Z_]+)?)\s*\((.*?)\)\s*$', raw)
                if bare_match:
                    func_name_raw = bare_match.group(1)
                    func_name = normalize_tool_name(func_name_raw)
                    
                    # Only process if it's a recognized tool (web_search)
                    if func_name == "web_search":
                        args_str = bare_match.group(2).strip()
                        args: Dict[str, Any] = {}
                        mq = re.search(r'query\s*=\s*\"(.*?)\"|query\s*=\s*\'(.*?)\'', args_str)
                        if mq:
                            args["query"] = mq.group(1) or mq.group(2) or ""
                        else:
                            mpos = re.match(r'^\"(.*)\"$|^\'(.*)\'$', args_str)
                            if mpos:
                                pos = mpos.group(1) or mpos.group(2)
                                if pos is not None:
                                    args["query"] = pos
                        print(f"✅ [PARSE] Bare tool call found: {func_name}({args})")
                        return {"name": func_name, "args": args}

                return None

            for new_text in streamer:
                if not new_text:
                    continue

                if not streaming_started:
                    buffer += new_text
                    stripped = buffer.strip()

                    # Detect tool call if the model responded with a tool request
                    tool_call = parse_tool_call(buffer)
                    if tool_call:
                        print(f"🧰 Tool call detected: {tool_call}")
                        break

                    # If tool call started but not finished, keep buffering
                    # BUT: if we have <|tool_call_start|> and a tool call pattern, try to extract it
                    if "<|tool_call_start|>" in buffer:
                        if "<|tool_call_end|>" not in buffer:
                            # Check if we can extract a tool call even without end token
                            tool_call = parse_tool_call(buffer)
                            if tool_call:
                                print(f"🧰 Tool call detected (incomplete token pair): {tool_call}")
                                break
                            
                            # Keep buffering if we haven't found a complete pattern yet
                            if not any(pattern in buffer for pattern in ["web_search(", "websearch("]):
                                if not debug_buffer_logged:
                                    print("🧪 Detected <|tool_call_start|> without end token. Buffering more...")
                                    debug_buffer_logged = True
                                continue
                    
                    # If it doesn't look like a tool call, begin streaming
                    if (not stripped.startswith("{")) or (len(stripped) > 400):
                        if tool_required and not tool_results:
                            print("⚠️ Tool required but model did not call tool. Proceeding without tool.")
                        if not debug_buffer_logged:
                            preview = stripped[:300].replace("\n", " ")
                            print(f"🧪 No tool call detected. Starting stream. Buffer preview: {preview}...")
                            debug_buffer_logged = True
                        streaming_started = True

                        # Send tool decision info to client
                        tool_decision = {
                            "use_tool": False,
                            "tool": None,
                            "args": {},
                            "tool_calls": tool_calls,
                            "max_tool_calls": max_tool_calls
                        }
                        yield json.dumps({
                            "type": "context",
                            "data": {
                                "tool_decision": tool_decision,
                                "processing_stage": "tool_decision_made"
                            }
                        }, ensure_ascii=False) + "\n\n"
                        
                        await asyncio.sleep(0.01)

                        # Send streaming start signal
                        yield json.dumps({
                            "type": "streaming_start",
                            "data": {
                                "timestamp": time.time(),
                                "model": LFM_MODEL_NAME
                            }
                        }, ensure_ascii=False) + "\n\n"

                        # Flush buffered text
                        if buffer:
                            token_count += 1
                            cleaned = strip_special_tokens(buffer)
                            if cleaned:
                                response_text += cleaned
                                yield json.dumps({
                                    "type": "token",
                                    "data": cleaned
                                }, ensure_ascii=False) + "\n\n"
                            buffer = ""

                        await asyncio.sleep(0.005)
                    continue

                # Streaming normal tokens
                token_count += 1
                cleaned = strip_special_tokens(new_text)
                if cleaned:
                    response_text += cleaned
                    yield json.dumps({
                        "type": "token",
                        "data": cleaned
                    }, ensure_ascii=False) + "\n\n"
                
                # Small delay for smooth streaming
                await asyncio.sleep(0.005)

            # If a tool call was detected, execute tool and restart loop
            if tool_call:
                thread.join(timeout=2.0)
                tool_calls += 1
                
                # ==== STORE THE TOOL CALL IN CONVERSATION HISTORY ====
                # Get the actual tool call text from buffer
                actual_tool_call = buffer.strip()
                assistant_tool_calls.append(actual_tool_call)
                
                # Add assistant's tool call to conversation history
                conversation_history.append({
                    "role": "assistant",
                    "content": actual_tool_call
                })
                
                print(f"✅ Stored tool call {tool_calls} in conversation history")

                if tool_calls > max_tool_calls:
                    error_msg = "Tool call limit exceeded (max 5)."
                    print(f"❌ {error_msg}")
                    yield json.dumps({
                        "type": "error",
                        "data": {"error": error_msg}
                    }, ensure_ascii=False) + "\n\n"
                    return

                tool_name = normalize_tool_name(tool_call.get("name", ""))
                args = tool_call.get("args", {})
                if not isinstance(args, dict):
                    args = {}

                # Fallback for missing query
                if tool_name == "web_search" and "query" not in args:
                    args["query"] = prompt[:100]

                try:
                    from tools import execute_tool
                    tool_output = execute_tool(tool_name, args)
                except Exception:
                    tool_output = None

                if not tool_output and tool_name == "web_search":
                    try:
                        tool_output = web_search(args.get("query", prompt[:50]))
                    except Exception as e:
                        print(f"Tool web_search error: {e}")
                        tool_output = None

                tool_output = tool_output if tool_output is not None else "No results available."

                last_tool_name = tool_name
                
                tool_results.append(f"TOOL {tool_calls} ({tool_name}):\n{tool_output}")
                
                # ==== ADD TOOL RESULT TO CONVERSATION HISTORY ====
                conversation_history.append({
                    "role": "tool",
                    "content": tool_output
                })
                
                print(f"✅ Tool '{tool_name}' executed and added to history (call {tool_calls}/{max_tool_calls})")
                print(f"📊 Conversation history now has {len(conversation_history)} messages")
                continue

            # If streaming never started (short response), start now and flush buffer
            if not streaming_started:
                # Send tool decision info to client
                tool_decision = {
                    "use_tool": False,
                    "tool": None,
                    "args": {},
                    "tool_calls": tool_calls,
                    "max_tool_calls": max_tool_calls
                }
                yield json.dumps({
                    "type": "context",
                    "data": {
                        "tool_decision": tool_decision,
                        "processing_stage": "tool_decision_made"
                    }
                }, ensure_ascii=False) + "\n\n"
                
                await asyncio.sleep(0.01)

                yield json.dumps({
                    "type": "streaming_start",
                    "data": {
                        "timestamp": time.time(),
                        "model": LFM_MODEL_NAME
                    }
                }, ensure_ascii=False) + "\n\n"

                if buffer:
                    token_count += 1
                    cleaned = strip_special_tokens(buffer)
                    if cleaned:
                        response_text += cleaned
                        yield json.dumps({
                            "type": "token",
                            "data": cleaned
                        }, ensure_ascii=False) + "\n\n"
                    buffer = ""

            # Calculate stats
            elapsed = time.time() - start_time
            tokens_per_sec = token_count / elapsed if elapsed > 0 else 0
            
            print(f"✅ Generation complete: {token_count} tokens, {tokens_per_sec:.1f} t/s")
            
            # Send completion with analysis
            yield json.dumps({
                "type": "done",
                "data": {
                    "tokens": token_count,
                    "chars": len(response_text),
                    "elapsed": round(elapsed, 2),
                    "tokens_per_sec": round(tokens_per_sec, 1),
                    "tool_used": last_tool_name,
                    "model": LFM_MODEL_NAME,
                    "response_characteristics": {
                        "context_informed": bool(tool_results or file_text)
                    }
                }
            }, ensure_ascii=False) + "\n\n"
            
            # Clean up thread
            thread.join(timeout=2.0)
            return
        
        except torch.cuda.OutOfMemoryError:
            error_msg = "CUDA out of memory. Try reducing max_tokens or using CPU."
            print(f"❌ {error_msg}")
            yield json.dumps({
                "type": "error",
                "data": {"error": error_msg}
            }, ensure_ascii=False) + "\n\n"
            return
        
        except Exception as e:
            error_msg = f"Generation error: {str(e)}"
            print(f"❌ {error_msg}")
            traceback.print_exc()
            yield json.dumps({
                "type": "error", 
                "data": {"error": error_msg}
            }, ensure_ascii=False) + "\n\n"
            return
# ================== ENDPOINTS =================

@app.get("/", response_class=HTMLResponse)
async def home(request: Request):
    """Serve homepage"""
    return templates.TemplateResponse("index.html", {
        "request": request,
        "model_name": "LFM2.5-1.2B-Instruct",
        "environment": ENVIRONMENT
    })

@app.post("/api/analyze")
async def analyze(
    prompt: str = Form(...),
    files: Optional[List[UploadFile]] = File(None),
    max_tokens: int = Form(MAX_TOKENS_DEFAULT),
    temperature: float = Form(DEFAULT_TEMPERATURE),
):
    """
    Main analysis endpoint with enhanced intent and emotion detection
    """
    if not prompt and not files:
        raise HTTPException(400, "Empty prompt and no files provided")
    
    # Validate
    max_tokens = min(max_tokens, MAX_TOKENS_HARD_LIMIT)
    temperature = max(0.1, min(temperature, 2.0))
    
    print(f"\n📨 Request: '{prompt[:50]}...' | tokens={max_tokens} | temp={temperature}")
    
    if not lfm_model:
        raise HTTPException(503, "GAKR model not loaded. Check server health.")
    
    async def event_stream():
        try:
            async for chunk in generate_reply_stream(prompt, files or [], max_tokens, temperature):
                yield f"data: {chunk}\n\n"
        except Exception as e:
            yield f"data: {json.dumps({'type': 'error', 'data': str(e)[:200]}, ensure_ascii=False)}\n\n"
    
    return StreamingResponse(
        event_stream(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",
        }
    )

@app.get("/health")
async def health():
    """Health check"""
    return {
        "status": "healthy" if lfm_model else "degraded",
        "models": {
            "main_model": LFM_MODEL_NAME
        },
        "loaded": {
            "lfm_loaded": lfm_model is not None
        },
        "device": DEVICE,
        "cuda_available": torch.cuda.is_available(),
        "max_context": MAX_CONTEXT_LENGTH,
        "timestamp": datetime.now(timezone.utc).isoformat()
    }

@app.get("/api/config")
async def config():
    """Configuration endpoint"""
    return {
        "models": {
            "main_model": LFM_MODEL_NAME
        },
        "capabilities": {
            "web_search": True,
            "tool_integration": True,
            "file_processing": True,
            "streaming_responses": True
        },
        "limits": {
            "max_tokens": MAX_TOKENS_HARD_LIMIT,
            "max_context": MAX_CONTEXT_LENGTH,
            "default_temperature": DEFAULT_TEMPERATURE
        },
        "features": [
            "tool_integration", 
            "file_upload", 
            "streaming",
            "conversation_history"
        ]
    }

@app.get("/api/model-info")
async def model_info():
    """Detailed model information"""
    info = {
        "models": {
            "lfm2.5": {
                "name": "LFM2.5-1.2B-Instruct",
                "publisher": "Liquid AI",
                "parameters": "1.2B",
                "context_length": "32,768 tokens",
                "license": "Apache 2.0",
                "description": "Hybrid model designed for on-device deployment with extended pre-training and reinforcement learning",
                "capabilities": ["Reasoning", "Coding", "Mathematics", "Multilingual (8 languages)", "Tool Use"],
                "optimized_for": ["Edge deployment", "Low-latency applications", "Memory-constrained environments"],
                "default_generation_params": {
                    "temperature": 0.1,
                    "top_p": 0.1,
                    "top_k": 50,
                    "repetition_penalty": 1.05
                },
                "chat_template": "ChatML format (<|im_start|>role\\ncontent<|im_end|>)",
                "special_tokens": ["<|im_start|>", "<|im_end|>", "<|endoftext|>"]
            }
        }
    }
    
    if lfm_model:
        info["models"]["lfm2.5"].update({
            "loaded": True,
            "device": str(lfm_model.device),
            "dtype": str(lfm_model.dtype),
            "tokenizer_chat_template": lfm_tokenizer.chat_template is not None
        })
    
    return info

@app.post("/api/test-lfm-template")
async def test_lfm_template(text: str = Form(...), system_prompt: str = Form(None)):
    """Test LFM2.5 chat template formatting"""
    if not lfm_tokenizer:
        raise HTTPException(503, "LFM2.5 tokenizer not loaded")
    
    system_msg = system_prompt or "You are a helpful AI assistant."
    
    messages = [
        {"role": "system", "content": system_msg},
        {"role": "user", "content": text}
    ]
    
    try:
        # Apply chat template
        formatted = lfm_tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True
        )
        
        # Decode to see the format
        decoded = lfm_tokenizer.decode(lfm_tokenizer.encode(formatted), skip_special_tokens=False)
        
        return {
            "formatted_template": formatted,
            "decoded": decoded,
            "messages": messages,
            "special_tokens": {
                "im_start": "<|im_start|>",
                "im_end": "<|im_end|>",
                "eos_token": lfm_tokenizer.eos_token,
                "pad_token": lfm_tokenizer.pad_token
            }
        }
    except Exception as e:
        raise HTTPException(500, f"Template error: {str(e)}")

@app.get("/api/debug-stream")
async def debug_stream():
    """Test streaming with LFM2.5"""
    async def event_stream():
        yield f"data: {json.dumps({'type': 'context', 'data': {'test': 'LFM2.5 streaming test'}}, ensure_ascii=False)}\n\n"
        
        test_msg = "This is a streaming test from LFM2.5-1.2B-Instruct backend. "
        words = test_msg.split()
        
        for i, word in enumerate(words):
            yield f"data: {json.dumps({'type': 'token', 'data': word + ' '}, ensure_ascii=False)}\n\n"
            await asyncio.sleep(0.1)
        
        yield f"data: {json.dumps({'type': 'done', 'data': {'test': 'complete'}}, ensure_ascii=False)}\n\n"
    
    return StreamingResponse(
        event_stream(),
        media_type="text/event-stream"
    )

# ================== ERROR HANDLERS =================

@app.exception_handler(HTTPException)
async def http_exception_handler(request: Request, exc: HTTPException):
    return JSONResponse(
        status_code=exc.status_code,
        content={"error": exc.detail}
    )

@app.exception_handler(Exception)
async def generic_exception_handler(request: Request, exc: Exception):
    print(f"❌ Unhandled: {exc}")
    traceback.print_exc()
    return JSONResponse(
        status_code=500,
        content={"error": "Internal server error"}
    )

# ================== MAIN =================

if __name__ == "__main__":
    import uvicorn
    
    print("\n" + "="*50)
    print("🔧 STARTING GAKR AI ASSISTANT")
    print("="*50)
    print(f"🧠 GAKR Model: {LFM_MODEL_NAME}")
    print("="*50)
    
    uvicorn.run(
        "app:app",
        host="0.0.0.0",
        port=7860,
        reload=True,
        log_level="info"
    )