Update main.py

main.py

@@ -1,138 +1,70 @@
 import os
-import json
-from fastapi import FastAPI, Request
-from fastapi.staticfiles import StaticFiles
+import httpx
+from fastapi import FastAPI, Request, BackgroundTasks
 import uvicorn
-from playwright.async_api import async_playwright
 from litellm import acompletion
 
 # --- CONFIG ---
 os.environ["GEMINI_API_KEY"] = os.getenv("GEMINI_API_KEY", "")
-SPACE_HOST = os.getenv("MY_SPACE_HOST", "")
-STATIC_DIR = "static"
+TELEGRAM_TOKEN = os.getenv("TELEGRAM_TOKEN", "")
+TELEGRAM_API_URL = f"https://api.telegram.org/bot{TELEGRAM_TOKEN}"
 
 app = FastAPI()
 
-if not os.path.exists(STATIC_DIR):
-    os.makedirs(STATIC_DIR)
-
-app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
-
-# --- TOOLS ---
-async def take_screenshot(url: str) -> str:
-    """Navigates to a URL and takes a screenshot. Returns the URL to view it."""
-    filepath = f"{STATIC_DIR}/screenshot.png"
-    async with async_playwright() as p:
-        browser = await p.chromium.launch(headless=True)
-        page = await browser.new_page()
-        await page.goto(url)
-        await page.screenshot(path=filepath, full_page=True)
-        await browser.close()
-    
-    return f"Screenshot successful. View it here: https://{SPACE_HOST}/static/screenshot.png"
-
-tools = [
-    {
-        "type": "function",
-        "function": {
-            "name": "take_screenshot",
-            "description": "Takes a screenshot of a URL.",
-            "parameters": {
-                "type": "object",
-                "properties": {
-                    "url": {"type": "string", "description": "The URL to screenshot"}
-                },
-                "required": ["url"]
-            }
-        }
-    }
-]
-
 system_prompt = (
-    "You are an expert coding/devops agent. Users interact via a WhatsApp bridge. "
-    "Keep replies concise and under 1000 chars. Explain before acting. "
-    "If you take a screenshot, provide the URL clearly."
+    "You are an expert coding/devops agent. "
+    "Keep replies concise and under 1000 chars. Use Markdown formatting."
 )
 
 # --- NATIVE AGENT ENGINE ---
 async def run_agent(prompt: str) -> str:
-    messages = [
-        {"role": "system", "content": system_prompt},
-        {"role": "user", "content": prompt}
-    ]
-    
     try:
         response = await acompletion(
             model="gemini/gemini-2.5-pro",
-            messages=messages,
-            tools=tools,
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": prompt}
+            ],
             temperature=0.2
         )
-        
-        msg = response.choices[0].message
-        
-        # Check if the LLM wants to use the screenshot tool
-        if hasattr(msg, 'tool_calls') and msg.tool_calls:
-            assistant_msg = {"role": "assistant", "content": msg.content or ""}
-            assistant_msg["tool_calls"] = [
-                {
-                    "id": tc.id,
-                    "type": "function",
-                    "function": {
-                        "name": tc.function.name,
-                        "arguments": tc.function.arguments
-                    }
-                } for tc in msg.tool_calls
-            ]
-            messages.append(assistant_msg)
-            
-            # Execute the tool
-            for tool_call in msg.tool_calls:
-                if tool_call.function.name == "take_screenshot":
-                    args = json.loads(tool_call.function.arguments)
-                    result = await take_screenshot(args["url"])
-                    
-                    messages.append({
-                        "role": "tool",
-                        "tool_call_id": tool_call.id,
-                        "name": "take_screenshot",
-                        "content": result
-                    })
-            
-            # Get the final answer after the screenshot
-            final_response = await acompletion(
-                model="gemini/gemini-2.5-pro",
-                messages=messages,
-                temperature=0.2
-            )
-            return final_response.choices[0].message.content
-            
-        return msg.content or "No response generated."
+        return response.choices[0].message.content or "No response generated."
     except Exception as e:
         return f"Agent Error: {str(e)}"
 
+# --- BACKGROUND TASK: PROCESS AND REPLY ---
+async def process_and_reply(chat_id: int, text: str):
+    response_text = await run_agent(text)
+    
+    # Send the final response back to Telegram
+    async with httpx.AsyncClient() as client:
+        await client.post(
+            f"{TELEGRAM_API_URL}/sendMessage",
+            json={
+                "chat_id": chat_id, 
+                "text": response_text, 
+                "parse_mode": "Markdown"
+            }
+        )
+
 # --- ENDPOINTS ---
 @app.get("/health")
 async def health_check():
     return {"status": "alive"}
 
-@app.api_route("/webhook", methods=["GET", "POST"])
-async def wa_webhook(request: Request):
-    user_msg = request.query_params.get("message", "")
+@app.post("/webhook")
+async def telegram_webhook(request: Request, background_tasks: BackgroundTasks):
+    data = await request.json()
     
-    if request.method == "POST":
-        try:
-            data = await request.json()
-            user_msg = data.get("message", user_msg)
-        except:
-            pass
-
-    if user_msg.lower().startswith("claw "):
-        prompt = user_msg[5:].strip()
-        response_text = await run_agent(prompt)
-        return {"replies": [{"message": response_text}]}
-    
-    return {"replies": []}
+    # Extract the message from Telegram's JSON payload
+    if "message" in data and "text" in data["message"]:
+        chat_id = data["message"]["chat"]["id"]
+        user_msg = data["message"]["text"]
+        
+        # Run the AI generation in the background so we can instantly reply "OK" to Telegram
+        background_tasks.add_task(process_and_reply, chat_id, user_msg)
+        
+    # Telegram requires an immediate 200 OK so it doesn't get stuck in a retry loop
+    return {"status": "ok"}
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)