Update app.py

app.py

@@ -2,17 +2,15 @@ import os
 import re
 import logging
 import uuid
-import time
-from datetime import datetime, timezone, timedelta # Added timedelta
-from typing import Optional, Dict, Any, List # Added List
-from collections import defaultdict, deque # Added defaultdict and deque
+from datetime import datetime, timezone
+from typing import Optional, Dict, Any
 
-from fastapi import FastAPI, HTTPException, Body, BackgroundTasks, Path, Request # Added Request
+from fastapi import FastAPI, HTTPException, Body, BackgroundTasks, Path
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel, Field
 
-import openai
-import google.generativeai as genai
+import openai # For your custom API
+import google.generativeai as genai # For Gemini API
 from google.generativeai.types import GenerationConfig
 
 # --- Logging Configuration ---
@@ -29,108 +27,7 @@ CUSTOM_API_MODEL_DEFAULT = "gemma3:27b"
 DEFAULT_GEMINI_MODEL = "gemini-1.5-flash-latest"
 GEMINI_REQUEST_TIMEOUT_SECONDS = 300
 
-# Rate limiting dictionary
-class RateLimiter:
-    def __init__(self, max_requests: int, time_window: timedelta):
-        self.max_requests = max_requests
-        self.time_window = time_window
-        self.requests: Dict[str, list] = defaultdict(list)
-    
-    def _cleanup_old_requests(self, user_ip: str) -> None:
-        """Remove requests that are outside the time window."""
-        current_time = time.time()
-        self.requests[user_ip] = [
-            timestamp for timestamp in self.requests[user_ip]
-            if current_time - timestamp < self.time_window.total_seconds()
-        ]
-    
-    def is_rate_limited(self, user_ip: str) -> bool:
-        """Check if the user has exceeded their rate limit."""
-        self._cleanup_old_requests(user_ip)
-        
-        # Get current count after cleanup
-        current_count = len(self.requests[user_ip])
-        
-        # Add current request timestamp (incrementing the count)
-        current_time = time.time()
-        self.requests[user_ip].append(current_time)
-        
-        # Check if user has exceeded the maximum requests
-        return (current_count + 1) > self.max_requests
-    
-    def get_current_count(self, user_ip: str) -> int:
-        """Get the current request count for an IP."""
-        self._cleanup_old_requests(user_ip)
-        return len(self.requests[user_ip])
-
-# Initialize rate limiter for the /chat endpoint
-# Max 12 requests per day, per IP.
-rate_limiter = RateLimiter(
-    max_requests=3,
-    time_window=timedelta(days=1)
-)
-
-def get_user_ip(request: Request) -> str:
-    """Helper function to get user's IP address."""
-    forwarded = request.headers.get("X-Forwarded-For")
-    if forwarded:
-        return forwarded.split(",")[0]
-    return request.client.host
-
-
-class ApiRotator:
-    def __init__(self, apis):
-        self.apis = apis
-        self.last_successful_index = None
-
-    def get_prioritized_apis(self):
-        if self.last_successful_index is not None:
-            # Move the last successful API to the front
-            rotated_apis = (
-                [self.apis[self.last_successful_index]] + 
-                self.apis[:self.last_successful_index] + 
-                self.apis[self.last_successful_index+1:]
-            )
-            return rotated_apis
-        return self.apis
-
-    def update_last_successful(self, index):
-        self.last_successful_index = index
-
-
-# --- API Rotator Class (as provided by you) ---
-class ApiRotator:
-    def __init__(self, apis: List[Any]): # Assuming apis is a list of API configurations
-        self.apis = apis
-        self.last_successful_index: Optional[int] = None
-        logger.info(f"ApiRotator initialized with {len(apis)} APIs.")
-
-    def get_prioritized_apis(self) -> List[Any]:
-        if self.last_successful_index is not None and 0 <= self.last_successful_index < len(self.apis):
-            # Move the last successful API to the front
-            rotated_apis = (
-                [self.apis[self.last_successful_index]] +
-                self.apis[:self.last_successful_index] +
-                self.apis[self.last_successful_index + 1:]
-            )
-            # logger.debug("Rotated APIs based on last successful index.")
-            return rotated_apis
-        # logger.debug("Returning APIs in original or non-rotated order.")
-        return list(self.apis) # Return a copy
-
-    def update_last_successful(self, api_config_used: Any):
-        try:
-            index = self.apis.index(api_config_used)
-            self.last_successful_index = index
-            # logger.info(f"Updated last successful API to index: {index}")
-        except ValueError:
-            # This might happen if api_config_used is not in the original list,
-            # or if the list was modified.
-            logger.warning(f"Could not find API config in rotator to update last successful index.")
-            self.last_successful_index = None # Reset if not found
-
-
-# --- In-Memory Task Storage (for Gemini async tasks) ---
+# --- In-Memory Task Storage ---
 tasks_db: Dict[str, Dict[str, Any]] = {}
 
 # --- Pydantic Models ---
