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

webscout/Provider/OPENAI/BLACKBOXAI.py

@@ -1,22 +1,204 @@
+# from pickle import NONE
+import requests
 import requests
 import random
 import string
 import base64
 from datetime import datetime, timedelta
 from typing import Generator, List, Dict, Optional, Any, Union
-import json # Not used directly in this snippet, but often useful
 import uuid
 import time
-
-# Import base classes and utility structures
-from webscout.Provider.OPENAI.base import OpenAICompatibleProvider, BaseChat, BaseCompletions
+import codecs
+import gzip
+import zstandard as zstd
+import brotli
+import zlib
 from webscout.Provider.OPENAI.utils import (
     ChatCompletion, Choice,
     ChatCompletionMessage, CompletionUsage, count_tokens,
-    ChatCompletionChunk, ChoiceDelta # Added for streaming return type
+    ChatCompletionChunk, ChoiceDelta  # Ensure ChoiceDelta is always imported at the top
 )
-from webscout.litagent import LitAgent, agent
-agent = LitAgent()
+try:
+    from webscout.litagent import LitAgent
+    agent = LitAgent()
+except ImportError:
+    print("Warning: LitAgent not available, using default user agent")
+    class MockAgent:
+        def random(self):
+            return 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/118.0.0.0 Safari/537.36'
+    agent = MockAgent()
+
+from webscout.Provider.OPENAI.base import OpenAICompatibleProvider, BaseChat, BaseCompletions
+
+class StreamingDecompressor:
+    """
+    A streaming decompressor that can handle partial compressed data in real-time.
+    This allows for true streaming decompression without buffering entire response.
+    """
+    def __init__(self, content_encoding: str):
+        self.encoding = content_encoding.lower().strip() if content_encoding else None
+        self.decompressor = None
+        self.text_decoder = codecs.getincrementaldecoder("utf-8")("replace")
+        self.zstd_buffer = b""  # Buffer for zstd incomplete frames
+        
+        if self.encoding == 'gzip':
+            self.decompressor = zlib.decompressobj(16 + zlib.MAX_WBITS)  # gzip format
+        elif self.encoding == 'deflate':
+            self.decompressor = zlib.decompressobj()  # deflate format
+        elif self.encoding == 'zstd':
+            self.decompressor = zstd.ZstdDecompressor()
+        elif self.encoding == 'br':
+            self.decompressor = brotli.Decompressor()
+    
+    def decompress_chunk(self, chunk: bytes) -> str:
+        """
+        Decompress a chunk of data and return decoded text.
+        Handles partial compressed data properly for real-time streaming.
+        """
+        try:
+            if not chunk:
+                return ""
+            
+            if not self.encoding or self.encoding not in ['gzip', 'deflate', 'zstd', 'br']:
+                # No compression or unsupported - decode directly
+                return self.text_decoder.decode(chunk, final=False)
+            
+            if self.encoding in ['gzip', 'deflate']:
+                # Use zlib decompressor for gzip/deflate
+                decompressed_data = self.decompressor.decompress(chunk)
+                return self.text_decoder.decode(decompressed_data, final=False)
+            
+            elif self.encoding == 'zstd':
+                # Zstandard streaming decompression with buffering for incomplete frames
+                self.zstd_buffer += chunk
+                try:
+                    # Try to decompress the current buffer
+                    decompressed_data = self.decompressor.decompress(self.zstd_buffer)
+                    # If successful, clear the buffer and return decoded text
+                    self.zstd_buffer = b""
+                    return self.text_decoder.decode(decompressed_data, final=False)
+                except zstd.ZstdError:
+                    # Frame is incomplete, keep buffering
+                    # Try to decompress any complete frames from buffer start
+                    try:
+                        # Process buffer in chunks to find complete frames
+                        buffer_len = len(self.zstd_buffer)
+                        if buffer_len > 4:  # Minimum zstd frame size
+                            # Try smaller chunks of the buffer
+                            for end_pos in range(4, buffer_len + 1):
+                                try:
+                                    partial_data = self.decompressor.decompress(self.zstd_buffer[:end_pos])
+                                    # If we got here, we found a complete frame
+                                    self.zstd_buffer = self.zstd_buffer[end_pos:]
+                                    return self.text_decoder.decode(partial_data, final=False)
+                                except zstd.ZstdError:
+                                    continue
+                    except Exception:
+                        pass
+                    return ""
+            
+            elif self.encoding == 'br':
+                # Brotli streaming decompression
+                try:
+                    decompressed_data = self.decompressor.decompress(chunk)
+                    return self.text_decoder.decode(decompressed_data, final=False)
+                except brotli.error:
+                    # If brotli fails, it might need more data or be at end
+                    return ""
+        
+        except Exception as e:
+            # If decompression fails, try to decode the chunk as-is (fallback)
+            try:
+                return self.text_decoder.decode(chunk, final=False)
+            except UnicodeDecodeError:
+                return ""
+    
+    def finalize(self) -> str:
+        """
+        Finalize the decompression and return any remaining decoded text.
+        """
+        try:
+            remaining_text = ""
+            
+            if self.encoding in ['gzip', 'deflate'] and self.decompressor:
+                # Flush any remaining compressed data
+                remaining_data = self.decompressor.flush()
+                if remaining_data:
+                    remaining_text = self.text_decoder.decode(remaining_data, final=True)
+                else:
+                    remaining_text = self.text_decoder.decode(b"", final=True)
+            elif self.encoding == 'zstd':
+                # Process any remaining buffered data
+                if self.zstd_buffer:
+                    try:
+                        remaining_data = self.decompressor.decompress(self.zstd_buffer)
+                        remaining_text = self.text_decoder.decode(remaining_data, final=True)
+                    except:
+                        # If buffered data can't be decompressed, finalize decoder
+                        remaining_text = self.text_decoder.decode(b"", final=True)
+                else:
+                    remaining_text = self.text_decoder.decode(b"", final=True)
+            else:
+                # Finalize the text decoder for other encodings
+                remaining_text = self.text_decoder.decode(b"", final=True)
+            
+            return remaining_text
+        except Exception:
+            # Ensure we always finalize the text decoder
+            try:
+                return self.text_decoder.decode(b"", final=True)
+            except:
+                return ""
+
+def decompress_response(response_content: bytes, content_encoding: str) -> str:
+    """
+    Decompress response content based on the Content-Encoding header.
+    
+    Args:
+        response_content: The raw response content as bytes
+        content_encoding: The Content-Encoding header value
+        
+    Returns:
+        str: The decompressed and decoded content as UTF-8 string
+        
+    Raises:
+        IOError: If decompression fails
+    """
+    try:
+        if not content_encoding:
+            # No compression, decode directly
+            return response_content.decode('utf-8')
+        
+        encoding = content_encoding.lower().strip()
+        
+        if encoding == 'zstd':
+            # Decompress using zstandard
+            dctx = zstd.ZstdDecompressor()
+            decompressed_data = dctx.decompress(response_content)
+            return decompressed_data.decode('utf-8')
+        
+        elif encoding == 'gzip':
+            # Decompress using gzip
+            decompressed_data = gzip.decompress(response_content)
+            return decompressed_data.decode('utf-8')
+        
+        elif encoding == 'br':
+            # Decompress using brotli
+            decompressed_data = brotli.decompress(response_content)
+            return decompressed_data.decode('utf-8')
+        
+        elif encoding == 'deflate':
+            # Decompress using zlib (deflate)
+            import zlib
+            decompressed_data = zlib.decompress(response_content)
+            return decompressed_data.decode('utf-8')
+        
+        else:
+            # Unknown or unsupported encoding, try to decode as-is
+            return response_content.decode('utf-8')
+            
+    except Exception as e:
+        raise IOError(f"Failed to decompress response with encoding '{content_encoding}': {str(e)}") from e
 
 def to_data_uri(image_data):
     """Convert image data to a data URI format"""
