Update main.py

main.py

@@ -1,76 +1,111 @@
 import os
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
+import asyncio
+from fastapi import FastAPI
+from telegram import Update
+from telegram.ext import Application, CommandHandler, MessageHandler, filters, ContextTypes
 from huggingface_hub import hf_hub_download
 from llama_cpp import Llama
 
-app = FastAPI()
-
-# --- CONFIGURATION ---
-# REPO_ID = "CoolShotSystems/Axiom-3.1-Sovereign"
-# FILENAME = "Meta-Llama-3.1-8B.Q4_K_M.gguf"
-
 # --- CONFIGURATION ---
-# OLD (Slow): "CoolShotSystems/Axiom-3.1-Sovereign"
-# NEW (Fast): "hugging-quants/Llama-3.2-3B-Instruct-Q4_K_M-GGUF"
-
+# We use the FAST model (Llama 3.2 3B) for Telegram speed
 REPO_ID = "hugging-quants/Llama-3.2-3B-Instruct-Q4_K_M-GGUF"
 FILENAME = "llama-3.2-3b-instruct-q4_k_m.gguf"
 
-# Global Model Variable
+# Global Variables
+app = FastAPI()
 axiom_model = None
+bot_app = None
 
-@app.on_event("startup")
-def load_model():
+# --- AXIOM BRAIN SETUP ---
+def load_brain():
     global axiom_model
-    print("📡 DOWNLOADING AXIOM...")
+    print("📡 DOWNLOADING AXIOM BRAIN...")
     try:
         model_path = hf_hub_download(
             repo_id=REPO_ID,
             filename=FILENAME,
-            token=os.environ.get("HF_TOKEN") 
+            token=os.environ.get("HF_TOKEN")
         )
-        print("🧠 LOADING (LIGHT MODE)...")
-        # OPTIMIZATION: n_ctx=512 makes it MUCH faster on Free Tier
+        print("🧠 LOADING AXIOM INTO RAM...")
         axiom_model = Llama(
             model_path=model_path,
-            n_ctx=512,  # Reduced from 2048 for speed
-            n_threads=2, # Optimal for Hugging Face Free Tier
+            n_ctx=2048, # 2k Context is fine for 3B model
+            n_threads=2,
             verbose=False
         )
-        print("✅ AXIOM ONLINE (FAST MODE)")
+        print("✅ AXIOM ONLINE")
     except Exception as e:
-        print(f"❌ ERROR: {e}")
+        print(f"❌ MODEL FAILURE: {e}")
 
-class ChatRequest(BaseModel):
-    messages: list
+# --- TELEGRAM BOT LOGIC ---
+async def start_command(update: Update, context: ContextTypes.DEFAULT_TYPE):
+    """The /start command"""
+    await update.message.reply_text("Axiom 3.1 Sovereign Interface Online.\nCool Shot Systems Proprietary.\n\nSend me a message.")
 
-@app.get("/")
-def home():
-    return {"status": "Axiom Online"}
-
-@app.post("/v1/chat/completions")
-async def chat(request: ChatRequest):
+async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE):
+    """Handles normal text messages"""
+    user_text = update.message.text
+    
     if not axiom_model:
-        raise HTTPException(status_code=503, detail="Model loading...")
+        await update.message.reply_text("⚠️ Axiom is still waking up... please wait 30 seconds.")
+        return
+
+    # Notify user we are thinking (Typing status)
+    await context.bot.send_chat_action(chat_id=update.effective_chat.id, action="typing")
+
+    # 1. Format Prompt with System Identity
+    prompt = f"""<|start_header_id|>system<|end_header_id|>
+
+You are Axiom 3.1, the Sovereign AI of Cool Shot Systems, created by Professor Heritage.
+You are helpful, strategic, and concise.<|eot_id|><|start_header_id|>user<|end_header_id|>
+
+{user_text}<|eot_id|><|start_header_id|>assistant<|end_header_id|>
+
+"""
+    # 2. Generate
+    # We run this in a thread so we don't block the bot
+    output = await asyncio.to_thread(
+        axiom_model,
+        prompt,
+        max_tokens=256,
+        stop=["<|eot_id|>", "<|end_of_text|>"],
+        echo=False
+    )
     
-    # --- PROMPT FORMATTING (Fixed) ---
-    # We removed <|begin_of_text|> to fix the warning
-    # We inject the Identity immediately
-    prompt = "<|start_header_id|>system<|end_header_id|>\n\nYou are Axiom 3.1, a Sovereign AI created by Professor Heritage at Cool Shot Systems.<|eot_id|>"
+    response = output['choices'][0]['text']
     
-    for msg in request.messages:
-        role = msg['role']
-        content = msg['content']
-        prompt += f"<|start_header_id|>{role}<|end_header_id|>\n\n{content}<|eot_id|>"
-    prompt += "<|start_header_id|>assistant<|end_header_id|>\n\n"
+    # 3. Reply
+    await update.message.reply_text(response)
+
+# --- STARTUP SEQUENCE ---
+@app.on_event("startup")
+async def startup_event():
+    # 1. Load Brain
+    load_brain()
     
-    # GENERATION SETTINGS
-    output = axiom_model(
-        prompt, 
-        max_tokens=128, # Limit output length to save time
-        stop=["<|eot_id|>", "<|end_of_text|>"], 
-        echo=False
-    )
+    # 2. Start Telegram Bot (Background Task)
+    token = os.environ.get("TELEGRAM_TOKEN")
+    if not token:
+        print("❌ NO TELEGRAM TOKEN FOUND!")
+        return
+
+    print("🤖 STARTING TELEGRAM BOT...")
+    global bot_app
+    bot_app = Application.builder().token(token).build()
+    
+    # Add Handlers
+    bot_app.add_handler(CommandHandler("start", start_command))
+    bot_app.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, handle_message))
+    
+    # Initialize and Start Polling
+    await bot_app.initialize()
+    await bot_app.start()
     
-    return {"choices": [{"message": {"role": "assistant", "content": output['choices'][0]['text']}}]}
+    # Run polling in a separate non-blocking task
+    asyncio.create_task(bot_app.updater.start_polling())
+    print("✅ TELEGRAM BOT LISTENING")
+
+# --- DUMMY SERVER (Keeps Space Alive) ---
+@app.get("/")
+def home():
+    return {"status": "Axiom Telegram Bot Running"}