@@ -156,23 +53,24 @@ class TaskStatusResponse(BaseModel):
     last_updated_at: datetime
     result: Optional[str] = None
     error: Optional[str] = None
+    # request_params: Optional[Dict[str, Any]] = None # Optionally return original params
+
 
 # --- FastAPI App Initialization ---
 app = FastAPI(
-    title="Dual Chat & Async Gemini API with Rate Limiting",
-    description="Provides rate-limited direct chat and asynchronous tasks for Gemini API.",
-    version="2.1.0"
+    title="Dual Chat & Async Gemini API",
+    description="Provides direct chat via custom API and asynchronous tasks for Gemini API.",
+    version="2.0.0"
 )
 
-# --- Helper Functions (Regex from previous step) ---
+# --- Helper Functions ---
 def is_video_url_for_gemini(url: Optional[str]) -> bool:
     if not url:
         return False
     youtube_regex = (
         r'(https_?://)?(www\.)?'
-        r'(youtube|youtu|youtube-nocookie)\.(com|be)/'
-        r'(watch\?v=|embed/|v/|.+\?v=)?([^&=%\?]{11})'
-    )
+        '(youtube|youtu|youtube-nocookie)\.(com|be)/'
+        '(watch\?v=|embed/|v/|.+\?v=)?([^&=%\?]{11})')
     googleusercontent_youtube_regex = r'https_?://googleusercontent\.com/youtube\.com/\w+'
     return re.match(youtube_regex, url) is not None or \
            re.match(googleusercontent_youtube_regex, url) is not None
@@ -184,37 +82,54 @@ async def process_gemini_request_background(
     requested_gemini_model: str, 
     gemini_key_to_use: str
 ):
-    # (Gemini background processing logic remains the same as before)
     logger.info(f"[Task {task_id}] Starting background Gemini processing. Model: {requested_gemini_model}, URL: {input_url}")
     tasks_db[task_id]["status"] = "PROCESSING"
     tasks_db[task_id]["last_updated_at"] = datetime.now(timezone.utc)
 
     try:
         genai.configure(api_key=gemini_key_to_use)
+        
         model_instance = genai.GenerativeModel(model_name=requested_gemini_model)
+        
         content_parts = [{"text": user_message}]
         if input_url and is_video_url_for_gemini(input_url):
             logger.info(f"[Task {task_id}] Adding video URL to Gemini content: {input_url}")
-            content_parts.append({"file_data": {"mime_type": "video/youtube", "file_uri": input_url}})
+            content_parts.append({
+                "file_data": {
+                    "mime_type": "video/youtube", # Or let Gemini infer
+                    "file_uri": input_url
+                }
+            })
         
         gemini_contents = [{"parts": content_parts}]
+        
         generation_config = GenerationConfig(candidate_count=1)
         request_options = {"timeout": GEMINI_REQUEST_TIMEOUT_SECONDS}
 
         logger.info(f"[Task {task_id}] Sending request to Gemini API...")
         response = await model_instance.generate_content_async(
-            gemini_contents, stream=False, generation_config=generation_config, request_options=request_options
+            gemini_contents,
+            stream=False, # Collect full response for async task
+            generation_config=generation_config,
+            request_options=request_options
         )
         
+        # Assuming response.text contains the full aggregated text
+        # If using a model version that streams even for non-stream call, aggregate it:
         full_response_text = ""
         if hasattr(response, 'text') and response.text:
             full_response_text = response.text
-        elif hasattr(response, 'parts'):
+        elif hasattr(response, 'parts'): # Check for newer API structures if .text is not primary
             for part in response.parts:
-                if hasattr(part, 'text'): full_response_text += part.text
-        else:
+                if hasattr(part, 'text'):
+                    full_response_text += part.text
+        else: # Fallback for safety if structure is unexpected or if it's an iterable of chunks
+            # This part might need adjustment based on actual non-streaming response object
+            # For now, assuming generate_content_async with stream=False gives a response with .text
+            # or we need to iterate if it's still a stream internally for some models
             logger.warning(f"[Task {task_id}] Gemini response structure not as expected or empty. Response: {response}")
 
+
         if not full_response_text and response.prompt_feedback and response.prompt_feedback.block_reason:
             block_reason_name = response.prompt_feedback.block_reason.name if hasattr(response.prompt_feedback.block_reason, 'name') else str(response.prompt_feedback.block_reason)
             logger.warning(f"[Task {task_id}] Gemini content blocked: {block_reason_name}")
