Update app.py

app.py

@@ -1,522 +1,28 @@
-from flask import Flask, request, jsonify
-from mistral_inference.transformer import Transformer
-from mistral_inference.generate import generate
-from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
-from mistral_common.protocol.instruct.messages import UserMessage, AssistantMessage, ToolMessage
-from mistral_common.protocol.instruct.request import ChatCompletionRequest
-from mistral_common.protocol.instruct.tool_calls import Function, Tool
-import requests
-import json
-import os
-import threading
-import queue
-import time
-from datetime import datetime
-from pathlib import Path
-import logging
-
-# Configure logging
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+# Load the model
+model_name = "Qwen/Qwen2.5-0.5B-Instruct"  # Or 1.5B, 3B, 7B
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype="auto",
+    device_map="auto"
 )
-logger = logging.getLogger(__name__)
-
-app = Flask(__name__)
-
-# Global variables
-model = None
-tokenizer = None
-MODEL_PATH = "/app/models"
-MEMORY_FILE = "/app/memory.json"
-MAX_MEMORY_ENTRIES_TOTAL = 500  # Total shared memory entries for the group
-MAX_MEMORY_CONTEXT = 20  # Number of recent messages to include in context
-
-# Request queue for handling concurrent requests
-request_queue = queue.Queue(maxsize=100)
-processing_lock = threading.Lock()
-
-class MemoryManager:
-    """Manages shared group memories with size limits"""
-    
-    def __init__(self, memory_file):
-        self.memory_file = memory_file
-        self.memory_lock = threading.Lock()
-        self.load_memory()
-    
-    def load_memory(self):
-        """Load memory from file"""
-        try:
-            if os.path.exists(self.memory_file):
-                with open(self.memory_file, 'r') as f:
-                    data = json.load(f)
-                    # Support both old format and new shared format
-                    if isinstance(data, dict) and "entries" in data:
-                        self.memory = data
-                    else:
-                        # Convert old format to new
-                        self.memory = {"entries": [], "users": {}}
-            else:
-                self.memory = {"entries": [], "users": {}}
-        except Exception as e:
-            logger.error(f"Error loading memory: {e}")
-            self.memory = {"entries": [], "users": {}}
-    
-    def save_memory(self):
-        """Save memory to file"""
-        try:
-            with open(self.memory_file, 'w') as f:
-                json.dump(self.memory, f, indent=2)
-        except Exception as e:
-            logger.error(f"Error saving memory: {e}")
-    
-    def add_memory(self, chat_id, username, name, message):
-        """Add a memory entry to shared group memory"""
-        with self.memory_lock:
-            chat_key = str(chat_id)
-            
-            # Track user info
-            if chat_key not in self.memory["users"]:
-                self.memory["users"][chat_key] = {
-                    "username": username,
-                    "name": name,
-                    "first_seen": datetime.now().isoformat(),
-                    "message_count": 0
-                }
-            
-            # Update user info
-            self.memory["users"][chat_key]["username"] = username
-            self.memory["users"][chat_key]["name"] = name
-            self.memory["users"][chat_key]["message_count"] += 1
-            self.memory["users"][chat_key]["last_seen"] = datetime.now().isoformat()
-            
-            # Add new entry to shared memory
-            entry = {
-                "chat_id": chat_key,
-                "message": message,
-                "timestamp": datetime.now().isoformat(),
-                "username": username,
-                "name": name
-            }
-            
-            self.memory["entries"].append(entry)
-            
-            # Trim if too large - keep most recent entries
-            if len(self.memory["entries"]) > MAX_MEMORY_ENTRIES_TOTAL:
-                # Keep first 50 (important early context) and last 450 (recent context)
-                important = self.memory["entries"][:50]
-                recent = self.memory["entries"][-(MAX_MEMORY_ENTRIES_TOTAL-50):]
-                self.memory["entries"] = important + recent
-            
-            self.save_memory()
-            
-            return {
-                "status": "success",
-                "total_entries": len(self.memory["entries"]),
-                "user_message_count": self.memory["users"][chat_key]["message_count"]
-            }
-    
-    def get_memory(self, chat_id):
-        """Get memory for a specific chat (for compatibility)"""
-        with self.memory_lock:
-            chat_key = str(chat_id)
-            if chat_key in self.memory["users"]:
-                # Return user-specific view
-                user_entries = [e for e in self.memory["entries"] if e.get("chat_id") == chat_key]
-                return {
-                    "username": self.memory["users"][chat_key]["username"],
-                    "name": self.memory["users"][chat_key]["name"],
-                    "entries": user_entries
-                }
-            return None
-    
-    def get_all_memory(self):
-        """Get all shared group memory"""
-        with self.memory_lock:
-            return self.memory
-    
-    def get_memory_summary(self, chat_id=None):
-        """Get a condensed summary of shared group memories for context"""
-        with self.memory_lock:
-            if not self.memory.get("entries"):
-                return ""
-            
-            # Get recent entries for context
-            recent_entries = self.memory["entries"][-MAX_MEMORY_CONTEXT:]
-            
-            summary = f"\n--- Shared Group Memory Context ---\n"
-            summary += f"Group has {len(self.memory['users'])} members\n"
-            summary += f"Recent conversation history:\n"
-            
-            for entry in recent_entries:
-                timestamp = entry.get('timestamp', '')[:16]  # Date + time
-                name = entry.get('name', 'Unknown')
-                username = entry.get('username', 'unknown')
-                msg = entry.get('message', '')[:150]  # Truncate long messages
-                summary += f"- [{timestamp}] {name} (@{username}): {msg}\n"
-            
-            summary += "--- End Memory Context ---\n"
-            return summary
-    
-    def search_memory(self, keyword, chat_id=None):
-        """Search for specific information in memory"""
-        with self.memory_lock:
-            results = []
-            for entry in self.memory["entries"]:
-                # If chat_id provided, filter to that user
-                if chat_id and entry.get("chat_id") != str(chat_id):
-                    continue
-                    
-                if keyword.lower() in entry["message"].lower():
-                    results.append(entry)
-            
-            return results
-    
-    def get_user_info(self, chat_id):
-        """Get information about a specific user"""
-        with self.memory_lock:
-            chat_key = str(chat_id)
-            return self.memory["users"].get(chat_key)
 
