webscout 8.3.1__py3-none-any.whl → 8.3.3__py3-none-any.whl

webscout/auth/rate_limiter.py

@@ -0,0 +1,254 @@
+# webscout/auth/rate_limiter.py
+
+import asyncio
+from datetime import datetime, timezone, timedelta
+from typing import Optional, Tuple
+import logging
+
+from .models import APIKey, RateLimitEntry
+from .database import DatabaseManager
+from .config import AppConfig
+
+logger = logging.getLogger(__name__)
+
+
+class RateLimiter:
+    """Rate limiter for API requests."""
+    
+    def __init__(self, database_manager: DatabaseManager):
+        self.db = database_manager
+        self.default_rate_limit = 10  # requests per minute
+        self.window_size = 60  # 1 minute in seconds
+    
+    async def check_rate_limit(self, api_key: APIKey) -> Tuple[bool, dict]:
+        """
+        Check if a request is allowed under the rate limit.
+        
+        Returns:
+            Tuple of (is_allowed, rate_limit_info)
+        """
+        now = datetime.now(timezone.utc)
+        window_start = now - timedelta(seconds=self.window_size)
+        
+        try:
+            # Get or create rate limit entry
+            entry = await self.db.get_rate_limit_entry(api_key.id)
+            
+            if not entry:
+                entry = RateLimitEntry(api_key_id=api_key.id, requests=[])
+            
+            # Clean old requests outside the window
+            entry.requests = [req for req in entry.requests if req > window_start]
+            
+            # Check if limit is exceeded
+            current_count = len(entry.requests)
+            limit = api_key.rate_limit or self.default_rate_limit
+            
+            if current_count >= limit:
+                # Rate limit exceeded
+                oldest_request = min(entry.requests) if entry.requests else now
+                reset_at = oldest_request + timedelta(seconds=self.window_size)
+                retry_after = int((reset_at - now).total_seconds())
+                
+                rate_limit_info = {
+                    "allowed": False,
+                    "limit": limit,
+                    "remaining": 0,
+                    "reset_at": reset_at,
+                    "retry_after": max(retry_after, 1)
+                }
+                
+                logger.warning(f"Rate limit exceeded for API key {api_key.id}: {current_count}/{limit}")
+                return False, rate_limit_info
+            
+            # Add current request
+            entry.requests.append(now)
+            
+            # Update database
+            await self.db.update_rate_limit_entry(entry)
+            
+            # Calculate next reset time
+            if entry.requests:
+                oldest_request = min(entry.requests)
+                reset_at = oldest_request + timedelta(seconds=self.window_size)
+            else:
+                reset_at = now + timedelta(seconds=self.window_size)
+            
+            rate_limit_info = {
+                "allowed": True,
+                "limit": limit,
+                "remaining": limit - len(entry.requests),
+                "reset_at": reset_at,
+                "retry_after": None
+            }
+            
+            return True, rate_limit_info
+            
+        except Exception as e:
+            logger.error(f"Error checking rate limit for API key {api_key.id}: {e}")
+            # In case of error, allow the request but log the issue
+            rate_limit_info = {
+                "allowed": True,
+                "limit": api_key.rate_limit or self.default_rate_limit,
+                "remaining": api_key.rate_limit or self.default_rate_limit,
+                "reset_at": now + timedelta(seconds=self.window_size),
+                "retry_after": None
+            }
+            return True, rate_limit_info
+    
+    async def reset_rate_limit(self, api_key_id: str) -> bool:
+        """Reset rate limit for an API key (admin function)."""
+        try:
+            entry = RateLimitEntry(api_key_id=api_key_id, requests=[])
+            await self.db.update_rate_limit_entry(entry)
+            logger.info(f"Reset rate limit for API key {api_key_id}")
+            return True
+        except Exception as e:
+            logger.error(f"Error resetting rate limit for API key {api_key_id}: {e}")
+            return False
+    
+    async def get_rate_limit_status(self, api_key: APIKey) -> dict:
+        """Get current rate limit status without making a request."""
+        now = datetime.now(timezone.utc)
+        window_start = now - timedelta(seconds=self.window_size)
+        
+        try:
+            entry = await self.db.get_rate_limit_entry(api_key.id)
+            
+            if not entry:
+                entry = RateLimitEntry(api_key_id=api_key.id, requests=[])
+            
+            # Clean old requests
+            entry.requests = [req for req in entry.requests if req > window_start]
+            
+            limit = api_key.rate_limit or self.default_rate_limit
+            current_count = len(entry.requests)
+            
+            # Calculate reset time
+            if entry.requests:
+                oldest_request = min(entry.requests)
+                reset_at = oldest_request + timedelta(seconds=self.window_size)
+            else:
+                reset_at = now + timedelta(seconds=self.window_size)
+            
+            return {
+                "limit": limit,
+                "remaining": max(0, limit - current_count),
+                "reset_at": reset_at,
+                "window_size": self.window_size
+            }
+            
+        except Exception as e:
+            logger.error(f"Error getting rate limit status for API key {api_key.id}: {e}")
+            return {
+                "limit": api_key.rate_limit or self.default_rate_limit,
+                "remaining": api_key.rate_limit or self.default_rate_limit,
+                "reset_at": now + timedelta(seconds=self.window_size),
+                "window_size": self.window_size
+            }
+    
+    async def cleanup_old_entries(self) -> int:
+        """Clean up old rate limit entries (maintenance function)."""
+        # Remove requests older than the window_size for all rate limit entries
+        try:
+            # Try to get all rate limit entries from the database
+            if hasattr(self.db, 'get_all_rate_limit_entries'):
+                entries = await self.db.get_all_rate_limit_entries()
+            else:
+                logger.warning("Database does not support get_all_rate_limit_entries; cleanup skipped.")
+                return 0
+
+            now = datetime.now(timezone.utc)
+            window_start = now - timedelta(seconds=self.window_size)
+            cleaned = 0
+
+            for entry in entries:
+                old_count = len(entry.requests)
+                entry.requests = [req for req in entry.requests if req > window_start]
+                if len(entry.requests) < old_count:
+                    await self.db.update_rate_limit_entry(entry)
+                    cleaned += 1
+            logger.info(f"Cleaned up {cleaned} old rate limit entries.")
+            return cleaned
+        except Exception as e:
+            logger.error(f"Error cleaning up old rate limit entries: {e}")
+            return 0
+
+    async def check_ip_rate_limit(self, client_ip: str) -> Tuple[bool, dict]:
+        """
+        Check rate limit for an IP address (used in no-auth mode).
+
+        Returns:
+            Tuple of (is_allowed, rate_limit_info)
+        """
+        now = datetime.now(timezone.utc)
+        window_start = now - timedelta(seconds=self.window_size)
+
+        # Create a pseudo API key for IP-based rate limiting
+        ip_key_id = f"ip_{client_ip}"
+        limit = AppConfig.default_rate_limit
+
+        try:
+            # Get or create rate limit entry for IP
+            entry = await self.db.get_rate_limit_entry(ip_key_id)
+
+            if not entry:
+                entry = RateLimitEntry(api_key_id=ip_key_id, requests=[])
+
+            # Clean old requests outside the window
+            entry.requests = [req for req in entry.requests if req > window_start]
+
+            # Check if limit is exceeded
+            current_count = len(entry.requests)
+
+            if current_count >= limit:
+                # Rate limit exceeded
+                oldest_request = min(entry.requests) if entry.requests else now
+                reset_at = oldest_request + timedelta(seconds=self.window_size)
+                retry_after = int((reset_at - now).total_seconds())
+
+                rate_limit_info = {
+                    "allowed": False,
+                    "limit": limit,
+                    "remaining": 0,
+                    "reset_at": reset_at,
+                    "retry_after": max(retry_after, 1)
+                }
+
+                logger.warning(f"Rate limit exceeded for IP {client_ip}: {current_count}/{limit}")
+                return False, rate_limit_info
+
+            # Add current request
+            entry.requests.append(now)
+
+            # Update database
+            await self.db.update_rate_limit_entry(entry)
+
+            # Calculate next reset time
+            if entry.requests:
+                oldest_request = min(entry.requests)
+                reset_at = oldest_request + timedelta(seconds=self.window_size)
+            else:
+                reset_at = now + timedelta(seconds=self.window_size)
+
+            rate_limit_info = {
+                "allowed": True,
+                "limit": limit,
+                "remaining": limit - len(entry.requests),
+                "reset_at": reset_at,
+                "retry_after": None
+            }
+
+            return True, rate_limit_info
+
+        except Exception as e:
+            logger.error(f"Error checking rate limit for IP {client_ip}: {e}")
+            # In case of error, allow the request but log the issue
+            rate_limit_info = {
+                "allowed": True,
+                "limit": limit,
+                "remaining": limit,
+                "reset_at": now + timedelta(seconds=self.window_size),
+                "retry_after": None
+            }
+            return True, rate_limit_info