@@ -238,21 +153,9 @@ async def process_gemini_request_background(
 
 # --- API Endpoints ---
 
-# Modified /chat endpoint with rate limiting
 @app.post("/chat", response_class=StreamingResponse)
-async def direct_chat(payload: ChatPayload, request: Request): # Added 'request: Request'
-    user_ip = get_user_ip(request)
-    logger.info(f"Direct chat request received from IP: {user_ip}. Temp: {payload.temperature}, Msg: '{payload.message[:50]}...'")
-    
-    if rate_limiter.is_rate_limited(user_ip):
-        current_count = rate_limiter.get_current_count(user_ip)
-        raise HTTPException(
-            status_code=429,
-            detail={
-                "error": "You have exceeded the maximum number of requests per day. Please try again tomorrow.",
-                "url": "https://t.me/chrunoss"
-            }
-        )
+async def direct_chat(payload: ChatPayload):
+    logger.info(f"Direct chat request received. Temperature: {payload.temperature}, Message: '{payload.message[:50]}...'")
     
     custom_api_key_secret = os.getenv("CUSTOM_API_SECRET_KEY")
     custom_api_base_url = os.getenv("CUSTOM_API_BASE_URL", CUSTOM_API_BASE_URL_DEFAULT)
@@ -266,7 +169,7 @@ async def direct_chat(payload: ChatPayload, request: Request): # Added 'request:
 
     async def custom_api_streamer():
         try:
-            logger.info(f"IP {user_ip}: Sending request to Custom API for /chat.")
+            logger.info("Sending request to Custom API for /chat.")
             stream = client.chat.completions.create(
                 model=custom_api_model,
                 temperature=payload.temperature,
@@ -282,43 +185,50 @@ async def direct_chat(payload: ChatPayload, request: Request): # Added 'request:
                 if content_to_yield:
                     yield content_to_yield
         except Exception as e:
-            logger.error(f"IP {user_ip}: Error during Custom API call for /chat: {e}", exc_info=True)
-            yield f"Error processing with Custom API: {str(e)}" # Consider a more generic error for client
+            logger.error(f"Error during Custom API call for /chat: {e}", exc_info=True)
+            yield f"Error processing with Custom API: {str(e)}"
     return StreamingResponse(custom_api_streamer(), media_type="text/plain")
 
 
 @app.post("/gemini/submit_task", response_model=TaskSubmissionResponse)
-async def submit_gemini_task(request_payload: GeminiTaskRequest, background_tasks: BackgroundTasks, http_request: Request): # Renamed request to request_payload and added http_request
-    user_ip = get_user_ip(http_request) # Potentially rate limit this endpoint too if needed
+async def submit_gemini_task(request: GeminiTaskRequest, background_tasks: BackgroundTasks):
     task_id = str(uuid.uuid4())
-    logger.info(f"IP {user_ip}: Received Gemini task submission. Assigning Task ID: {task_id}. Msg: '{request_payload.message[:50]}...'")
+    logger.info(f"Received Gemini task submission. Assigning Task ID: {task_id}. Message: '{request.message[:50]}...'")
 
-    gemini_api_key_from_request = request_payload.api_key
+    gemini_api_key_from_request = request.api_key
     gemini_api_key_secret = os.getenv("GEMINI_API_KEY")
-    key_to_use = gemini_api_key_from_request #or gemini_api_key_secret
+    key_to_use = gemini_api_key_from_request or gemini_api_key_secret
 
     if not key_to_use:
-        logger.error(f"[Task {task_id}] IP {user_ip}: Gemini API Key missing for task submission.")
+        logger.error(f"[Task {task_id}] Gemini API Key missing for task submission.")
         raise HTTPException(status_code=400, detail="Gemini API Key required.")
 
-    requested_model = request_payload.gemini_model or DEFAULT_GEMINI_MODEL
+    requested_model = request.gemini_model or DEFAULT_GEMINI_MODEL
     
     current_time = datetime.now(timezone.utc)
     tasks_db[task_id] = {
-        "status": "PENDING", "result": None, "error": None,
-        "submitted_at": current_time, "last_updated_at": current_time,
-        "request_params": request_payload.model_dump()
+        "status": "PENDING",
+        "result": None,
+        "error": None,
+        "submitted_at": current_time,
+        "last_updated_at": current_time,
+        "request_params": request.model_dump() # Store original request
     }
 
     background_tasks.add_task(
-        process_gemini_request_background, task_id, request_payload.message,
-        request_payload.url, requested_model, key_to_use
+        process_gemini_request_background,
+        task_id,
+        request.message,
+        request.url,
+        requested_model,
+        key_to_use
     )
     
-    logger.info(f"[Task {task_id}] IP {user_ip}: Task submitted to background processing.")
+    logger.info(f"[Task {task_id}] Task submitted to background processing.")
     return TaskSubmissionResponse(
-        task_id=task_id, status="PENDING",
-        task_detail_url=f"/gemini/task/{task_id}"
+        task_id=task_id,
+        status="PENDING",
+        task_detail_url=f"/gemini/task/{task_id}" # Provide the URL to poll
     )
 
 
@@ -342,7 +252,6 @@ async def get_gemini_task_status(task_id: str = Path(..., description="The ID of
     )
 
 @app.get("/")
-async def read_root(request: Request): # Added request
-    user_ip = get_user_ip(request)
-    logger.info(f"IP {user_ip}: Root endpoint '/' accessed (health check).")
-    return {"message": "API for Direct Chat and Async Gemini Tasks with Rate Limiting is running."}
+async def read_root():
+    logger.info("Root endpoint '/' accessed (health check).")
+    return {"message": "API for Direct Chat and Async Gemini Tasks is running."}