webscout 8.2.2__py3-none-any.whl → 2026.1.19__py3-none-any.whl

webscout/Provider/OPENAI/heckai.py

@@ -1,27 +1,36 @@
 import time
 import uuid
+from typing import Any, Dict, Generator, List, Optional, Union, cast
+
 import requests
-from typing import List, Dict, Optional, Union, Generator, Any
 
-from webscout.litagent import LitAgent
-from .base import BaseChat, BaseCompletions, OpenAICompatibleProvider
-from .utils import (
+from webscout.Provider.OPENAI.base import (
+    BaseChat,
+    BaseCompletions,
+    OpenAICompatibleProvider,
+    SimpleModelList,
+)
+from webscout.Provider.OPENAI.utils import (
     ChatCompletion,
     ChatCompletionChunk,
-    Choice,
     ChatCompletionMessage,
+    Choice,
     ChoiceDelta,
     CompletionUsage,
-    format_prompt
+    count_tokens,
+    format_prompt,
 )
 
+from ...litagent import LitAgent
+
 # ANSI escape codes for formatting
 BOLD = "\033[1m"
 RED = "\033[91m"
 RESET = "\033[0m"
 
+
 class Completions(BaseCompletions):
-    def __init__(self, client: 'HeckAI'):
+    def __init__(self, client: "HeckAI"):
         self._client = client
 
     def create(
@@ -33,7 +42,9 @@ class Completions(BaseCompletions):
         stream: bool = False,
         temperature: Optional[float] = None,  # Not used by HeckAI but kept for compatibility
         top_p: Optional[float] = None,  # Not used by HeckAI but kept for compatibility
-        **kwargs: Any  # Not used by HeckAI but kept for compatibility
+        timeout: Optional[int] = None,
+        proxies: Optional[Dict[str, str]] = None,
+        **kwargs: Any,  # Not used by HeckAI but kept for compatibility
     ) -> Union[ChatCompletion, Generator[ChatCompletionChunk, None, None]]:
         """
         Creates a model response for the given chat conversation.
@@ -45,26 +56,36 @@ class Completions(BaseCompletions):
         question = format_prompt(messages, add_special_tokens=True)
 
         # Prepare the payload for HeckAI API
+        model = self._client.convert_model_name(model)
         payload = {
             "model": model,
             "question": question,
             "language": self._client.language,
             "sessionId": self._client.session_id,
-            "previousQuestion": None,  # Not needed when using format_prompt
-            "previousAnswer": None,    # Not needed when using format_prompt
-            "imgUrls": []
+            "previousQuestion": None,
+            "previousAnswer": None,
+            "imgUrls": [],
+            "superSmartMode": False,
         }
 
         request_id = f"chatcmpl-{uuid.uuid4()}"
         created_time = int(time.time())
 
         if stream:
-            return self._create_stream(request_id, created_time, model, payload)
+            return self._create_stream(request_id, created_time, model, payload, timeout, proxies)
         else:
-            return self._create_non_stream(request_id, created_time, model, payload)
+            return self._create_non_stream(
+                request_id, created_time, model, payload, timeout, proxies
+            )
 
     def _create_stream(
-        self, request_id: str, created_time: int, model: str, payload: Dict[str, Any]
+        self,
+        request_id: str,
+        created_time: int,
+        model: str,
+        payload: Dict[str, Any],
+        timeout: Optional[int] = None,
+        proxies: Optional[Dict[str, str]] = None,
     ) -> Generator[ChatCompletionChunk, None, None]:
         try:
             response = self._client.session.post(
@@ -72,151 +93,113 @@ class Completions(BaseCompletions):
                 headers=self._client.headers,
                 json=payload,
                 stream=True,
-                timeout=self._client.timeout
+                timeout=timeout or self._client.timeout,
+                proxies=proxies or getattr(self._client, "proxies", None),
             )
             response.raise_for_status()
 
-            # Track token usage across chunks
-            completion_tokens = 0
-
-            streaming_text = ""
+            streaming_text = []
             in_answer = False
 
             for line in response.iter_lines(decode_unicode=True):
                 if not line:
                     continue
-
-                # Remove "data: " prefix
                 if line.startswith("data: "):
                     data = line[6:]
                 else:
                     continue
-
-                # Check for control markers
                 if data == "[ANSWER_START]":
                     in_answer = True
                     continue
-
                 if data == "[ANSWER_DONE]":
                     in_answer = False
                     continue
-
-                if data == "[RELATE_Q_START]" or data == "[RELATE_Q_DONE]":
+                if data.startswith("[") and data.endswith("]"):
                     continue
-
-                # Process content if we're in an answer section
                 if in_answer:
-                    streaming_text += data
-                    completion_tokens += len(data) // 4  # Rough estimate
-
-                    # Create a delta object for this chunk
-                    delta = ChoiceDelta(content=data)
+                    # Fix encoding issues (e.g., emoji) for each chunk
+                    try:
+                        data_fixed = data.encode("latin1").decode("utf-8")
+                    except (UnicodeEncodeError, UnicodeDecodeError):
+                        data_fixed = data
+                    streaming_text.append(data_fixed)
+                    delta = ChoiceDelta(content=data_fixed)
                     choice = Choice(index=0, delta=delta, finish_reason=None)
-
                     chunk = ChatCompletionChunk(
                         id=request_id,
                         choices=[choice],
                         created=created_time,
                         model=model,
                     )
-
                     yield chunk
-
-            # Store the response for future context
-            # We don't need to store previous_question/answer as we're using format_prompt
-            # which handles the conversation formatting
-
             # Final chunk with finish_reason
             delta = ChoiceDelta(content=None)
             choice = Choice(index=0, delta=delta, finish_reason="stop")
-
             chunk = ChatCompletionChunk(
                 id=request_id,
                 choices=[choice],
                 created=created_time,
                 model=model,
             )
-
             yield chunk
-
         except requests.exceptions.RequestException as e:
             print(f"{RED}Error during HeckAI stream request: {e}{RESET}")
             raise IOError(f"HeckAI request failed: {e}") from e
 
     def _create_non_stream(
-        self, request_id: str, created_time: int, model: str, payload: Dict[str, Any]
+        self,
+        request_id: str,
+        created_time: int,
+        model: str,
+        payload: Dict[str, Any],
+        timeout: Optional[int] = None,
+        proxies: Optional[Dict[str, str]] = None,
     ) -> ChatCompletion:
         try:
-            full_text = ""
-            streaming_text = ""
+            answer_lines = []
             in_answer = False
-
             response = self._client.session.post(
                 self._client.url,
                 headers=self._client.headers,
                 json=payload,
                 stream=True,
-                timeout=self._client.timeout
+                timeout=timeout or self._client.timeout,
+                proxies=proxies or getattr(self._client, "proxies", None),
             )
             response.raise_for_status()
-
             for line in response.iter_lines(decode_unicode=True):
                 if not line:
                     continue
-
-                # Remove "data: " prefix
                 if line.startswith("data: "):
                     data = line[6:]
                 else:
                     continue
-
-                # Check for control markers
                 if data == "[ANSWER_START]":
                     in_answer = True
                     continue
-
                 if data == "[ANSWER_DONE]":
                     in_answer = False
                     continue
-
-                if data == "[RELATE_Q_START]" or data == "[RELATE_Q_DONE]":
+                if data.startswith("[") and data.endswith("]"):
                     continue
-
-                # Process content if we're in an answer section
                 if in_answer:
-                    streaming_text += data
-
-            full_text = streaming_text
-
-            # Store the response for future context
-            # We don't need to store previous_question/answer as we're using format_prompt
-            # which handles the conversation formatting
-
-            # Create usage statistics (estimated)
-            prompt_tokens = len(payload["question"]) // 4
-            completion_tokens = len(full_text) // 4
+                    answer_lines.append(data)
+            full_text = " ".join(x.strip() for x in answer_lines if x.strip())
+            # Fix encoding issues (e.g., emoji)
+            try:
+                full_text = full_text.encode("latin1").decode("utf-8")
+            except (UnicodeEncodeError, UnicodeDecodeError):
+                pass
+            prompt_tokens = count_tokens(payload.get("question", ""))
+            completion_tokens = count_tokens(full_text)
             total_tokens = prompt_tokens + completion_tokens
-
             usage = CompletionUsage(
                 prompt_tokens=prompt_tokens,
                 completion_tokens=completion_tokens,
-                total_tokens=total_tokens
-            )
-
-            # Create the message object
-            message = ChatCompletionMessage(
-                role="assistant",
-                content=full_text
+                total_tokens=total_tokens,
             )
-
-            # Create the choice object
-            choice = Choice(
-                index=0,
-                message=message,
-                finish_reason="stop"
-            )
-
-            # Create the completion object
+            message = ChatCompletionMessage(role="assistant", content=full_text)
+            choice = Choice(index=0, message=message, finish_reason="stop")
             completion = ChatCompletion(
                 id=request_id,
                 choices=[choice],
@@ -224,17 +207,17 @@ class Completions(BaseCompletions):
                 model=model,
                 usage=usage,
             )
-
             return completion
-
         except Exception as e:
             print(f"{RED}Error during HeckAI non-stream request: {e}{RESET}")
             raise IOError(f"HeckAI request failed: {e}") from e
 
+
 class Chat(BaseChat):
-    def __init__(self, client: 'HeckAI'):
+    def __init__(self, client: "HeckAI"):
         self.completions = Completions(client)
 
+
 class HeckAI(OpenAICompatibleProvider):
     """
     OpenAI-compatible client for HeckAI API.
@@ -248,18 +231,20 @@ class HeckAI(OpenAICompatibleProvider):
         print(response.choices[0].message.content)
     """
 
+    required_auth = False
     AVAILABLE_MODELS = [
+        "google/gemini-2.5-flash-preview",
         "deepseek/deepseek-chat",
-        "openai/gpt-4o-mini",
         "deepseek/deepseek-r1",
-        "google/gemini-2.0-flash-001"
+        "openai/gpt-4o-mini",
+        "openai/gpt-4.1-mini",
+        "x-ai/grok-3-mini-beta",
+        "meta-llama/llama-4-scout",
+        "openai/gpt-5-mini",
+        "openai/gpt-5-nano",
     ]
 
-    def __init__(
-        self,
-        timeout: int = 30,
-        language: str = "English"
-    ):
+    def __init__(self, timeout: int = 30, language: str = "English"):
         """
         Initialize the HeckAI client.
 
@@ -275,11 +260,11 @@ class HeckAI(OpenAICompatibleProvider):
         # Use LitAgent for user-agent
         agent = LitAgent()
         self.headers = {
-            'User-Agent': agent.random(),
-            'Content-Type': 'application/json',
-            'Origin': 'https://heck.ai',
-            'Referer': 'https://heck.ai/',
-            'Connection': 'keep-alive'
+            "User-Agent": agent.random(),
+            "Content-Type": "application/json",
+            "Origin": "https://heck.ai",
+            "Referer": "https://heck.ai/",
+            "Connection": "keep-alive",
         }
 
         self.session = requests.Session()
@@ -301,9 +286,15 @@ class HeckAI(OpenAICompatibleProvider):
                 return available_model
 
         # Default to gemini if no match
-        print(f"{BOLD}Warning: Model '{model}' not found, using default model 'google/gemini-2.0-flash-001'{RESET}")
+        print(
+            f"{BOLD}Warning: Model '{model}' not found, using default model 'google/gemini-2.0-flash-001'{RESET}"
+        )
         return "google/gemini-2.0-flash-001"
 
+    @property
+    def models(self) -> SimpleModelList:
+        return SimpleModelList(type(self).AVAILABLE_MODELS)
+
 
 # Simple test if run directly
 if __name__ == "__main__":
@@ -321,10 +312,15 @@ if __name__ == "__main__":
                     {"role": "system", "content": "You are a helpful assistant."},
                     {"role": "user", "content": "Say 'Hello' in one word"},
                 ],
-                stream=False
+                stream=False,
             )
 
-            if response and response.choices and response.choices[0].message.content:
+            if (
+                isinstance(response, ChatCompletion)
+                and response.choices
+                and response.choices[0].message
+                and response.choices[0].message.content
+            ):
                 status = "✓"
                 # Truncate response if too long
                 display_text = response.choices[0].message.content.strip()