webscout/auth/request_models.py

@@ -0,0 +1,127 @@
+"""
+Pydantic models for API requests and responses.
+"""
+
+from typing import List, Dict, Optional, Union, Any, Literal
+from webscout.Provider.OPENAI.pydantic_imports import BaseModel, Field
+
+
+# Define Pydantic models for multimodal content parts, aligning with OpenAI's API
+class TextPart(BaseModel):
+    """Text content part for multimodal messages."""
+    type: Literal["text"]
+    text: str
+
+
+class ImageURL(BaseModel):
+    """Image URL configuration for multimodal messages."""
+    url: str  # Can be http(s) or data URI
+    detail: Optional[Literal["auto", "low", "high"]] = Field(
+        "auto",
+        description="Specifies the detail level of the image."
+    )
+
+
+class ImagePart(BaseModel):
+    """Image content part for multimodal messages."""
+    type: Literal["image_url"]
+    image_url: ImageURL
+
+
+MessageContentParts = Union[TextPart, ImagePart]
+
+
+class Message(BaseModel):
+    """Chat message model compatible with OpenAI API."""
+    role: Literal["system", "user", "assistant", "function", "tool"]
+    content: Optional[Union[str, List[MessageContentParts]]] = Field(
+        None,
+        description="The content of the message. Can be a string, a list of content parts (for multimodal), or null."
+    )
+    name: Optional[str] = None
+
+
+class ChatCompletionRequest(BaseModel):
+    """Request model for chat completions."""
+    model: str = Field(..., description="ID of the model to use. See the model endpoint for the available models.")
+    messages: List[Message] = Field(..., description="A list of messages comprising the conversation so far.")
+    temperature: Optional[float] = Field(None, description="What sampling temperature to use, between 0 and 2.")
+    top_p: Optional[float] = Field(None, description="An alternative to sampling with temperature, called nucleus sampling.")
+    n: Optional[int] = Field(1, description="How many chat completion choices to generate for each input message.")
+    stream: Optional[bool] = Field(False, description="If set, partial message deltas will be sent, like in ChatGPT.")
+    max_tokens: Optional[int] = Field(None, description="The maximum number of tokens to generate in the chat completion.")
+    presence_penalty: Optional[float] = Field(None, description="Number between -2.0 and 2.0. Positive values penalize new tokens based on whether they appear in the text so far.")
+    frequency_penalty: Optional[float] = Field(None, description="Number between -2.0 and 2.0. Positive values penalize new tokens based on their existing frequency in the text so far.")
+    logit_bias: Optional[Dict[str, float]] = Field(None, description="Modify the likelihood of specified tokens appearing in the completion.")
+    user: Optional[str] = Field(None, description="A unique identifier representing your end-user.")
+    stop: Optional[Union[str, List[str]]] = Field(None, description="Up to 4 sequences where the API will stop generating further tokens.")
+
+    class Config:
+        extra = "ignore"
+        schema_extra = {
+            "example": {
+                "model": "Cloudflare/@cf/meta/llama-4-scout-17b-16e-instruct",
+                "messages": [
+                    {"role": "system", "content": "You are a helpful assistant."},
+                    {"role": "user", "content": "Hello, how are you?"}
+                ],
+                "temperature": 0.7,
+                "max_tokens": 150,
+                "stream": False
+            }
+        }
+
+
+class ImageGenerationRequest(BaseModel):
+    """Request model for OpenAI-compatible image generation endpoint."""
+    prompt: str = Field(..., description="A text description of the desired image(s). The maximum length is 1000 characters.")
+    model: str = Field(..., description="The model to use for image generation.")
+    n: Optional[int] = Field(1, description="The number of images to generate. Must be between 1 and 10.")
+    size: Optional[str] = Field("1024x1024", description="The size of the generated images. Must be one of: '256x256', '512x512', or '1024x1024'.")
+    response_format: Optional[Literal["url", "b64_json"]] = Field("url", description="The format in which the generated images are returned.")
+    user: Optional[str] = Field(None, description="A unique identifier representing your end-user.")
+    style: Optional[str] = Field(None, description="Optional style for the image (provider/model-specific).")
+    aspect_ratio: Optional[str] = Field(None, description="Optional aspect ratio for the image (provider/model-specific).")
+    timeout: Optional[int] = Field(None, description="Optional timeout for the image generation request in seconds.")
+    image_format: Optional[str] = Field(None, description="Optional image format (e.g., 'png', 'jpeg').")
+    seed: Optional[int] = Field(None, description="Optional random seed for reproducibility.")
+
+    class Config:
+        extra = "ignore"
+        schema_extra = {
+            "example": {
+                "prompt": "A futuristic cityscape at sunset, digital art",
+                "model": "PollinationsAI/turbo",
+                "n": 1,
+                "size": "1024x1024",
+                "response_format": "url",
+                "user": "user-1234"
+            }
+        }
+
+
+class ModelInfo(BaseModel):
+    """Model information for the models endpoint."""
+    id: str
+    object: str = "model"
+    created: int
+    owned_by: str
+
+
+class ModelListResponse(BaseModel):
+    """Response model for the models list endpoint."""
+    object: str = "list"
+    data: List[ModelInfo]
+
+
+class ErrorDetail(BaseModel):
+    """Error detail structure compatible with OpenAI API."""
+    message: str
+    type: str = "server_error"
+    param: Optional[str] = None
+    code: Optional[str] = None
+
+
+class ErrorResponse(BaseModel):
+    """Error response structure compatible with OpenAI API."""
+    error: ErrorDetail

