webscout 8.2.8__py3-none-any.whl → 8.3__py3-none-any.whl

webscout/Provider/OPENAI/textpollinations.py

@@ -8,7 +8,7 @@ from typing import List, Dict, Optional, Union, Generator, Any
 from .base import OpenAICompatibleProvider, BaseChat, BaseCompletions
 from .utils import (
     ChatCompletionChunk, ChatCompletion, Choice, ChoiceDelta,
-    ChatCompletionMessage, CompletionUsage, ToolCall, ToolFunction
+    ChatCompletionMessage, CompletionUsage, ToolCall, ToolFunction, count_tokens
 )
 
 # Import LitAgent for browser fingerprinting
@@ -34,6 +34,8 @@ class Completions(BaseCompletions):
         top_p: Optional[float] = None,
         tools: Optional[List[Dict[str, Any]]] = None,
         tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
+        timeout: Optional[int] = None,
+        proxies: Optional[Dict[str, str]] = None,
         **kwargs: Any
     ) -> Union[ChatCompletion, Generator[ChatCompletionChunk, None, None]]:
         """
@@ -62,16 +64,18 @@ class Completions(BaseCompletions):
         created_time = int(time.time())
 
         if stream:
-            return self._create_streaming(request_id, created_time, model, payload)
+            return self._create_streaming(request_id, created_time, model, payload, timeout, proxies)
         else:
-            return self._create_non_streaming(request_id, created_time, model, payload)
+            return self._create_non_streaming(request_id, created_time, model, payload, timeout, proxies)
 
     def _create_streaming(
         self,
         request_id: str,
         created_time: int,
         model: str,
-        payload: Dict[str, Any]
+        payload: Dict[str, Any],
+        timeout: Optional[int] = None,
+        proxies: Optional[Dict[str, str]] = None
     ) -> Generator[ChatCompletionChunk, None, None]:
         """Implementation for streaming chat completions."""
         try:
@@ -82,7 +86,8 @@ 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)
             )
 
             if not response.ok:
@@ -163,7 +168,9 @@ class Completions(BaseCompletions):
         request_id: str,
         created_time: int,
         model: str,
-        payload: Dict[str, Any]
+        payload: Dict[str, Any],
+        timeout: Optional[int] = None,
+        proxies: Optional[Dict[str, str]] = None
     ) -> ChatCompletion:
         """Implementation for non-streaming chat completions."""
         try:
@@ -173,7 +180,8 @@ class Completions(BaseCompletions):
                 self._client.api_endpoint,
                 headers=self._client.headers,
                 json=payload,
-                timeout=self._client.timeout
+                timeout=timeout or self._client.timeout,
+                proxies=proxies or getattr(self._client, "proxies", None)
             )
 
             if not response.ok:
@@ -226,9 +234,9 @@ class Completions(BaseCompletions):
                 finish_reason="stop"
             )
 
-            # Estimate token usage (very rough estimate)
-            prompt_tokens = sum(len(msg.get("content", "")) // 4 for msg in payload.get("messages", []))
-            completion_tokens = len(full_content) // 4
+            # Estimate token usage using count_tokens
+            prompt_tokens = count_tokens([msg.get("content", "") for msg in payload.get("messages", [])])
+            completion_tokens = count_tokens(full_content)
             usage = CompletionUsage(
                 prompt_tokens=prompt_tokens,
                 completion_tokens=completion_tokens,
@@ -336,4 +344,3 @@ class TextPollinations(OpenAICompatibleProvider):
             def list(inner_self):
                 return type(self).AVAILABLE_MODELS
         return _ModelList()
-