Create app.py

app.py

@@ -0,0 +1,500 @@
+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'
+)
+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
+    
+    if model is None or tokenizer is None:
+        logger.info("Loading Mistral model...")
+        try:
+            tokenizer = MistralTokenizer.from_file(f"{MODEL_PATH}/tokenizer.model.v3")
+            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"""
+    
+    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
+
+@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)