@@ -38,6 +220,19 @@ def to_data_uri(image_data):
 
     return f"data:{mime_type};base64,{encoded}"
 
+def clean_text(text):
+    """Clean text by removing null bytes and control characters except newlines and tabs."""
+    import re
+    if not isinstance(text, str):
+        return text
+    
+    # Remove null bytes
+    text = text.replace('\x00', '')
+    
+    # Keep newlines, tabs, and other printable characters, remove other control chars
+    # This regex matches control characters except \n, \r, \t
+    return re.sub(r'[\x01-\x08\x0b\x0c\x0e-\x1f\x7f]', '', text)
+
 
 class Completions(BaseCompletions):
     def __init__(self, client: 'BLACKBOXAI'):
@@ -52,6 +247,8 @@ class Completions(BaseCompletions):
         stream: bool = False,
         temperature: Optional[float] = None,
         top_p: Optional[float] = None,
+        timeout: Optional[int] = None,
+        proxies: Optional[dict] = None,
         **kwargs: Any
     ) -> Union[ChatCompletion, Generator[ChatCompletionChunk, None, None]]:
         """
@@ -117,7 +314,9 @@ class Completions(BaseCompletions):
                 max_tokens=max_tokens,
                 temperature=temperature,
                 top_p=top_p,
-                media=media
+                media=media,
+                timeout=timeout,
+                proxies=proxies
             )
 
         # Use non-streaming implementation
@@ -130,7 +329,9 @@ class Completions(BaseCompletions):
             max_tokens=max_tokens,
             temperature=temperature,
             top_p=top_p,
-            media=media
+            media=media,
+            timeout=timeout,
+            proxies=proxies
         )
 
 
@@ -145,9 +346,18 @@ class Completions(BaseCompletions):
         max_tokens: Optional[int] = None,
         temperature: Optional[float] = None,
         top_p: Optional[float] = None,
-        media: List = None
+        media: List = None,
+        timeout: Optional[int] = None,
+        proxies: Optional[dict] = None
     ) -> ChatCompletion:
         """Implementation for non-streaming chat completions."""
+        original_proxies = self._client.session.proxies.copy()
+        # Only use proxies if they are explicitly provided and not causing issues
+        if proxies is not None and proxies:
+            self._client.session.proxies.update(proxies)
+        else:
+            # Clear proxies to avoid connection issues
+            self._client.session.proxies = {}
         try:
             # Prepare user messages for BlackboxAI API format
             blackbox_messages = []
@@ -191,20 +401,39 @@ class Completions(BaseCompletions):
                 model=model
             )
 
-            # Make the API request with cookies
-            response = self._client.session.post(
-                self._client.api_endpoint,
-                json=payload,
-                headers=self._client.headers,
-                cookies=self._client.cookies,
-                timeout=self._client.timeout
-            )
+            # Make the API request with cookies and retry logic
+            max_retries = 3
+            for attempt in range(max_retries):
+                try:
+                    response = self._client.session.post(
+                        self._client.api_endpoint,
+                        json=payload,
+                        headers=self._client.headers,
+                        cookies=self._client.cookies,
+                        timeout=timeout if timeout is not None else self._client.timeout
+                    )
+                    break  # Success, exit retry loop
+                except (requests.exceptions.ConnectionError, requests.exceptions.ProxyError) as e:
+                    if attempt == max_retries - 1:
+                        raise IOError(f"BlackboxAI connection failed after {max_retries} attempts: {str(e)}") from e
+                    # Clear proxies and retry
+                    self._client.session.proxies = {}
+                    time.sleep(1)  # Wait before retry
 
             # Process the response
             full_content = ""
             if response.status_code == 200:
-                # Extract content from response text
-                response_text = response.text
+                # Check for Content-Encoding header
+                content_encoding = response.headers.get('Content-Encoding')
+                
+                # Decompress the response if needed
+                try:
+                    response_text = decompress_response(response.content, content_encoding)
+                except IOError as e:
+                    # If decompression fails, fall back to the original method
+                    print(f"Warning: {e}. Falling back to original decoding method.")
+                    decoder = codecs.getincrementaldecoder("utf-8")("replace")
+                    response_text = decoder.decode(response.content, final=True)
 
                 # Handle possible SSE format in response
                 if "data: " in response_text:
@@ -227,10 +456,11 @@ class Completions(BaseCompletions):
                 # Handle error response
                 raise IOError(f"BlackboxAI request failed with status code {response.status_code}")
 
-            # Create the completion message
+            # Clean and create the completion message
+            cleaned_content = clean_text(full_content)
             message = ChatCompletionMessage(
                 role="assistant",
-                content=full_content
+                content=cleaned_content
             )
 
             # Create the choice with the message
@@ -242,7 +472,7 @@ class Completions(BaseCompletions):
 
             # Estimate token usage using count_tokens
             prompt_tokens = count_tokens([str(msg.get("content", "")) for msg in messages])
-            completion_tokens = count_tokens(full_content)
+            completion_tokens = count_tokens(cleaned_content)
 
             # Create the final completion object
             completion = ChatCompletion(
@@ -261,6 +491,9 @@ class Completions(BaseCompletions):
 
         except Exception as e:
             raise IOError(f"BlackboxAI request failed: {str(e)}") from e
+        finally:
+            # Restore original proxies
+            self._client.session.proxies = original_proxies
 
     def _create_streaming(
         self,
@@ -273,91 +506,146 @@ class Completions(BaseCompletions):
         max_tokens: Optional[int] = None,
         temperature: Optional[float] = None,
         top_p: Optional[float] = None,
-        media: List = None
+        media: List = None,
+        timeout: Optional[int] = None,
+        proxies: Optional[dict] = None
     ):
         """Implementation for streaming chat completions (OpenAI-compatible chunks)."""
-        # Prepare user messages for BlackboxAI API format
-        blackbox_messages = []
-        for i, msg in enumerate(messages):
-            if msg["role"] == "system":
-                continue  # System message handled separately
-            msg_id = self._client.generate_id() if i > 0 else request_id
-            blackbox_messages.append({
-                "id": msg_id,
-                "content": msg["content"],
-                "role": msg["role"]
-            })
-        # Add image data if provided
-        if media and blackbox_messages:
-            blackbox_messages[-1]['data'] = {
-                "imagesData": [
-                    {
-                        "filePath": f"/",
-                        "contents": to_data_uri(image[0])
-                    } for image in media
-                ],
-                "fileText": "",
-                "title": ""
-            }
-        # Generate request payload with session
-        request_email = f"{self._client.generate_random_string(8)}@blackbox.ai"
-        session_data = self._client.generate_session(request_email)
-        payload = self._client.create_request_payload(
-            messages=blackbox_messages,
-            chat_id=request_id,
-            system_message=system_message,
-            max_tokens=max_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            session_data=session_data,
-            model=model
-        )
-        # Make the API request with cookies, stream=True
-        response = self._client.session.post(
-            self._client.api_endpoint,
-            json=payload,
-            headers=self._client.headers,
-            cookies=self._client.cookies,
-            stream=True,
-            timeout=self._client.timeout
-        )
-        # Blackbox streams as raw text, no line breaks, so chunk manually
-        buffer = ""
-        chunk_size = 32  # Tune as needed for smoothness
-        from webscout.Provider.OPENAI.utils import ChatCompletionChunk, Choice, ChoiceDelta
-        for chunk in response.iter_content(chunk_size=chunk_size):
-            if not chunk:
-                continue
-            text = chunk.decode(errors="ignore")
-            buffer += text
-            # Yield in small pieces, but only non-empty
-            while len(buffer) >= chunk_size:
-                out = buffer[:chunk_size]
-                buffer = buffer[chunk_size:]
-                if out.strip():
-                    # Wrap the chunk in OpenAI-compatible structure
-                    delta = ChoiceDelta(content=out, role="assistant")
-                    choice = Choice(index=0, delta=delta, finish_reason=None)
-                    chunk_obj = ChatCompletionChunk(
-                        id=request_id,
-                        choices=[choice],
-                        created=created_time,
-                        model=model,
-                        system_fingerprint=None
+        original_proxies = self._client.session.proxies.copy()
+        # Only use proxies if they are explicitly provided and not causing issues
+        if proxies is not None and proxies:
+            self._client.session.proxies.update(proxies)
+        else:
+            # Clear proxies to avoid connection issues
+            self._client.session.proxies = {}
+        try:
+            # Prepare user messages for BlackboxAI API format
+            blackbox_messages = []
+            for i, msg in enumerate(messages):
+                if msg["role"] == "system":
+                    continue  # System message handled separately
+                msg_id = self._client.generate_id() if i > 0 else request_id
+                blackbox_messages.append({
+                    "id": msg_id,
+                    "content": msg["content"],
+                    "role": msg["role"]
+                })
+            # Add image data if provided
+            if media and blackbox_messages:
+                blackbox_messages[-1]['data'] = {
+                    "imagesData": [
+                        {
+                            "filePath": f"/",
+                            "contents": to_data_uri(image[0])
+                        } for image in media
+                    ],
+                    "fileText": "",
+                    "title": ""
+                }
+            # Generate request payload with session
+            request_email = f"{self._client.generate_random_string(8)}@blackbox.ai"
+            session_data = self._client.generate_session(request_email)
+            payload = self._client.create_request_payload(
+                messages=blackbox_messages,
+                chat_id=request_id,
+                system_message=system_message,
+                max_tokens=max_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                session_data=session_data,
+                model=model
+            )
+            # Make the API request with cookies, stream=True and retry logic
+            max_retries = 3
+            for attempt in range(max_retries):
+                try:
+                    response = self._client.session.post(
+                        self._client.api_endpoint,
+                        json=payload,
+                        headers=self._client.headers,
+                        cookies=self._client.cookies,
+                        stream=True,
+                        timeout=timeout if timeout is not None else self._client.timeout
                     )
-                    yield chunk_obj
-        # Yield any remaining buffer
-        if buffer.strip():
-            delta = ChoiceDelta(content=buffer, role="assistant")
-            choice = Choice(index=0, delta=delta, finish_reason=None)
-            chunk_obj = ChatCompletionChunk(
+                    break  # Success, exit retry loop
+                except (requests.exceptions.ConnectionError, requests.exceptions.ProxyError) as e:
+                    if attempt == max_retries - 1:
+                        raise IOError(f"BlackboxAI connection failed after {max_retries} attempts: {str(e)}") from e
+                    # Clear proxies and retry
+                    self._client.session.proxies = {}
+                    time.sleep(1)  # Wait before retry
+            # Blackbox streams as raw text, no line breaks, so chunk manually
+            chunk_size = 32  # Tune as needed for smoothness
+            # ChoiceDelta is already imported at the top of the file
+            
+            # Check if the response is compressed and create appropriate decompressor
+            content_encoding = response.headers.get('Content-Encoding')
+            streaming_decompressor = StreamingDecompressor(content_encoding)
+            
+            # Stream with real-time decompression
+            for chunk in response.iter_content(chunk_size=chunk_size):
+                if not chunk:
+                    continue
+                
+                # Decompress chunk in real-time
+                text = streaming_decompressor.decompress_chunk(chunk)
+                
+                if text:
+                    cleaned_chunk = clean_text(text)
+                    if cleaned_chunk.strip():
+                        delta = ChoiceDelta(content=cleaned_chunk, role="assistant")
+                        choice = Choice(index=0, delta=delta, finish_reason=None)
+                        chunk_obj = ChatCompletionChunk(
+                            id=request_id,
+                            choices=[choice],
+                            created=created_time,
+                            model=model,
+                            system_fingerprint=None
+                        )
+                        yield chunk_obj
+            
+            # Finalize decompression and get any remaining text
+            final_text = streaming_decompressor.finalize()
+            if final_text.strip():
+                cleaned_final = clean_text(final_text)
+                delta = ChoiceDelta(content=cleaned_final, role="assistant")
+                choice = Choice(index=0, delta=delta, finish_reason=None)
+                chunk_obj = ChatCompletionChunk(
+                    id=request_id,
+                    choices=[choice],
+                    created=created_time,
+                    model=model,
+                    system_fingerprint=None
+                )
+                yield chunk_obj
+            
+            # Send final chunk with finish_reason="stop"
+            delta = ChoiceDelta(content="", role="assistant")
+            choice = Choice(index=0, delta=delta, finish_reason="stop")
+            final_chunk = ChatCompletionChunk(
                 id=request_id,
                 choices=[choice],
                 created=created_time,
                 model=model,
                 system_fingerprint=None
             )
-            yield chunk_obj
+            yield final_chunk
+            
+        except Exception as e:
+            # Handle errors gracefully by yielding an error chunk
+            error_delta = ChoiceDelta(content=f"Error: {str(e)}", role="assistant")
+            error_choice = Choice(index=0, delta=error_delta, finish_reason="stop")
+            error_chunk = ChatCompletionChunk(
+                id=request_id,
+                choices=[error_choice],
+                created=created_time,
+                model=model,
+                system_fingerprint=None
+            )
+            yield error_chunk
+        finally:
+            # Restore original proxies
+            self._client.session.proxies = original_proxies
 
 
 class Chat(BaseChat):
@@ -381,12 +669,7 @@ class BLACKBOXAI(OpenAICompatibleProvider):
     default_model = "GPT-4.1"
     default_vision_model = default_model
     api_endpoint = "https://www.blackbox.ai/api/chat"
-    timeout = 30
-
-
-    # Default model (remains the same as per original class)
-    default_model = "GPT-4.1"
-    default_vision_model = default_model
+    timeout = None
 
     # New OpenRouter models list
     openrouter_models = [
@@ -437,6 +720,7 @@ class BLACKBOXAI(OpenAICompatibleProvider):
 
     # New base models list
     models = [
+        "gpt-4.1-mini", # Added new model
         default_model,
         "o3-mini",
         "gpt-4.1-nano",
@@ -525,6 +809,7 @@ class BLACKBOXAI(OpenAICompatibleProvider):
         default_model: {'mode': True, 'id': "openai/gpt-4.1", 'name': default_model}, # Assuming GPT-4.1 is agent-compatible
         'o3-mini': {'mode': True, 'id': "o3-mini", 'name': "o3-mini"}, # Assuming o3-mini is agent-compatible
         'gpt-4.1-nano': {'mode': True, 'id': "gpt-4.1-nano", 'name': "gpt-4.1-nano"}, # Assuming gpt-4.1-nano is agent-compatible
+        'gpt-4.1-mini': {'mode': True, 'id': "gpt-4.1-mini", 'name': "gpt-4.1-mini"}, # Added agent mode for gpt-4.1-mini
     }
 
     # Trending agent modes
@@ -565,17 +850,15 @@ class BLACKBOXAI(OpenAICompatibleProvider):
 
 
     def __init__(
-        self,
-        proxies: dict = {}
+        self
     ):
         """
         Initialize the BlackboxAI provider with OpenAI compatibility.
-
-        Args:
-            proxies: Optional proxy configuration
         """
         # Initialize session
         self.session = requests.Session()
+        # Remove any proxy configuration to avoid connection issues
+        self.session.proxies = {}
 
         # Set headers based on GitHub reference
         self.headers = {
@@ -602,9 +885,6 @@ class BLACKBOXAI(OpenAICompatibleProvider):
             '__cf_bm': self.generate_id(32),
         }
 
-        # Set proxies if provided
-        self.session.proxies = proxies
-
         # Initialize chat interface with completions
         self.chat = Chat(self)
 
@@ -764,3 +1044,6 @@ if __name__ == "__main__":
     )
     for chunk in response:
         print(chunk.choices[0].delta.content, end='', flush=True)
+    print()
+    print("Proxies on instance:", client.proxies)
+    print("Proxies on session:", client.session.proxies)