-# Initialize memory manager
-memory_manager = MemoryManager(MEMORY_FILE)
-
-def initialize_model():
-    """Initialize the Mistral model and tokenizer"""
-    global model, tokenizer
+# Chat function
+def chat(prompt):
+    messages = [{"role": "user", "content": prompt}]
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
     
-    if model is None or tokenizer is None:
-        logger.info("Loading Mistral model...")
-        try:
-            # Download model if not present
-            tokenizer_path = f"{MODEL_PATH}/tokenizer.model.v3"
-            if not os.path.exists(tokenizer_path):
-                logger.info("Model not found. Downloading from HuggingFace...")
-                logger.info("This will take several minutes on first run...")
-                
-                os.makedirs(MODEL_PATH, exist_ok=True)
-                
-                import subprocess
-                result = subprocess.run(
-                    ["git", "clone", "https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.3", MODEL_PATH],
-                    capture_output=True,
-                    text=True,
-                    timeout=1800  # 30 minute timeout
-                )
-                
-                if result.returncode != 0:
-                    logger.error(f"Failed to download model: {result.stderr}")
-                    raise Exception(f"Model download failed: {result.stderr}")
-                
-                logger.info("Model downloaded successfully!")
-            
-            tokenizer = MistralTokenizer.from_file(tokenizer_path)
-            model = Transformer.from_folder(MODEL_PATH)
-            logger.info("Model loaded successfully")
-        except Exception as e:
-            logger.error(f"Error loading model: {e}")
-            raise
-
-def check_health_api():
-    """Call the health API"""
-    try:
-        response = requests.get("https://bridge.zone.id/api/health", timeout=5)
-        return response.json()
-    except Exception as e:
-        logger.error(f"Health API error: {e}")
-        return {"error": str(e), "status": "error"}
-
-def extract_tool_calls(text):
-    """Extract tool calls from model output more robustly"""
-    try:
-        # Look for function call patterns
-        if "[TOOL_CALLS]" in text or "check_health" in text.lower():
-            return True
-        return False
-    except Exception as e:
-        logger.error(f"Error extracting tool calls: {e}")
-        return False
-
-def process_with_mistral(query, chat_id=None, use_tools=False, max_retries=2):
-    """Process query with Mistral model with error handling"""
+    inputs = tokenizer([text], return_tensors="pt").to(model.device)
+    outputs = model.generate(**inputs, max_new_tokens=512)
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     
-    for attempt in range(max_retries):
-        try:
-            initialize_model()
-            
-            # Prepare tools if needed
-            tools = None
-            if use_tools:
-                tools = [
-                    Tool(
-                        function=Function(
-                            name="check_health",
-                            description="Check the health status of the bridge.zone.id API service",
-                            parameters={
-                                "type": "object",
-                                "properties": {},
-                                "required": [],
-                            },
-                        )
-                    )
-                ]
-            
-            # Build context with shared group memory
-            context = ""
-            memory_context = memory_manager.get_memory_summary(chat_id)
-            if memory_context:
-                context = memory_context + "\n"
-            
-            # Create user message with context
-            user_content = context + query if context else query
-            messages = [UserMessage(content=user_content)]
-            
-            completion_request = ChatCompletionRequest(
-                messages=messages,
-                tools=tools if use_tools else None
-            )
-            
-            tokens = tokenizer.encode_chat_completion(completion_request).tokens
-            
-            # Limit context to prevent overflow
-            max_context_tokens = 4096  # Adjust based on your model
-            if len(tokens) > max_context_tokens:
-                tokens = tokens[-max_context_tokens:]
-            
-            with processing_lock:  # Ensure single-threaded model access
-                out_tokens, _ = generate(
-                    [tokens], 
-                    model, 
-                    max_tokens=512, 
-                    temperature=0.7, 
-                    eos_id=tokenizer.instruct_tokenizer.tokenizer.eos_id
-                )
-            
-            result = tokenizer.instruct_tokenizer.tokenizer.decode(out_tokens[0])
-            
-            # Handle tool calls
-            if use_tools and extract_tool_calls(result):
-                logger.info("Tool call detected, executing health check")
-                health_data = check_health_api()
-                
-                # Create follow-up request with tool result
-                messages.append(AssistantMessage(content=result))
-                messages.append(ToolMessage(
-                    content=json.dumps(health_data),
-                    tool_call_id="check_health_1"
-                ))
-                
-                completion_request = ChatCompletionRequest(messages=messages)
-                tokens = tokenizer.encode_chat_completion(completion_request).tokens
-                
-                if len(tokens) > max_context_tokens:
-                    tokens = tokens[-max_context_tokens:]
-                
-                with processing_lock:
-                    out_tokens, _ = generate(
-                        [tokens], 
-                        model, 
-                        max_tokens=512, 
-                        temperature=0.7, 
-                        eos_id=tokenizer.instruct_tokenizer.tokenizer.eos_id
-                    )
-                
-                result = tokenizer.instruct_tokenizer.tokenizer.decode(out_tokens[0])
-            
-            return {"success": True, "response": result}
-            
-        except Exception as e:
-            logger.error(f"Attempt {attempt + 1} failed: {e}")
-            if attempt == max_retries - 1:
-                return {"success": False, "error": str(e)}
-            time.sleep(1)  # Brief pause before retry
-    
-    return {"success": False, "error": "Max retries exceeded"}
-
-@app.route('/api/chat', methods=['POST'])
-def chat():
-    """Main chat endpoint with queue management"""
-    try:
-        data = request.get_json()
-        
-        if not data or 'query' not in data:
-            return jsonify({
-                "error": "Missing 'query' parameter"
-            }), 400
-        
-        query = data['query'].strip()
-        chat_id = data.get('chat_id')
-        
-        if not query:
-            return jsonify({
-                "error": "Empty query"
-            }), 400
-        
-        # Check queue size
-        if request_queue.qsize() >= 95:
-            return jsonify({
-                "error": "Server is busy, please try again later"
-            }), 503
-        
-        # Determine if we should use function calling
-        health_keywords = ['health', 'status', 'api status', 'bridge', 'check service', 'api check']
-        use_tools = any(keyword in query.lower() for keyword in health_keywords)
-        
-        # Process with Mistral
-        start_time = time.time()
-        result = process_with_mistral(query, chat_id=chat_id, use_tools=use_tools)
-        processing_time = time.time() - start_time
-        
-        if not result["success"]:
-            return jsonify({
-                "error": result.get("error", "Processing failed")
-            }), 500
-        
-        logger.info(f"Processed query in {processing_time:.2f}s")
-        
-        return jsonify({
-            "query": query,
-            "response": result["response"],
-            "used_tools": use_tools,
-            "processing_time": processing_time,
-            "has_memory": len(memory_manager.memory.get("entries", [])) > 0,
-            "total_group_memories": len(memory_manager.memory.get("entries", []))
-        })
-        
-    except Exception as e:
-        logger.error(f"Chat endpoint error: {e}")
-        return jsonify({
-            "error": str(e)
-        }), 500
-
-@app.route('/api/update-memory', methods=['POST'])
-def update_memory():
-    """Update memory for a user"""
-    try:
-        data = request.get_json()
-        
-        required_fields = ['chat_id', 'username', 'name', 'message']
-        if not all(field in data for field in required_fields):
-            return jsonify({
-                "error": f"Missing required fields: {required_fields}"
-            }), 400
-        
-        result = memory_manager.add_memory(
-            chat_id=data['chat_id'],
-            username=data['username'],
-            name=data['name'],
-            message=data['message']
-        )
-        
-        return jsonify(result)
-        
-    except Exception as e:
-        logger.error(f"Update memory error: {e}")
-        return jsonify({
-            "error": str(e)
-        }), 500
+    return response
 
