Update app.py

app.py

@@ -1,267 +1,138 @@
 import os
 import json
-import re
 from typing import List, Dict, Any
-from fastapi import FastAPI, HTTPException, BackgroundTasks
-from pydantic import BaseModel
+from fastapi import FastAPI, HTTPException
 from dotenv import load_dotenv
 from supabase import create_client, Client
 from groq import Groq
 from google import genai
 import uvicorn
-import csv
 
-# Load environment variables
 load_dotenv()
 
-app = FastAPI(title="Supabase Data Linker API (Structured Edition)")
+app = FastAPI(title="Movie Linker API (Production)")
 
-# Configuration
+# --- Configuration ---
 SUPABASE_URL = os.getenv("SUPABASE_URL")
 SUPABASE_KEY = os.getenv("SUPABASE_ANON_KEY")
-GROQ_API_KEYS = os.getenv("GROQ_API_KEYS", "").split(",")
-GEMINI_API_KEYS = os.getenv("GEMINI_API_KEYS", "").split(",")
+GROQ_API_KEYS = [k.strip() for k in os.getenv("GROQ_API_KEYS", "").split(",") if k.strip()]
+GEMINI_API_KEYS = [k.strip() for k in os.getenv("GEMINI_API_KEYS", "").split(",") if k.strip()]
 
-# Initialize Supabase
 if not SUPABASE_URL or not SUPABASE_KEY:
-    raise ValueError("Missing SUPABASE environment variables")
+    raise ValueError("Missing Supabase credentials in environment variables.")
+
 supabase: Client = create_client(SUPABASE_URL, SUPABASE_KEY)
 
-class MultiProviderManager:
-    """Manages rotation for both Groq and Gemini keys."""
-    def __init__(self, groq_keys: List[str], gemini_keys: List[str]):
-        self.providers = {
-            "groq": [k.strip() for k in groq_keys if k.strip()],
-            "gemini": [k.strip() for k in gemini_keys if k.strip()]
-        }
+class APIKeyManager:
+    """Manages rotation between Groq and Gemini providers and their keys."""
+    def __init__(self, groq_keys, gemini_keys):
+        self.keys = {"groq": groq_keys, "gemini": gemini_keys}
         self.indices = {"groq": 0, "gemini": 0}
-        self.current_provider = "groq" if self.providers["groq"] else "gemini"
-
-    def get_session(self):
-        if not self.providers[self.current_provider]:
-            if self.current_provider == "groq" and self.providers["gemini"]:
-                self.current_provider = "gemini"
-                print("🔀 Groq not available, switching to Gemini...")
-                return self.get_session()
-            raise ValueError("No API keys found for the current provider or fallback.")
-
-        key = self.providers[self.current_provider][self.indices[self.current_provider]]
+        self.provider = "groq" if groq_keys else "gemini"
+
+    def get_client(self):
+        if not self.keys[self.provider]:
+            if self.provider == "groq" and self.keys["gemini"]:
+                self.provider = "gemini"
+                return self.get_client()
+            raise Exception("No API keys available.")
         
-        if self.current_provider == "groq":
-            return Groq(api_key=key), "groq"
-        else:
-            return genai.Client(api_key=key), "gemini"
-
-    def rotate_current(self):
-        """Try next key in current provider, or switch provider."""
-        current_keys = self.providers[self.current_provider]
-        if self.indices[self.current_provider] + 1 < len(current_keys):
-            self.indices[self.current_provider] += 1
-            print(f"🔄 Rotating to {self.current_provider.upper()} Key #{self.indices[self.current_provider] + 1}")
+        current_key = self.keys[self.provider][self.indices[self.provider]]
+        if self.provider == "groq":
+            return Groq(api_key=current_key), "groq"
+        return genai.Client(api_key=current_key), "gemini"
+
+    def rotate(self):
+        if self.indices[self.provider] + 1 < len(self.keys[self.provider]):
+            self.indices[self.provider] += 1
+            print(f"🔄 Switched to {self.provider} key #{self.indices[self.provider]+1}")
+        elif self.provider == "groq" and self.keys["gemini"]:
+            self.provider = "gemini"
+            print("⚠️ Groq exhausted, falling back to Gemini.")
         else:
-            if self.current_provider == "groq" and self.providers["gemini"]:
-                self.current_provider = "gemini"
-                self.indices["gemini"] = 0
-                print("⚠️ Groq Keys exhausted. Switching to Gemini provider.")
-            else:
-                raise Exception(f"All keys for {self.current_provider} exhausted.")
-
-api_manager = MultiProviderManager(GROQ_API_KEYS, GEMINI_API_KEYS)
-
-# --- Helper Functions ---
+            raise Exception("All API keys for all providers exhausted.")
 
