Spaces:
Sleeping
Sleeping
| import os | |
| import json | |
| from typing import List, Dict, Any | |
| 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 | |
| load_dotenv() | |
| app = FastAPI(title="Movie Linker API (Production)") | |
| # --- Configuration --- | |
| SUPABASE_URL = os.getenv("SUPABASE_URL") | |
| SUPABASE_KEY = os.getenv("SUPABASE_ANON_KEY") | |
| 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()] | |
| if not SUPABASE_URL or not SUPABASE_KEY: | |
| raise ValueError("Missing Supabase credentials in environment variables.") | |
| supabase: Client = create_client(SUPABASE_URL, SUPABASE_KEY) | |
| 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.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.") | |
| 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: | |
| raise Exception("All API keys for all providers exhausted.") | |
| key_manager = APIKeyManager(GROQ_API_KEYS, GEMINI_API_KEYS) | |
| # --- Core Functions --- | |
| def ai_match(client, provider, prompt): | |
| if provider == "groq": | |
| res = client.chat.completions.create( | |
| messages=[{"role": "user", "content": prompt}], | |
| model="llama-3.3-70b-versatile", | |
| response_format={"type": "json_object"} | |
| ) | |
| return json.loads(res.choices[0].message.content) | |
| 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 --- | |
| def home(): | |
| return {"status": "online", "engine": "Groq+Gemini"} | |
| def status(): | |
| linked = get_db_data("linked_results", "short_id") | |
| return { | |
| "linked_count": len(linked), | |
| "provider": key_manager.provider, | |
| "active_keys": {p: len(k) for p, k in key_manager.keys.items()} | |
| } | |
| 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")} | |
| # 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": len(all_matches), "total": len(linked_ids) + len(all_matches)} | |
| if __name__ == "__main__": | |
| uvicorn.run(app, host="0.0.0.0", port=7860) | |