-@app.route('/api/get-memory/<chat_id>', methods=['GET'])
-def get_memory(chat_id):
-    """Get memory for a specific user"""
-    try:
-        memory_data = memory_manager.get_memory(chat_id)
-        
-        if memory_data is None:
-            return jsonify({
-                "error": "No memory found for this user"
-            }), 404
-        
-        return jsonify(memory_data)
-        
-    except Exception as e:
-        logger.error(f"Get memory error: {e}")
-        return jsonify({
-            "error": str(e)
-        }), 500
-
-@app.route('/api/search-memory/<chat_id>', methods=['POST'])
-def search_memory(chat_id):
-    """Search memory for specific information"""
-    try:
-        data = request.get_json()
-        
-        if not data or 'keyword' not in data:
-            return jsonify({
-                "error": "Missing 'keyword' parameter"
-            }), 400
-        
-        # If chat_id is "all", search entire group memory
-        search_chat_id = None if chat_id == "all" else chat_id
-        results = memory_manager.search_memory(data['keyword'], chat_id=search_chat_id)
-        
-        return jsonify({
-            "results": results,
-            "count": len(results)
-        })
-        
-    except Exception as e:
-        logger.error(f"Search memory error: {e}")
-        return jsonify({
-            "error": str(e)
-        }), 500
-
-@app.route('/api/get-all-memory', methods=['GET'])
-def get_all_memory():
-    """Get all shared group memory"""
-    try:
-        memory_data = memory_manager.get_all_memory()
-        return jsonify(memory_data)
-        
-    except Exception as e:
-        logger.error(f"Get all memory error: {e}")
-        return jsonify({
-            "error": str(e)
-        }), 500
-
-@app.route('/api/get-user-info/<chat_id>', methods=['GET'])
-def get_user_info(chat_id):
-    """Get information about a specific user"""
-    try:
-        user_info = memory_manager.get_user_info(chat_id)
-        
-        if user_info is None:
-            return jsonify({
-                "error": "User not found"
-            }), 404
-        
-        return jsonify(user_info)
-        
-    except Exception as e:
-        logger.error(f"Get user info error: {e}")
-        return jsonify({
-            "error": str(e)
-        }), 500
-
-@app.route('/health', methods=['GET'])
-def health():
-    """Health check endpoint"""
-    return jsonify({
-        "status": "ok",
-        "model_loaded": model is not None and tokenizer is not None,
-        "queue_size": request_queue.qsize(),
-        "total_group_members": len(memory_manager.memory.get("users", {})),
-        "total_group_messages": len(memory_manager.memory.get("entries", []))
-    })
-
-@app.route('/api/check-bridge-health', methods=['GET'])
-def check_bridge_health():
-    """Direct endpoint to check bridge API health"""
-    health_data = check_health_api()
-    return jsonify(health_data)
-
-if __name__ == '__main__':
-    # Initialize model on startup
-    try:
-        initialize_model()
-    except Exception as e:
-        logger.error(f"Warning: Could not initialize model on startup: {e}")
-        logger.info("Model will be loaded on first request")
-    
-    # Run with threading enabled
-    app.run(host='0.0.0.0', port=7860, debug=False, threaded=True)
+# Test it
+print(chat("Hello, how are you?"))