-def fetch_supabase_data(table_name: str):
-    try:
-        response = supabase.table(table_name).select("*").execute()
-        return response.data
-    except Exception as e:
-        print(f"Error fetching {table_name}: {e}")
-        return []
+key_manager = APIKeyManager(GROQ_API_KEYS, GEMINI_API_KEYS)
 
-def load_existing_links_from_db():
-    try:
-        response = supabase.table("linked_results").select("short_id").execute()
-        return response.data
-    except Exception as e:
-        print(f"Error loading links from DB: {e}")
-        return []
+# --- Core Functions ---
 
-def save_links_to_db(links: List[Dict[str, Any]]):
-    if not links: return 0
-    try:
-        data = [{
-            "short_id": l["short_id"],
-            "media_id": l["media_id"],
-            "short_title": l["short_title"],
-            "media_title": l["media_title"],
-            "short_link": l["short_link"],
-            "media_link": l["media_link"]
-        } for l in links]
-        response = supabase.table("linked_results").insert(data).execute()
-        return len(response.data)
-    except Exception as e:
-        print(f"❌ DB Save Error: {e}")
-        return 0
-
-# --- Linking Logic ---
-
-def call_ai(session, provider, prompt):
+def ai_match(client, provider, prompt):
     if provider == "groq":
-        completion = session.chat.completions.create(
+        res = client.chat.completions.create(
             messages=[{"role": "user", "content": prompt}],
             model="llama-3.3-70b-versatile",
             response_format={"type": "json_object"}
         )
-        return json.loads(completion.choices[0].message.content)
-    else:
-        response = session.models.generate_content(
-            model="gemini-3-flash-preview",
-            contents=prompt,
-            config={'response_mime_type': 'application/json'}
-        )
-        return json.loads(response.text)
-
-def perform_smart_linking(shorts, movies, batch_size=20):
-    all_matches = []
+        return json.loads(res.choices[0].message.content)
     
-    # Process media for AI
-    movies_ref = [
-        {
-            "id": m["id"], 
-            "title": m["title"], 
-            "type": m.get("type", "unknown"),
-            "year": str(m.get("releasdate", ""))[:4],
-            "description": m.get("dec", "")[:150],
-            "url": m.get("url") or f"https://flixhq.to/{m.get('id_slug','')}"
-        } for m in movies
-    ]
-
-    # Process shorts for AI (Optimized for structured data)
-    shorts_list = []
-    for s in shorts:
-        shorts_list.append({
-            "id": s.get("id"),
-            "title": s.get("title", s.get("yt_title", "Unknown")),
-            "description": s.get("description", "")[:100], # Pass some description for context
-            "url": s.get("url") or s.get("yt_link", "N/A")
-        })
-
-    for i in range(0, len(shorts_list), batch_size):
-        batch = shorts_list[i : i + batch_size]
-        success = False
-        
-        while not success:
-            try:
-                session, provider = api_manager.get_session()
-                prompt = f"""
-                Match these 'shorts' to their corresponding 'movies/tv series' (media).
-                Reference Media: {json.dumps(movies_ref[:120])}
-                Target Shorts: {json.dumps(batch)}
-                Respond ONLY with JSON: {{"matches": [{{"short_id", "media_id", "short_title", "media_title", "short_link", "media_link"}}]}}
-                """
-                
-                result = call_ai(session, provider, prompt)
-                if "matches" in result:
-                    all_matches.extend(result["matches"])
-                success = True
-            except Exception as e:
-                err = str(e).lower()
-                if any(x in err for x in ["rate_limit", "429", "413"]):
-                    try:
-                        api_manager.rotate_current()
-                    except:
-                        print("🛑 All Keys/Providers exhausted for this batch.")
-                        break
-                else:
-                    print(f"❌ Error in batch {i}: {e}")
-                    break
-    return all_matches
+    res = client.models.generate_content(
+        model="gemini-3-flash-preview",
+        contents=prompt,
+        config={'response_mime_type': 'application/json'}
+    )
+    return json.loads(res.text)
+
+def get_db_data(table: str, select: str = "*"):
+    return supabase.table(table).select(select).execute().data
+
+def process_batch(shorts_batch, movies_ref):
+    while True:
+        try:
+            client, provider = key_manager.get_client()
+            prompt = f"Match these shorts to media. Media: {json.dumps(movies_ref[:120])}. Shorts: {json.dumps(shorts_batch)}. Respond with JSON: {{'matches': [{{'short_id', 'media_id', 'short_title', 'media_title', 'short_link', 'media_link'}}]}}"
+            return ai_match(client, provider, prompt).get("matches", [])
+        except Exception as e:
+            if any(x in str(e).lower() for x in ["rate_limit", "429", "limit_reached"]):
+                key_manager.rotate()
+            else:
+                print(f"❌ Batch Error: {e}")
+                return []
 
 # --- Endpoints ---
 
 @app.get("/")
 def home():
-    return {"status": "online", "message": "Supabase Linked API (Structured Shorts Edition)"}
+    return {"status": "online", "engine": "Groq+Gemini"}
 
 @app.get("/status")
 def status():
-    existing = load_existing_links_from_db()
+    linked = get_db_data("linked_results", "short_id")
     return {
-        "db_linked_count": len(existing),
-        "groq_keys": len(api_manager.providers["groq"]),
-        "gemini_keys": len(api_manager.providers["gemini"]),
-        "current_provider": api_manager.current_provider
+        "linked_count": len(linked),
+        "provider": key_manager.provider,
+        "active_keys": {p: len(k) for p, k in key_manager.keys.items()}
     }
 
-@app.get("/test-db")
-def test_db():
-    try:
-        test_data = {
-            "short_id": 999999,
-            "media_id": 1,
-            "short_title": "Test DB Access",
-            "media_title": "Test Movie",
-            "short_link": "http://test.com",
-            "media_link": "http://test.com"
-        }
-        supabase.table("linked_results").insert(test_data).execute()
-        supabase.table("linked_results").delete().eq("short_id", 999999).execute()
-        return {"status": "success", "message": "Database is working!"}
-    except Exception as e:
-        return {"status": "failed", "error": str(e)}
-
-@app.get("/test-integration")
-def test_integration():
-    if not api_manager.providers["gemini"]:
-        return {"error": "No Gemini Keys configured"}
-    try:
-        client = genai.Client(api_key=api_manager.providers["gemini"][0])
-        prompt = """Generate one realistic mock linked result JSON."""
-        response = client.models.generate_content(
-            model="gemini-3-flash-preview",
-            contents=prompt,
-            config={'response_mime_type': 'application/json'}
-        )
-        mock_data = json.loads(response.text)
-        # Handle cases where AI returns a list or direct object
-        if isinstance(mock_data, list): mock_data = mock_data[0]
-        
-        # Ensure keys match DB
-        test_entry = {
-            "short_id": mock_data.get("short_id", 888888),
-            "media_id": mock_data.get("media_id", 12345),
-            "short_title": mock_data.get("short_title", "Mock Short"),
-            "media_title": mock_data.get("media_title", "Mock Media"),
-            "short_link": mock_data.get("short_link", "http://yt.com"),
-            "media_link": mock_data.get("media_link", "http://flix.com")
-        }
-        res = supabase.table("linked_results").insert(test_entry).execute()
-        return {"status": "success", "data_generated": test_entry}
-    except Exception as e:
-        return {"status": "failed", "error": str(e)}
-
 @app.post("/link-data")
-def trigger():
-    shorts = fetch_supabase_data("shorts")
-    media = fetch_supabase_data("media")
-    
-    if not shorts or not media:
-        raise HTTPException(status_code=500, detail="Data missing")
-        
-    existing_ids = {str(l["short_id"]) for l in load_existing_links_from_db()}
-    new_shorts = [s for s in shorts if str(s.get("id")) not in existing_ids]
+def link_data():
+    shorts = get_db_data("shorts")
+    media = get_db_data("media")
+    linked_ids = {str(l["short_id"]) for l in get_db_data("linked_results", "short_id")}
     
-    if not new_shorts:
-        return {"message": "All caught up!"}
-        
-    print(f"🚀 Processing {len(new_shorts)} new structured shorts...")
-    matches = perform_smart_linking(new_shorts, media)
-    saved = save_links_to_db(matches)
+    # Filter only new shorts
+    new_shorts = [s for s in shorts if str(s.get("id")) not in linked_ids]
+    if not new_shorts: return {"message": "Everything is already linked."}
+
+    # Prepare data for AI
+    movies_ref = [{
+        "id": m["id"], "title": m["title"], "type": m.get("type"),
+        "year": str(m.get("releasdate", ""))[:4], "desc": m.get("dec", "")[:100]
+    } for m in media]
+
+    batch_shorts = [{
+        "id": s["id"], "title": s.get("title"), "desc": s.get("description", "")[:100]
+    } for s in new_shorts]
+
+    # Matching Loop
+    all_matches = []
+    for i in range(0, len(batch_shorts), 20):
+        all_matches.extend(process_batch(batch_shorts[i : i + 20], movies_ref))
+
+    # Save Results
+    if all_matches:
+        supabase.table("linked_results").insert(all_matches).execute()
     
-    return {
-        "newly_linked": saved,
-        "total_in_db": len(load_existing_links_from_db())
-    }
+    return {"newly_linked": len(all_matches), "total": len(linked_ids) + len(all_matches)}
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)