Update app.py

app.py

@@ -1,312 +1,197 @@
-import os
-import sys
-import gc
-import json
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request
+from fastapi.responses import HTMLResponse
+from fastapi.templating import Jinja2Templates
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+import asyncio
 import logging
-import traceback
-from datetime import datetime
-from pathlib import Path
-from typing import Dict, Optional, Set, List, Any
-from collections import deque
-from contextlib import contextmanager
-import psutil
-import tempfile
-
-# Set base directory to current file's directory
-BASE_DIR = Path(__file__).parent
-
-# Set environment variables before any other imports
-temp_cache_dir = tempfile.gettempdir()
-os.environ["HF_HOME"] = str(temp_cache_dir)
-if "TRANSFORMERS_CACHE" in os.environ:
-    del os.environ["TRANSFORMERS_CACHE"]
-
-os.environ["TRANSFORMERS_PARALLELISM"] = "false"
-os.environ["TORCH_HOME"] = str(Path(temp_cache_dir) / "torch")
-
-# Critical ML imports
+import httpx
 import torch
-import transformers
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 
-# Additional imports
-import aiofiles
-import aiodns
-import httpx
-import uvicorn
-from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request, HTTPException, Response, Query
-from fastapi.responses import HTMLResponse
-from fastapi.staticfiles import StaticFiles
-from fastapi.templating import Jinja2Templates
-from prometheus_client import Counter, Gauge, Histogram, CollectorRegistry
-from prometheus_fastapi_instrumentator import Instrumentator
-import hazm
-import backoff
-from logging.handlers import RotatingFileHandler
-from cachetools import TTLCache
-
-class AppConfig:
-    BASE_DIR = Path("/app")
-    STATIC_DIR = BASE_DIR / "static"
-    TEMPLATES_DIR = BASE_DIR / "templates"
-    CACHE_DIR = Path(temp_cache_dir) / "cache"
-    LOG_DIR = Path(temp_cache_dir) / "logs"
-
-    # Performance settings
-    MAX_RETRIES = 3
-    TIMEOUT = 60
-    MAX_CONNECTIONS = 5
-    CACHE_TTL = 300
-    MEMORY_THRESHOLD = 85
-    CLEANUP_THRESHOLD = 80
-    MAX_MESSAGE_LENGTH = 512
-    MAX_THREADS = 1
-    MODEL_MAX_LENGTH = 128  # Increased from 16 to allow longer responses
-
-    # Model settings
-    MODEL_NAME = "bigscience/bloom-560m"  # Upgraded from tiny-gpt2
-    MODEL_BATCH_SIZE = 1
-    MODEL_MAX_LENGTH = 128
-
-    # Network settings
-    PROXY_URL = "http://your-proxy-url:port"  # Replace with your proxy if needed
-
-    @classmethod
-    def setup_directories(cls) -> None:
-        try:
-            for path in [cls.CACHE_DIR, cls.LOG_DIR, cls.STATIC_DIR, cls.TEMPLATES_DIR]:
-                path.mkdir(exist_ok=True, parents=True)
-            logging.info("Directory setup completed successfully")
-        except Exception as e:
-            logging.error(f"Failed to setup directories: {e}")
-            raise
+# Initialize FastAPI
+app = FastAPI()
 
-class BloomAI:
+# Logging setup
+logging.basicConfig(level=logging.INFO)
+
+# Telegram Token and Chat ID (Replace with your actual values)
+TELEGRAM_TOKEN = "7437859619:AAGeGG3ZkLM0OVaw-Exx1uMRE55JtBCZZCY"
+CHAT_ID = "-1002228627548"
+
+# Templating setup
+templates = Jinja2Templates(directory="templates")
+
+
+# WebSocket Manager
+class WebSocketManager:
     def __init__(self):
