Update app.py

app.py

@@ -1,127 +1,19 @@
-# ================================
-# AVIA AI BACKEND (Single File)
-# ================================
-
 import os
-import re
 import uuid
-import logging
-import asyncio
-from typing import List, Optional, Dict
-from urllib.parse import quote
-
+import requests
 import uvicorn
-from fastapi import FastAPI, HTTPException
+import edge_tts
+from fastapi import FastAPI, Request
 from fastapi.responses import FileResponse
 from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
-
-# HF
-from huggingface_hub import InferenceClient
-
-# Search
 from duckduckgo_search import DDGS
+import google.generativeai as genai  # 🟢 NEW: Gemini Import
 
-# TTS
-import edge_tts
-
-
-# ================================
-# CONFIG
-# ================================
-
-HF_TOKEN = os.getenv("HF_TOKEN")
-
-MODEL_LIST = [
-    "microsoft/Phi-3.5-mini-instruct",
-    "Qwen/Qwen2-1.5B-Instruct",
-]
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger("avia")
-
-ddgs = DDGS()
-
-SYSTEM_PROMPT = """
-You are Avia AI assistant created by Tanveer Ali.
-Speak in Hinglish.
-If user wants image → output only:
-[IMAGE_PROMPT: detailed english prompt]
-"""
-
-
-# ================================
-# UTIL
-# ================================
-
-def contains_hindi(text: str) -> bool:
-    return bool(re.search(r'[\u0900-\u097F]', text))
-
-
-async def search_internet(query: str) -> str:
-    try:
-        results = await asyncio.to_thread(lambda: list(ddgs.text(query, max_results=3)))
-        if not results:
-            return ""
-        return "\n".join([r.get("body", "") for r in results])
-    except Exception as e:
-        logger.error(e)
-        return ""
-
-
-def extract_image_prompt(text: str):
-    m = re.search(r"\[IMAGE_PROMPT:\s*(.*?)\]", text, re.I)
-    return m.group(1) if m else None
-
-
-def generate_image_url(prompt: str):
-    return f"https://image.pollinations.ai/prompt/{quote(prompt)}?seed={uuid.uuid4().int%1000}"
-
-
-def user_wants_image(msg: str):
-    words = ["image", "photo", "picture", "tasveer", "banao", "draw"]
-    return any(w in msg.lower() for w in words)
-
-
-async def call_llm(messages):
-
-    for model in MODEL_LIST:
-        try:
-            logger.info(f"Trying model {model}")
-            client = InferenceClient(model=model, token=HF_TOKEN)
-
-            resp = await asyncio.to_thread(
-                lambda: client.chat_completion(
-                    messages=messages,
-                    max_tokens=400,
-                    temperature=0.7,
-                )
-            )
-
-            return resp.choices[0].message.content
-
-        except Exception as e:
-            logger.error(f"{model} failed {e}")
-
-    return None
-
-
-# ================================
-# SCHEMA
-# ================================
-
-class ChatRequest(BaseModel):
-    message: str
-    history: Optional[List[Dict[str, str]]] = None
-
-
-class ChatResponse(BaseModel):
-    reply: str
-
-
-# ================================
-# APP
-# ================================
+# ✅ 1. SETUP & CONFIGURATION
+# Ab HF_TOKEN ki jagah hum GEMINI_API_KEY use karenge
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
+if GEMINI_API_KEY:
+    genai.configure(api_key=GEMINI_API_KEY)
 
 app = FastAPI()
 
@@ -132,75 +24,121 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-
-# ================================
-# ROUTES
-# ================================
-
 @app.get("/")
-def root():
-    return {"status": "ok", "avia": True}
-
-
-@app.post("/chat", response_model=ChatResponse)
-async def chat(req: ChatRequest):
-
-    msg = req.message.strip()
-    history = req.history or []
-
-    # search trigger
-    search_words = ["news", "weather", "price", "who", "latest"]
-    search_ctx = ""
-
-    if any(w in msg.lower() for w in search_words):
-        search_ctx = await search_internet(msg)
-
-    system = SYSTEM_PROMPT
-    if search_ctx:
-        system += "\n" + search_ctx
+async def home():
+    return {"status": "Avia AI (Gemini Powered Stable Version) is Ready! 🚀"}
 
