""" 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" )