-        try:
-            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-            torch.set_num_threads(AppConfig.MAX_THREADS)
-            if self.device.type == 'cpu':
-                torch.set_num_interop_threads(1)
-            self.pipeline = None
-            self.model_loaded = False
-            self.load_attempts = 0
-            self.max_attempts = AppConfig.MAX_RETRIES
-            self.model_name = AppConfig.MODEL_NAME
-            self.tokenizer = None
-            logging.info(f"BloomAI initialized with device: {self.device}")
-        except Exception as e:
-            logging.error(f"Failed to initialize BloomAI: {e}")
-            raise
-
-    def initialize(self) -> bool:
-        try:
-            logging.info(f"Loading AI model {self.model_name} on {self.device}")
-
-            # Load tokenizer first
-            self.tokenizer = AutoTokenizer.from_pretrained(
-                self.model_name,
-                cache_dir=AppConfig.CACHE_DIR,
-                use_fast=True
-            )
-
-            # Configure model loading with optimizations
-            model = AutoModelForCausalLM.from_pretrained(
-                self.model_name,
-                cache_dir=AppConfig.CACHE_DIR,
-                torch_dtype=torch.float32,
-                low_cpu_mem_usage=True,
-                device_map="auto" if torch.cuda.is_available() else None
-            )
-
-            # Create pipeline with optimized configuration
-            self.pipeline = pipeline(
-                "text-generation",
-                model=model,
-                tokenizer=self.tokenizer,
-                device=self.device,
-                framework="pt",
-                model_kwargs={
-                    "pad_token_id": self.tokenizer.eos_token_id
-                }
-            )
-
-            # Clean up memory
-            gc.collect()
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
-
-            self.model_loaded = True
-            return True
-
-        except Exception as e:
-            logging.error(f"Model load error: {e}")
-            traceback.print_exc()
-            return False
-
-    def generate_response(self, text: str) -> str:
-        try:
-            if not self.pipeline:
-                return "مدل در دسترس نیست."
-
-            with RESPONSE_TIME.labels(endpoint="/generate_response").time():
-                outputs = self.pipeline(
-                    text,
-                    max_length=AppConfig.MODEL_MAX_LENGTH,
-                    min_length=20,
-                    do_sample=True,
-                    top_k=50,
-                    top_p=0.92,
-                    temperature=0.7,
-                    num_return_sequences=1,
-                    no_repeat_ngram_size=3,
-                    pad_token_id=self.tokenizer.eos_token_id,
-                    attention_mask=None,
-                    early_stopping=True,
-                    repetition_penalty=1.2
-                )
-
-            # Post-process the response
-            response = outputs[0]['generated_text']
-            response = response.strip()
-
-            # Remove the input prompt if included
-            if response.startswith(text):
-                response = response[len(text):].strip()
-
-            REQUESTS.labels(endpoint="/generate_response").inc()
-            return response
-
-        except Exception as e:
-            logging.error(f"Error generating response: {e}", exc_info=True)
-            return "خطا در تولید پاسخ."
-
-# FastAPI Application Setup
-app = FastAPI(title="BLOOM AI Assistant")
-
-# Mount static files and templates
-app.mount("/static", StaticFiles(directory=AppConfig.STATIC_DIR), name="static")
-templates = Jinja2Templates(directory=AppConfig.TEMPLATES_DIR)
-
-# Setup Prometheus metrics
-custom_registry = CollectorRegistry()
-
-REQUESTS = Counter(
-    "bloom_http_requests_total",
-    "Total number of HTTP requests made",
-    labelnames=["endpoint"],
-    registry=custom_registry
-)
-
-RESPONSE_TIME = Histogram(
-    "bloom_http_response_seconds",
-    "HTTP response time in seconds",
-    labelnames=["endpoint"],
-    registry=custom_registry
-)
-
-MEMORY_USAGE = Gauge(
-    "bloom_system_memory_bytes",
-    "System memory usage in bytes",
-    registry=custom_registry
-)
-
-# Setup Prometheus instrumentation
-instrumentator = Instrumentator(
-    registry=custom_registry,
-    should_group_status_codes=False,
-    should_ignore_untemplated=True,
-    should_respect_env_var=True,
-    should_instrument_requests_inprogress=True,
-    excluded_handlers=["/metrics"]
-)
-
-@app.middleware("http")
-def add_custom_metrics(request: Request, call_next):
-    response = call_next(request)
-    REQUESTS.labels(endpoint=request.url.path).inc()
-    return response
-
-instrumentator.instrument(app).expose(app)
-
-# Default HTML template
-DEFAULT_HTML = """
-<!DOCTYPE html>
-<html dir="rtl" lang="fa">
-<head>
-    <meta charset="UTF-8">
-    <title>هوش مصنوعی BLOOM</title>
-    <style>
-        body {
-            font-family: Tahoma, Arial;
-            text-align: center;
-            margin-top: 50px;
-            background-color: #f0f2f5;
-        }
-        .container {
-            max-width: 600px;
-            margin: 0 auto;
-            padding: 20px;
-            background-color: white;
-            border-radius: 10px;
-            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
-        }
-        h1 { color: #1a73e8; }
-        .status { margin: 20px 0; }
-    </style>
-</head>
-<body>
-    <div class="container">
-        <h1>سیستم هوش مصنوعی BLOOM</h1>
-        <div class="status">
-            <p>وضعیت سیستم: <span id="status">در حال بارگذاری...</span></p>
-        </div>
-    </div>
-</body>
-</html>
-"""
-
-# Route handlers
-@app.get("/", response_class=HTMLResponse)
-def home(request: Request):
-    try:
-        ai_ready = request.app.state.assistant.is_ready.is_set()
-        return templates.TemplateResponse(
-            "index.html",
-            {
-                "request": request,
-                "ai_status": "آماده" if ai_ready else "در حال آماده‌سازی"
-            }
-        )
-    except Exception as e:
-        logging.error(f"Template error: {e}")
-        return HTMLResponse(content=DEFAULT_HTML)
+        self.active_connection: WebSocket = None
+
+    async def connect(self, websocket: WebSocket):
+        """Connects the WebSocket"""
+        await websocket.accept()
+        self.active_connection = websocket
+        logging.info("WebSocket connected.")
+
+    async def disconnect(self):
+        """Disconnects the WebSocket"""
+        if self.active_connection:
+            await self.active_connection.close()
+            self.active_connection = None
+            logging.info("WebSocket disconnected.")
+
+    async def send_message(self, message: str):
+        """Sends a message through WebSocket"""
+        if self.active_connection:
+            await self.active_connection.send_text(message)
+            logging.info(f"Sent via WebSocket: {message}")
+
 