-    messages = [{"role": "system", "content": system}]
-    messages.extend(history[-4:])
-    messages.append({"role": "user", "content": msg})
-
-    # call model
-    reply = await call_llm(messages)
-
-    if not reply:
-        return ChatResponse(reply="Server busy. Try again.")
-
-    # image logic
-    img_prompt = extract_image_prompt(reply)
-
-    if img_prompt:
-        return ChatResponse(reply="IMAGE_URL:" + generate_image_url(img_prompt))
-
-    if user_wants_image(msg):
-        return ChatResponse(reply="IMAGE_URL:" + generate_image_url(msg))
+# 🌍 2. INTERNET SEARCH (Real-time Data)
+def search_web(query):
+    try:
+        with DDGS() as ddgs:
+            results = [r for r in ddgs.text(query, region='in-en', max_results=3)]
+            if results:
+                summary = "\n".join([f"• {r['title']}: {r['body']}" for r in results])
+                return f"\n[INTERNET DATA - Use this if relevant]:\n{summary}\n"
+    except: pass
+    return ""
+
+# 🧠 3. SYSTEM PROMPT (The Persona)
+def get_avia_prompt(search_context):
+    return f"""
+    You are Avia, a smart and friendly AI Assistant created by Tanveer Ali.
+    
+    YOUR RULES:
+    1. Language: Answer in a natural mix of Hindi and English (Hinglish).
+    2. Personality: Helpful, intelligent, and polite.
+    3. Creator: Tanveer Ali.
+    
+    CAPABILITIES:
+    1. IMAGE GENERATION: 
+       - If the user asks to "generate", "create", "draw", "make" an image or photo:
+       - You MUST output a special tag: [IMAGE_PROMPT: <detailed English description>].
+       - Example: User: "Ek futuristic car banao" -> Response: "Sure! [IMAGE_PROMPT: A futuristic sports car, neon lights, cyberpunk city background, 8k resolution]"
+    
+    2. INTERNET: 
+       - You have access to real-time information. Use the data below if needed.
+    
+    {search_context}
+    """
+
+# 📨 4. CHAT API (The Main Logic)
+@app.post("/chat")
+async def chat(request: Request):
+    try:
+        data = await request.json()
+        user_msg = data.get("message", "")
+        history = data.get("history", [])
+        
+        # --- Step A: Internet Search Check ---
+        search_context = ""
+        triggers = ["news", "price", "who is", "weather", "today", "search", "score", "match", "latest"]
+        if "?" in user_msg or any(w in user_msg.lower() for w in triggers):
+            web_data = search_web(user_msg)
+            if web_data: search_context = web_data
+
+        # --- Step B: Initialize Gemini Brain with Persona ---
+        system_instruction = get_avia_prompt(search_context)
+        
+        # 🟢 NEW: Setup Gemini Model
+        model = genai.GenerativeModel(
+            model_name="gemini-1.5-flash",
+            system_instruction=system_instruction
+        )
+        
+        # 🟢 NEW: Convert history to Gemini format (Last 4 messages)
+        gemini_history = []
+        for h in history[-4:]:
+            role = "model" if h['role'] in ["assistant", "model"] else "user"
+            content = h.get('content', '')
+            if content:
+                gemini_history.append({"role": role, "parts": [content]})
+        
+        # --- Step C: Call Internal Brain (Gemini API) ---
+        try:
+            chat_session = model.start_chat(history=gemini_history)
+            response = chat_session.send_message(user_msg)
+            ai_reply = response.text
+            
+            # --- Step D: Image Generation Logic (Pollinations) ---
+            
+            # 1. Check if AI gave the tag
+            if "[IMAGE_PROMPT" in ai_reply:
+                try:
+                    start = ai_reply.find("[IMAGE_PROMPT:") + 14
+                    end = ai_reply.find("]", start)
+                    if end != -1:
+                        prompt = ai_reply[start:end].strip()
+                        img_url = f"https://image.pollinations.ai/prompt/{requests.utils.quote(prompt)}?nologo=true&seed={uuid.uuid4().int % 1000}"
+                        return {"reply": f"IMAGE_URL:{img_url}"}
+                except: pass
+
+            # 2. Backup Logic
+            keywords = ["generate", "draw", "create", "banao", "tasveer", "photo", "image"]
+            if any(w in user_msg.lower() for w in keywords) and ("image" in user_msg.lower() or "photo" in user_msg.lower() or "tasveer" in user_msg.lower()):
+                 img_url = f"https://image.pollinations.ai/prompt/{requests.utils.quote(user_msg)}?nologo=true&seed={uuid.uuid4().int % 1000}"
+                 return {"reply": f"IMAGE_URL:{img_url}"}
+
+            return {"reply": ai_reply}
 
-    return ChatResponse(reply=reply)
+        except Exception as e:
+            print(f"Brain Error: {e}")
+            return {"reply": "Sorry Tanveer, mera server abhi thoda busy hai. Please try again in 5 seconds."}
 
+    except Exception as e:
+        return {"reply": f"Error: {str(e)}"}
 
+# 🔊 5. TTS API (Voice)
 @app.get("/tts")
 async def tts(text: str):
-
-    if not text:
-        raise HTTPException(400)
-
-    voice = "hi-IN-SwaraNeural" if contains_hindi(text) else "en-US-AriaNeural"
-
-    filename = f"voice_{uuid.uuid4()}.mp3"
-
     try:
-        await edge_tts.Communicate(text, voice).save(filename)
+        is_hindi = any('\u0900' <= c <= '\u097f' for c in text)
+        voice = "hi-IN-SwaraNeural" if is_hindi else "en-US-AriaNeural"
+        
+        communicate = edge_tts.Communicate(text, voice)
+        filename = f"voice_{uuid.uuid4()}.mp3"
+        await communicate.save(filename)
         return FileResponse(filename, media_type="audio/mpeg")
-    except Exception:
-        raise HTTPException(500)
-
-
-# ================================
-# RUN
-# ================================
+    except: return {"error": "TTS Error"}
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)