webscout/auth/request_processing.py

@@ -0,0 +1,226 @@
+"""
+Request processing utilities for the Webscout API.
+"""
+
+import json
+import time
+import uuid
+from typing import List, Dict, Any
+from starlette.status import HTTP_422_UNPROCESSABLE_ENTITY, HTTP_500_INTERNAL_SERVER_ERROR
+from fastapi.responses import StreamingResponse
+
+from webscout.Provider.OPENAI.utils import ChatCompletion, Choice, ChatCompletionMessage, CompletionUsage
+from webscout.Litlogger import Logger, LogLevel, LogFormat, ConsoleHandler
+import sys
+
+from .request_models import Message, ChatCompletionRequest
+from .exceptions import APIError, clean_text
+
+# Setup logger
+logger = Logger(
+    name="webscout.api",
+    level=LogLevel.INFO,
+    handlers=[ConsoleHandler(stream=sys.stdout)],
+    fmt=LogFormat.DEFAULT
+)
+
+
+def process_messages(messages: List[Message]) -> List[Dict[str, Any]]:
+    """Process and validate chat messages."""
+    processed_messages = []
+
+    for i, msg_in in enumerate(messages):
+        try:
+            message_dict_out = {"role": msg_in.role}
+
+            if msg_in.content is None:
+                message_dict_out["content"] = None
+            elif isinstance(msg_in.content, str):
+                message_dict_out["content"] = msg_in.content
+            else:  # List[MessageContentParts]
+                message_dict_out["content"] = [
+                    part.model_dump(exclude_none=True) for part in msg_in.content
+                ]
+
+            if msg_in.name:
+                message_dict_out["name"] = msg_in.name
+
+            processed_messages.append(message_dict_out)
+
+        except Exception as e:
+            raise APIError(
+                f"Invalid message at index {i}: {str(e)}",
+                HTTP_422_UNPROCESSABLE_ENTITY,
+                "invalid_request_error",
+                param=f"messages[{i}]"
+            )
+
+    return processed_messages
+
+
+def prepare_provider_params(chat_request: ChatCompletionRequest, model_name: str,
+                          processed_messages: List[Dict[str, Any]]) -> Dict[str, Any]:
+    """Prepare parameters for the provider."""
+    params = {
+        "model": model_name,
+        "messages": processed_messages,
+        "stream": chat_request.stream,
+    }
+
+    # Add optional parameters if present
+    optional_params = [
+        "temperature", "max_tokens", "top_p", "presence_penalty",
+        "frequency_penalty", "stop", "user"
+    ]
+
+    for param in optional_params:
+        value = getattr(chat_request, param, None)
+        if value is not None:
+            params[param] = value
+
+    return params
+
+
+async def handle_streaming_response(provider: Any, params: Dict[str, Any], request_id: str) -> StreamingResponse:
+    """Handle streaming chat completion response."""
+    async def streaming():
+        try:
+            logger.debug(f"Starting streaming response for request {request_id}")
+            completion_stream = provider.chat.completions.create(**params)
+
+            # Check if it's iterable (generator, iterator, or other iterable types)
+            if hasattr(completion_stream, '__iter__') and not isinstance(completion_stream, (str, bytes, dict)):
+                try:
+                    for chunk in completion_stream:
+                        # Standardize chunk format before sending
+                        if hasattr(chunk, 'model_dump'):  # Pydantic v2
+                            chunk_data = chunk.model_dump(exclude_none=True)
+                        elif hasattr(chunk, 'dict'):  # Pydantic v1
+                            chunk_data = chunk.dict(exclude_none=True)
+                        elif isinstance(chunk, dict):
+                            chunk_data = chunk
+                        else:  # Fallback for unknown chunk types
+                            chunk_data = chunk
+                        
+                        # Clean text content in the chunk to remove control characters
+                        if isinstance(chunk_data, dict) and 'choices' in chunk_data:
+                            for choice in chunk_data.get('choices', []):
+                                if isinstance(choice, dict):
+                                    # Handle delta for streaming
+                                    if 'delta' in choice and isinstance(choice['delta'], dict) and 'content' in choice['delta']:
+                                        choice['delta']['content'] = clean_text(choice['delta']['content'])
+                                    # Handle message for non-streaming
+                                    elif 'message' in choice and isinstance(choice['message'], dict) and 'content' in choice['message']:
+                                        choice['message']['content'] = clean_text(choice['message']['content'])
+                        
+                        yield f"data: {json.dumps(chunk_data, ensure_ascii=False)}\n\n"
+                except TypeError as te:
+                    logger.error(f"Error iterating over completion_stream: {te}")
+                    # Fall back to treating as non-generator response
+                    if hasattr(completion_stream, 'model_dump'):
+                        response_data = completion_stream.model_dump(exclude_none=True)
+                    elif hasattr(completion_stream, 'dict'):
+                        response_data = completion_stream.dict(exclude_none=True)
+                    else:
+                        response_data = completion_stream
+                    
+                    # Clean text content in the response
+                    if isinstance(response_data, dict) and 'choices' in response_data:
+                        for choice in response_data.get('choices', []):
+                            if isinstance(choice, dict):
+                                if 'delta' in choice and isinstance(choice['delta'], dict) and 'content' in choice['delta']:
+                                    choice['delta']['content'] = clean_text(choice['delta']['content'])
+                                elif 'message' in choice and isinstance(choice['message'], dict) and 'content' in choice['message']:
+                                    choice['message']['content'] = clean_text(choice['message']['content'])
+                    
+                    yield f"data: {json.dumps(response_data, ensure_ascii=False)}\n\n"
+            else:  # Non-generator response
+                if hasattr(completion_stream, 'model_dump'):
+                    response_data = completion_stream.model_dump(exclude_none=True)
+                elif hasattr(completion_stream, 'dict'):
+                    response_data = completion_stream.dict(exclude_none=True)
+                else:
+                    response_data = completion_stream
+                
+                # Clean text content in the response
+                if isinstance(response_data, dict) and 'choices' in response_data:
+                    for choice in response_data.get('choices', []):
+                        if isinstance(choice, dict):
+                            if 'delta' in choice and isinstance(choice['delta'], dict) and 'content' in choice['delta']:
+                                choice['delta']['content'] = clean_text(choice['delta']['content'])
+                            elif 'message' in choice and isinstance(choice['message'], dict) and 'content' in choice['message']:
+                                choice['message']['content'] = clean_text(choice['message']['content'])
+                
+                yield f"data: {json.dumps(response_data, ensure_ascii=False)}\n\n"
+
+        except Exception as e:
+            logger.error(f"Error in streaming response for request {request_id}: {e}")
+            error_message = clean_text(str(e))
+            error_data = {
+                "error": {
+                    "message": error_message,
+                    "type": "server_error",
+                    "code": "streaming_error"
+                }
+            }
+            yield f"data: {json.dumps(error_data, ensure_ascii=False)}\n\n"
+        finally:
+            yield "data: [DONE]\n\n"
+    return StreamingResponse(streaming(), media_type="text/event-stream")
+
+
+async def handle_non_streaming_response(provider: Any, params: Dict[str, Any],
+                                      request_id: str, start_time: float) -> Dict[str, Any]:
+    """Handle non-streaming chat completion response."""
+    try:
+        logger.debug(f"Starting non-streaming response for request {request_id}")
+        completion = provider.chat.completions.create(**params)
+
+        if completion is None:
+            # Return a valid OpenAI-compatible error response
+            return ChatCompletion(
+                id=request_id,
+                created=int(time.time()),
+                model=params.get("model", "unknown"),
+                choices=[Choice(
+                    index=0,
+                    message=ChatCompletionMessage(role="assistant", content="No response generated."),
+                    finish_reason="error"
+                )],
+                usage=CompletionUsage(prompt_tokens=0, completion_tokens=0, total_tokens=0)
+            ).model_dump(exclude_none=True)
+
+        # Standardize response format
+        if hasattr(completion, "model_dump"):  # Pydantic v2
+            response_data = completion.model_dump(exclude_none=True)
+        elif hasattr(completion, "dict"):  # Pydantic v1
+            response_data = completion.dict(exclude_none=True)
+        elif isinstance(completion, dict):
+            response_data = completion
+        else:
+            raise APIError(
+                "Invalid response format from provider",
+                HTTP_500_INTERNAL_SERVER_ERROR,
+                "provider_error"
+            )
+        
+        # Clean text content in the response to remove control characters
+        if isinstance(response_data, dict) and 'choices' in response_data:
+            for choice in response_data.get('choices', []):
+                if isinstance(choice, dict) and 'message' in choice:
+                    if isinstance(choice['message'], dict) and 'content' in choice['message']:
+                        choice['message']['content'] = clean_text(choice['message']['content'])
+
+        elapsed = time.time() - start_time
+        logger.info(f"Completed non-streaming request {request_id} in {elapsed:.2f}s")
+
+        return response_data
+
+    except Exception as e:
+        logger.error(f"Error in non-streaming response for request {request_id}: {e}")
+        error_message = clean_text(str(e))
+        raise APIError(
+            f"Provider error: {error_message}",
+            HTTP_500_INTERNAL_SERVER_ERROR,
+            "provider_error"
+        )