-@app.exception_handler(404)
-def not_found_handler(request: Request, exc: Exception):
-    return HTMLResponse(content=DEFAULT_HTML, status_code=404)
+websocket_manager = WebSocketManager()
 
-@app.exception_handler(500)
-def server_error_handler(request: Request, exc: Exception):
-    return HTMLResponse(content=DEFAULT_HTML, status_code=500)
 
+# BLOOM Model Manager
+class BloomAI:
+    def __init__(self):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.pipeline = None
+
+    def load_model(self):
+        """Loads BLOOM AI Model"""
+        logging.info("Loading BLOOM model...")
+        tokenizer = AutoTokenizer.from_pretrained("bigscience/bloom-560m")
+        model = AutoModelForCausalLM.from_pretrained("bigscience/bloom-560m")
+        self.pipeline = pipeline(
+            "text-generation",
+            model=model,
+            tokenizer=tokenizer,
+            device=0 if torch.cuda.is_available() else -1
+        )
+        logging.info("BLOOM model loaded successfully.")
+
+    async def generate_response(self, prompt: str) -> str:
+        """Generates a response using BLOOM"""
+        if not prompt.strip():
+            return "⚠️ Please send a valid message."
+        logging.info(f"Generating response for prompt: {prompt}")
+        outputs = self.pipeline(
+            prompt,
+            max_length=100,
+            do_sample=True,
+            temperature=0.7,
+            top_k=50,
+            top_p=0.9,
+            num_return_sequences=1,
+            no_repeat_ngram_size=2
+        )
+        response = outputs[0]["generated_text"]
+        return response.strip()
+
+
+# Initialize BLOOM
+bloom_ai = BloomAI()
+bloom_ai.load_model()
+
+
+# Telegram Message Handling
+async def send_telegram_message(text: str):
+    """Sends a message to Telegram"""
+    async with httpx.AsyncClient() as client:
+        url = f"https://api.telegram.org/bot{TELEGRAM_TOKEN}/sendMessage"
+        payload = {"chat_id": CHAT_ID, "text": text}
+        response = await client.post(url, json=payload)
+        if response.status_code == 200:
+            logging.info(f"Sent to Telegram: {text}")
+        else:
+            logging.error(f"Failed to send message to Telegram: {response.text}")
+
+
+@app.post("/telegram")
+async def telegram_webhook(update: dict):
+    """Handles Telegram Webhook messages"""
+    if "message" in update:
+        chat_id = str(update["message"]["chat"]["id"])
+        if chat_id != CHAT_ID:
+            return {"status": "Unauthorized"}
+        
+        user_message = update["message"]["text"]
+        logging.info(f"Received from Telegram: {user_message}")
+        
+        # Process the message
+        response = await bloom_ai.generate_response(user_message)
+        await send_telegram_message(response)
+    return {"status": "ok"}
+
+
+# WebSocket Endpoint
 @app.websocket("/ws")
