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

webscout/Provider/OPENAI/base.py

@@ -78,6 +78,8 @@ class BaseCompletions(ABC):
         top_p: Optional[float] = None,
         tools: Optional[List[Union[Tool, Dict[str, Any]]]] = None,  # Support for tool definitions
         tool_choice: Optional[Union[str, Dict[str, Any]]] = None,  # Support for tool_choice parameter
+        timeout: Optional[int] = None,
+        proxies: Optional[dict] = None,
         **kwargs: Any
     ) -> Union[ChatCompletion, Generator[ChatCompletionChunk, None, None]]:
         """

webscout/Provider/OPENAI/c4ai.py

@@ -10,7 +10,7 @@ from .base import OpenAICompatibleProvider, BaseChat, BaseCompletions
 from .utils import (
     ChatCompletionChunk, ChatCompletion, Choice, ChoiceDelta,
     ChatCompletionMessage, CompletionUsage,
-    get_system_prompt, get_last_user_message, format_prompt # Import format_prompt
+    get_system_prompt, get_last_user_message, format_prompt, count_tokens # Import format_prompt
 )
 
 # Attempt to import LitAgent, fallback if not available
@@ -49,6 +49,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]]:
         """
@@ -88,14 +90,21 @@ class Completions(BaseCompletions):
 
         # Pass the formatted conversation prompt
         if stream:
-            return self._create_stream(request_id, created_time, model, conversation_id, conversation_prompt, system_prompt)
+            return self._create_stream(request_id, created_time, model, conversation_id, conversation_prompt, system_prompt, timeout=timeout, proxies=proxies)
         else:
-            return self._create_non_stream(request_id, created_time, model, conversation_id, conversation_prompt, system_prompt)
+            return self._create_non_stream(request_id, created_time, model, conversation_id, conversation_prompt, system_prompt, timeout=timeout, proxies=proxies)
 
     def _create_stream(
-        self, request_id: str, created_time: int, model: str, conversation_id: str, prompt: str, system_prompt: str
+        self, request_id: str, created_time: int, model: str, conversation_id: str, prompt: str, system_prompt: str,
+        timeout: Optional[int] = None, proxies: Optional[dict] = None
     ) -> Generator[ChatCompletionChunk, None, None]:
+        original_proxies = self._client.session.proxies
+        if proxies is not None:
+            self._client.session.proxies = proxies
+        else:
+            self._client.session.proxies = {}
         try:
+            timeout_val = timeout if timeout is not None else self._client.timeout
             message_id = self._client._conversation_data[model]["messageId"]
             url = f"{self._client.url}/api/chat/message"
             payload = {
@@ -117,7 +126,8 @@ class Completions(BaseCompletions):
                 headers=self._client.headers,
                 json=payload,
                 stream=True,
-                timeout=self._client.timeout
+                timeout=timeout_val,
+                proxies=proxies or getattr(self._client, "proxies", None)
             )
             response.raise_for_status()
 
@@ -160,11 +170,20 @@ class Completions(BaseCompletions):
         except Exception as e:
             print(f"Error during C4AI stream request: {e}")
             raise IOError(f"C4AI request failed: {e}") from e
+        finally:
+            self._client.session.proxies = original_proxies
 
     def _create_non_stream(
-        self, request_id: str, created_time: int, model: str, conversation_id: str, prompt: str, system_prompt: str
+        self, request_id: str, created_time: int, model: str, conversation_id: str, prompt: str, system_prompt: str,
+        timeout: Optional[int] = None, proxies: Optional[dict] = None
     ) -> ChatCompletion:
+        original_proxies = self._client.session.proxies
+        if proxies is not None:
+            self._client.session.proxies = proxies
+        else:
+            self._client.session.proxies = {}
         try:
+            timeout_val = timeout if timeout is not None else self._client.timeout
             message_id = self._client._conversation_data[model]["messageId"]
             url = f"{self._client.url}/api/chat/message"
             payload = {
@@ -185,7 +204,8 @@ class Completions(BaseCompletions):
                 url,
                 headers=self._client.headers,
                 json=payload,
-                timeout=self._client.timeout
+                timeout=timeout_val,
+                proxies=proxies or getattr(self._client, "proxies", None)
             )
             response.raise_for_status()
 
@@ -194,8 +214,8 @@ class Completions(BaseCompletions):
             message = ChatCompletionMessage(role="assistant", content=response_text)
             choice = Choice(index=0, message=message, finish_reason="stop")
             # Estimate tokens based on the formatted prompt
-            prompt_tokens = len(prompt) // 4
-            completion_tokens = len(response_text) // 4
+            prompt_tokens = count_tokens(prompt)
+            completion_tokens = count_tokens(response_text)
             usage = CompletionUsage(
                 prompt_tokens=prompt_tokens,
                 completion_tokens=completion_tokens,
@@ -213,6 +233,8 @@ class Completions(BaseCompletions):
         except Exception as e:
             print(f"Error during C4AI non-stream request: {e}")
             raise IOError(f"C4AI request failed: {e}") from e
+        finally:
+            self._client.session.proxies = original_proxies
 
 class Chat(BaseChat):
     def __init__(self, client: 'C4AI'):
@@ -242,19 +264,18 @@ class C4AI(OpenAICompatibleProvider):
 
     def __init__(
         self,
-        timeout: Optional[int] = None,
         browser: str = "chrome"
     ):
         """
         Initialize the C4AI client.
 
         Args:
-            timeout: Request timeout in seconds.
             browser: Browser name for LitAgent to generate User-Agent.
         """
-        self.timeout = timeout
+        self.timeout = 30
         self.url = "https://cohereforai-c4ai-command.hf.space"
         self.session = requests.Session()
+        self.session.proxies = {}
         self.max_tokens_to_sample = 2000
 
         agent = LitAgent()
@@ -370,4 +391,4 @@ class C4AI(OpenAICompatibleProvider):
         class _ModelList:
             def list(inner_self):
                 return type(self).AVAILABLE_MODELS
-        return _ModelList()
+        return _ModelList()