-def websocket_endpoint(websocket: WebSocket, client_id: str = Query(...)):
-    websocket.accept()
+async def websocket_endpoint(websocket: WebSocket):
+    """WebSocket communication for real-time interaction"""
+    await websocket_manager.connect(websocket)
     try:
         while True:
-            data = websocket.receive_text()
-            response = app.state.assistant.process_message(data, client_id)
-            websocket.send_text(response)
+            # Receive message from WebSocket
+            data = await websocket.receive_text()
+            logging.info(f"Received from WebSocket: {data}")
+            
+            # Process the message
+            response = await bloom_ai.generate_response(data)
+            
+            # Send response back through WebSocket
+            await websocket_manager.send_message(response)
     except WebSocketDisconnect:
-        app.state.assistant.websocket_manager.disconnect(client_id)
-
-if __name__ == "__main__":
-    uvicorn.run(
-        app,
-        host="0.0.0.0",
-        port=7860,
-        reload=False,
-        workers=1,
-        log_level="info"
-    )
+        # Handle WebSocket disconnection
+        await websocket_manager.disconnect()
+
+
+# HTML Test UI
+@app.get("/")
+async def get_ui(request: Request):
+    """Displays the WebSocket HTML UI"""
+    return templates.TemplateResponse("index.html", {"request": request})
+
+
+# Simple UI (fallback in case templates folder is not available)
+@app.get("/simple-ui")
+async def simple_ui():
+    """Fallback HTML for WebSocket Test"""
+    return HTMLResponse(content="""
+    <!DOCTYPE html>
+    <html>
+    <head>
+        <title>WebSocket Test</title>
+        <script>
+            let ws = new WebSocket("ws://localhost:8000/ws");
+
+            ws.onopen = () => {
+                console.log("WebSocket connection opened.");
+            };
+
+            ws.onmessage = (event) => {
+                console.log("Message from server:", event.data);
+                const msgContainer = document.getElementById("messages");
+                const msg = document.createElement("div");
+                msg.innerText = event.data;
+                msgContainer.appendChild(msg);
+            };
+
+            ws.onclose = () => {
+                console.log("WebSocket connection closed.");
+            };
+
+            function sendMessage() {
+                const input = document.getElementById("messageInput");
+                const message = input.value;
+                ws.send(message);
+                input.value = "";
+            }
+        </script>
+    </head>
+    <body>
+        <h1>WebSocket Test</h1>
+        <div id="messages" style="border: 1px solid black; height: 200px; overflow-y: scroll;"></div>
+        <input id="messageInput" type="text" />
+        <button onclick="sendMessage()">Send</button>
+    </body>
+    </html>
+    """)
+