camel-ai 0.2.75a5__py3-none-any.whl → 0.2.76a0__py3-none-any.whl

camel/toolkits/hybrid_browser_toolkit/ws_wrapper.py

@@ -13,6 +13,7 @@
 # ========= Copyright 2023-2024 @ CAMEL-AI.org. All Rights Reserved. =========
 
 import asyncio
+import contextlib
 import datetime
 import json
 import os
@@ -115,24 +116,34 @@ class WebSocketBrowserWrapper:
         self._pending_responses: Dict[
             str, asyncio.Future[Dict[str, Any]]
         ] = {}  # Message ID -> Future
+        self._server_ready_future = None  # Future to track server ready state
 
         # Logging configuration
         self.browser_log_to_file = (config or {}).get(
             'browser_log_to_file', False
         )
+        self.log_dir = (config or {}).get('log_dir', 'browser_log')
         self.session_id = (config or {}).get('session_id', 'default')
         self.log_file_path: Optional[str] = None
         self.log_buffer: List[Dict[str, Any]] = []
+        self.ts_log_file_path: Optional[str] = None
+        self.ts_log_file = None  # File handle for TypeScript logs
+        self._log_reader_task = None  # Task for reading and logging stdout
 
-        # Set up log file if needed
+        # Set up log files if needed
         if self.browser_log_to_file:
-            log_dir = "browser_log"
+            log_dir = self.log_dir if self.log_dir else "browser_log"
             os.makedirs(log_dir, exist_ok=True)
             timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
             self.log_file_path = os.path.join(
                 log_dir,
                 f"hybrid_browser_toolkit_ws_{timestamp}_{self.session_id}.log",
             )
+            # Add TypeScript console log file
+            self.ts_log_file_path = os.path.join(
+                log_dir,
+                f"typescript_console_{timestamp}_{self.session_id}.log",
+            )
 
     async def __aenter__(self):
         """Async context manager entry."""
@@ -207,39 +218,67 @@ class WebSocketBrowserWrapper:
             ['node', 'websocket-server.js'],
             cwd=self.ts_dir,
             stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
+            stderr=subprocess.STDOUT,  # Redirect stderr to stdout
             text=True,
+            bufsize=1,  # Line buffered
         )
 
+        # Create a future to wait for server ready (before starting log reader)
+        self._server_ready_future = asyncio.get_running_loop().create_future()
+
+        # Start log reader task immediately after process starts
+        self._log_reader_task = asyncio.create_task(
+            self._read_and_log_output()
+        )
+
+        if self.browser_log_to_file and self.ts_log_file_path:
+            logger.info(
+                f"TypeScript console logs will be written to: "
+                f"{self.ts_log_file_path}"
+            )
+
         # Wait for server to output the port
         server_ready = False
         timeout = 10  # 10 seconds timeout
-        start_time = time.time()
-
-        while not server_ready and time.time() - start_time < timeout:
-            if self.process.poll() is not None:
-                # Process died
-                stderr = self.process.stderr.read()
-                raise RuntimeError(
-                    f"WebSocket server failed to start: {stderr}"
-                )
 
-            try:
-                line = self.process.stdout.readline()
-                if line.startswith('SERVER_READY:'):
-                    self.server_port = int(line.split(':')[1].strip())
-                    server_ready = True
-                    logger.info(
-                        f"WebSocket server ready on port {self.server_port}"
-                    )
-            except (ValueError, IndexError):
-                continue
+        # Wait for the server to be ready
+        try:
+            await asyncio.wait_for(self._server_ready_future, timeout=timeout)
+            server_ready = True
+        except asyncio.TimeoutError:
+            server_ready = False
 
         if not server_ready:
-            self.process.kill()
-            raise RuntimeError(
-                "WebSocket server failed to start within timeout"
-            )
+            with contextlib.suppress(ProcessLookupError, Exception):
+                self.process.kill()
+            with contextlib.suppress(Exception):
+                # Ensure the process fully exits
+                self.process.wait(timeout=2)
+            # Cancel and await the log reader task
+            if self._log_reader_task and not self._log_reader_task.done():
+                self._log_reader_task.cancel()
+                with contextlib.suppress(asyncio.CancelledError):
+                    await self._log_reader_task
+            # Close TS log file if open
+            if getattr(self, 'ts_log_file', None):
+                with contextlib.suppress(Exception):
+                    self.ts_log_file.close()
+                self.ts_log_file = None
+            self.process = None
+
+            error_msg = "WebSocket server failed to start within timeout"
+            import psutil
+
+            mem = psutil.virtual_memory()
+            if mem.available < 1024**3:  # Less than 1GB available
+                error_msg = (
+                    f"WebSocket server failed to start"
+                    f"(likely due to insufficient memory). "
+                    f"Available memory: {mem.available / 1024**3:.2f}GB "
+                    f"({mem.percent}% used)"
+                )
+
+            raise RuntimeError(error_msg)
 
         # Connect to the WebSocket server
         try:
@@ -251,10 +290,34 @@ class WebSocketBrowserWrapper:
             )
             logger.info("Connected to WebSocket server")
         except Exception as e:
-            self.process.kill()
-            raise RuntimeError(
-                f"Failed to connect to WebSocket server: {e}"
-            ) from e
+            with contextlib.suppress(ProcessLookupError, Exception):
+                self.process.kill()
+            with contextlib.suppress(Exception):
+                self.process.wait(timeout=2)
+            if self._log_reader_task and not self._log_reader_task.done():
+                self._log_reader_task.cancel()
+                with contextlib.suppress(asyncio.CancelledError):
+                    await self._log_reader_task
+            if getattr(self, 'ts_log_file', None):
+                with contextlib.suppress(Exception):
+                    self.ts_log_file.close()
+                self.ts_log_file = None
+            self.process = None
+
+            error_msg = f"Failed to connect to WebSocket server: {e}"
+            import psutil
+
+            mem = psutil.virtual_memory()
+            if mem.available < 1024**3:  # Less than 1GB available
+                error_msg = (
+                    f"Failed to connect to WebSocket server"
+                    f"(likely due to insufficient memory). "
+                    f"Available memory: {mem.available / 1024**3:.2f}GB"
+                    f"({mem.percent}% used). "
+                    f"Original error: {e}"
+                )
+
+            raise RuntimeError(error_msg) from e
 
         # Start the background receiver task
         self._receive_task = asyncio.create_task(self._receive_loop())
@@ -264,34 +327,59 @@ class WebSocketBrowserWrapper:
 
     async def stop(self):
         """Stop the WebSocket connection and server."""
-        # Cancel the receiver task
-        if self._receive_task and not self._receive_task.done():
-            self._receive_task.cancel()
-            try:
-                await self._receive_task
-            except asyncio.CancelledError:
-                pass
-
+        # First, send shutdown command while receive task is still running
         if self.websocket:
-            try:
-                await self._send_command('shutdown', {})
+            with contextlib.suppress(asyncio.TimeoutError, Exception):
+                # Send shutdown command with a short timeout
+                await asyncio.wait_for(
+                    self._send_command('shutdown', {}),
+                    timeout=2.0,  # 2 second timeout for shutdown
+                )
+                # Note: TimeoutError is expected as server may close
+                # before responding
+
+            # Close websocket connection
+            with contextlib.suppress(Exception):
                 await self.websocket.close()
-            except Exception as e:
-                logger.warning(f"Error during websocket shutdown: {e}")
-            finally:
-                self.websocket = None
+            self.websocket = None
 
+        # Gracefully stop the Node process before cancelling the log reader
         if self.process:
             try:
-                self.process.terminate()
-                self.process.wait(timeout=5)
+                # give the process a short grace period to exit after shutdown
+                self.process.wait(timeout=2)
             except subprocess.TimeoutExpired:
-                self.process.kill()
-                self.process.wait()
+                try:
+                    self.process.terminate()
+                    self.process.wait(timeout=3)
+                except subprocess.TimeoutExpired:
+                    with contextlib.suppress(ProcessLookupError, Exception):
+                        self.process.kill()
+                        self.process.wait()
+                except Exception as e:
+                    logger.warning(f"Error terminating process: {e}")
             except Exception as e:
-                logger.warning(f"Error terminating process: {e}")
-            finally:
-                self.process = None
+                logger.warning(f"Error waiting for process: {e}")
+
+        # Now cancel background tasks (reader won't block on readline)
+        tasks_to_cancel = [
+            ('_receive_task', self._receive_task),
+            ('_log_reader_task', self._log_reader_task),
+        ]
+        for _, task in tasks_to_cancel:
+            if task and not task.done():
+                task.cancel()
+                with contextlib.suppress(asyncio.CancelledError):
+                    await task
+
+        # Close TS log file if open
+        if getattr(self, 'ts_log_file', None):
+            with contextlib.suppress(Exception):
+                self.ts_log_file.close()
+            self.ts_log_file = None
+
+        # Ensure process handle cleared
+        self.process = None
 
     async def _log_action(
         self,
@@ -379,16 +467,42 @@ class WebSocketBrowserWrapper:
     async def _ensure_connection(self) -> None:
         """Ensure WebSocket connection is alive."""
         if not self.websocket:
-            raise RuntimeError("WebSocket not connected")
+            error_msg = "WebSocket not connected"
+            import psutil
+
+            mem = psutil.virtual_memory()
+            if mem.available < 1024**3:  # Less than 1GB available
+                error_msg = (
+                    f"WebSocket not connected "
+                    f"(likely due to insufficient memory). "
+                    f"Available memory: {mem.available / 1024**3:.2f}GB "
+                    f"({mem.percent}% used)"
+                )
+
+            raise RuntimeError(error_msg)
 
         # Check if connection is still alive
         try:
-            # Send a ping to check connection
-            await self.websocket.ping()
+            # Send a ping and wait for the corresponding pong (bounded wait)
+            pong_waiter = await self.websocket.ping()
+            await asyncio.wait_for(pong_waiter, timeout=5.0)
         except Exception as e:
             logger.warning(f"WebSocket ping failed: {e}")
             self.websocket = None
-            raise RuntimeError("WebSocket connection lost")
+
+            error_msg = "WebSocket connection lost"
+            import psutil
+
+            mem = psutil.virtual_memory()
+            if mem.available < 1024**3:  # Less than 1GB available
+                error_msg = (
+                    f"WebSocket connection lost "
+                    f"(likely due to insufficient memory). "
+                    f"Available memory: {mem.available / 1024**3:.2f}GB "
+                    f"({mem.percent}% used)"
+                )
+
+            raise RuntimeError(error_msg)
 
     async def _send_command(
         self, command: str, params: Dict[str, Any]
@@ -403,7 +517,8 @@ class WebSocketBrowserWrapper:
         message = {'id': message_id, 'command': command, 'params': params}
 
         # Create a future for this message
-        future: asyncio.Future[Dict[str, Any]] = asyncio.Future()
+        loop = asyncio.get_running_loop()
+        future: asyncio.Future[Dict[str, Any]] = loop.create_future()
         self._pending_responses[message_id] = future
 
         try:
@@ -507,9 +622,14 @@ class WebSocketBrowserWrapper:
         return ToolResult(text=response['text'], images=response['images'])
 
     def _ensure_ref_prefix(self, ref: str) -> str:
-        """Ensure ref has 'e' prefix."""
-        if ref and not ref.startswith('e'):
+        """Ensure ref has proper prefix"""
+        if not ref:
+            return ref
+
+        # If ref is purely numeric, add 'e' prefix for main frame
+        if ref.isdigit():
             return f'e{ref}'
+
         return ref
 
     def _process_refs_in_params(
@@ -676,3 +796,73 @@ class WebSocketBrowserWrapper:
             'wait_user', {'timeout': timeout_sec}
         )
         return response
+
+    async def _read_and_log_output(self):
+        """Read stdout from Node.js process & handle SERVER_READY + logging."""
+        if not self.process:
+            return
+
+        try:
+            with contextlib.ExitStack() as stack:
+                if self.ts_log_file_path:
+                    self.ts_log_file = stack.enter_context(
+                        open(self.ts_log_file_path, 'w', encoding='utf-8')
+                    )
+                    self.ts_log_file.write(
+                        f"TypeScript Console Log - Started at "
+                        f"{time.strftime('%Y-%m-%d %H:%M:%S')}\n"
+                    )
+                    self.ts_log_file.write("=" * 80 + "\n")
+                    self.ts_log_file.flush()
+
+                while self.process and self.process.poll() is None:
+                    try:
+                        line = (
+                            await asyncio.get_running_loop().run_in_executor(
+                                None, self.process.stdout.readline
+                            )
+                        )
+                        if not line:  # EOF
+                            break
+
+                        # Check for SERVER_READY message
+                        if line.startswith('SERVER_READY:'):
+                            try:
+                                self.server_port = int(
+                                    line.split(':', 1)[1].strip()
+                                )
+                                logger.info(
+                                    f"WebSocket server ready on port "
+                                    f"{self.server_port}"
+                                )
+                                if (
+                                    self._server_ready_future
+                                    and not self._server_ready_future.done()
+                                ):
+                                    self._server_ready_future.set_result(True)
+                            except (ValueError, IndexError) as e:
+                                logger.error(
+                                    f"Failed to parse SERVER_READY: {e}"
+                                )
+
+                        # Write all output to log file
+                        if self.ts_log_file:
+                            timestamp = time.strftime('%H:%M:%S')
+                            self.ts_log_file.write(f"[{timestamp}] {line}")
+                            self.ts_log_file.flush()
+
+                    except Exception as e:
+                        logger.warning(f"Error reading stdout: {e}")
+                        break
+
+                # Footer if we had a file
+                if self.ts_log_file:
+                    self.ts_log_file.write("\n" + "=" * 80 + "\n")
+                    self.ts_log_file.write(
+                        f"TypeScript Console Log - Ended at "
+                        f"{time.strftime('%Y-%m-%d %H:%M:%S')}\n"
+                    )
+                # ExitStack closes file; clear handle
+                self.ts_log_file = None
+        except Exception as e:
+            logger.warning(f"Error in _read_and_log_output: {e}")

camel/toolkits/hybrid_browser_toolkit_py/hybrid_browser_toolkit.py

@@ -95,6 +95,7 @@ class HybridBrowserToolkit(BaseToolkit, RegisteredAgentToolkit):
         cache_dir: str = "tmp/",
         enabled_tools: Optional[List[str]] = None,
         browser_log_to_file: bool = False,
+        log_dir: Optional[str] = None,
         session_id: Optional[str] = None,
         default_start_url: str = "https://google.com/",
         default_timeout: Optional[int] = None,
@@ -144,6 +145,8 @@ class HybridBrowserToolkit(BaseToolkit, RegisteredAgentToolkit):
                 and page loading times.
                 Logs are saved to an auto-generated timestamped file.
                 Defaults to `False`.
+            log_dir (Optional[str]): Custom directory path for log files.
+                If None, defaults to "browser_log". Defaults to `None`.
             session_id (Optional[str]): A unique identifier for this browser
                 session. When multiple HybridBrowserToolkit instances are
                 used
@@ -201,6 +204,7 @@ class HybridBrowserToolkit(BaseToolkit, RegisteredAgentToolkit):
         self._web_agent_model = web_agent_model
         self._cache_dir = cache_dir
         self._browser_log_to_file = browser_log_to_file
+        self._log_dir = log_dir
         self._default_start_url = default_start_url
         self._session_id = session_id or "default"
         self._viewport_limit = viewport_limit
@@ -237,7 +241,7 @@ class HybridBrowserToolkit(BaseToolkit, RegisteredAgentToolkit):
         # Set up log file if needed
         if self.log_to_file:
             # Create log directory if it doesn't exist
-            log_dir = "browser_log"
+            log_dir = self._log_dir if self._log_dir else "browser_log"
             os.makedirs(log_dir, exist_ok=True)
 
             timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")

camel/toolkits/{openai_image_toolkit.py → image_generation_toolkit.py}

@@ -15,7 +15,7 @@
 import base64
 import os
 from io import BytesIO
-from typing import List, Literal, Optional, Union
+from typing import ClassVar, List, Literal, Optional, Tuple, Union
 
 from openai import OpenAI
 from PIL import Image
@@ -29,21 +29,32 @@ logger = get_logger(__name__)
 
 
 @MCPServer()
-class OpenAIImageToolkit(BaseToolkit):
-    r"""A class toolkit for image generation using OpenAI's
-    Image Generation API.
-    """
-
-    @api_keys_required(
-        [
-            ("api_key", "OPENAI_API_KEY"),
-        ]
-    )
+class ImageGenToolkit(BaseToolkit):
+    r"""A class toolkit for image generation using Grok and OpenAI models."""
+
+    GROK_MODELS: ClassVar[List[str]] = [
+        "grok-2-image",
+        "grok-2-image-latest",
+        "grok-2-image-1212",
+    ]
+    OPENAI_MODELS: ClassVar[List[str]] = [
+        "gpt-image-1",
+        "dall-e-3",
+        "dall-e-2",
+    ]
+
     def __init__(
         self,
         model: Optional[
-            Literal["gpt-image-1", "dall-e-3", "dall-e-2"]
-        ] = "gpt-image-1",
+            Literal[
+                "gpt-image-1",
+                "dall-e-3",
+                "dall-e-2",
+                "grok-2-image",
+                "grok-2-image-latest",
+                "grok-2-image-1212",
+            ]
+        ] = "dall-e-3",
         timeout: Optional[float] = None,
         api_key: Optional[str] = None,
         url: Optional[str] = None,
@@ -72,12 +83,12 @@ class OpenAIImageToolkit(BaseToolkit):
         # NOTE: Some arguments are set in the constructor to prevent the agent
         # from making invalid API calls with model-specific parameters. For
         # example, the 'style' argument is only supported by 'dall-e-3'.
-        r"""Initializes a new instance of the OpenAIImageToolkit class.
+        r"""Initializes a new instance of the ImageGenToolkit class.
 
         Args:
             api_key (Optional[str]): The API key for authenticating
-                with the OpenAI service. (default: :obj:`None`)
-            url (Optional[str]): The url to the OpenAI service.
+                with the image model service. (default: :obj:`None`)
+            url (Optional[str]): The url to the image model service.
                 (default: :obj:`None`)
             model (Optional[str]): The model to use.
                 (default: :obj:`"dall-e-3"`)
@@ -103,9 +114,23 @@ class OpenAIImageToolkit(BaseToolkit):
                 image.(default: :obj:`"image_save"`)
         """
         super().__init__(timeout=timeout)
-        api_key = api_key or os.environ.get("OPENAI_API_KEY")
-        url = url or os.environ.get("OPENAI_API_BASE_URL")
-        self.client = OpenAI(api_key=api_key, base_url=url)
+        if model not in self.GROK_MODELS + self.OPENAI_MODELS:
+            available_models = sorted(self.OPENAI_MODELS + self.GROK_MODELS)
+            raise ValueError(
+                f"Unsupported model: {model}. "
+                f"Supported models are: {available_models}"
+            )
+
+        # Set default url for Grok models
+        url = "https://api.x.ai/v1" if model in self.GROK_MODELS else url
+
+        api_key, base_url = (
+            self.get_openai_credentials(url, api_key)
+            if model in self.OPENAI_MODELS
+            else self.get_grok_credentials(url, api_key)
+        )
+
+        self.client = OpenAI(api_key=api_key, base_url=base_url)
         self.model = model
         self.size = size
         self.quality = quality
@@ -139,7 +164,7 @@ class OpenAIImageToolkit(BaseToolkit):
             return None
 
     def _build_base_params(self, prompt: str, n: Optional[int] = None) -> dict:
-        r"""Build base parameters dict for OpenAI API calls.
+        r"""Build base parameters dict for Image Model API calls.
 
         Args:
             prompt (str): The text prompt for the image operation.
@@ -153,6 +178,10 @@ class OpenAIImageToolkit(BaseToolkit):
         # basic parameters supported by all models
         if n is not None:
             params["n"] = n  # type: ignore[assignment]
+
+        if self.model in self.GROK_MODELS:
+            return params
+
         if self.size is not None:
             params["size"] = self.size
 
@@ -179,16 +208,18 @@ class OpenAIImageToolkit(BaseToolkit):
                 params["quality"] = self.quality
             if self.background is not None:
                 params["background"] = self.background
-
         return params
 
     def _handle_api_response(
-        self, response, image_name: Union[str, List[str]], operation: str
+        self,
+        response,
+        image_name: Union[str, List[str]],
+        operation: str,
     ) -> str:
-        r"""Handle API response from OpenAI image operations.
+        r"""Handle API response from image operations.
 
         Args:
-            response: The response object from OpenAI API.
+            response: The response object from image model API.
             image_name (Union[str, List[str]]): Name(s) for the saved image
                 file(s). If str, the same name is used for all images (will
                 cause error for multiple images). If list, must have exactly
@@ -198,8 +229,9 @@ class OpenAIImageToolkit(BaseToolkit):
         Returns:
             str: Success message with image path/URL or error message.
         """
+        source = "Grok" if self.model in self.GROK_MODELS else "OpenAI"
         if response.data is None or len(response.data) == 0:
-            error_msg = "No image data returned from OpenAI API."
+            error_msg = f"No image data returned from {source} API."
             logger.error(error_msg)
             return error_msg
 
@@ -283,7 +315,7 @@ class OpenAIImageToolkit(BaseToolkit):
         image_name: Union[str, List[str]] = "image.png",
         n: int = 1,
     ) -> str:
-        r"""Generate an image using OpenAI's Image Generation models.
+        r"""Generate an image using image models.
         The generated image will be saved locally (for ``b64_json`` response
         formats) or an image URL will be returned (for ``url`` response
         formats).
@@ -309,15 +341,50 @@ class OpenAIImageToolkit(BaseToolkit):
             logger.error(error_msg)
             return error_msg
 
+    @api_keys_required([("api_key", "XAI_API_KEY")])
+    def get_grok_credentials(self, url, api_key) -> Tuple[str, str]:  # type: ignore[return-value]
+        r"""Get API credentials for the specified Grok model.
+
+        Args:
+            url (str): The base URL for the Grok API.
+            api_key (str): The API key for the Grok API.
+
+        Returns:
+            tuple: (api_key, base_url)
+        """
+
+        # Get credentials based on model type
+        api_key = api_key or os.getenv("XAI_API_KEY")
+        return api_key, url
+
+    @api_keys_required([("api_key", "OPENAI_API_KEY")])
+    def get_openai_credentials(self, url, api_key) -> Tuple[str, str | None]:  # type: ignore[return-value]
+        r"""Get API credentials for the specified OpenAI model.
+
+        Args:
+            url (str): The base URL for the OpenAI API.
+            api_key (str): The API key for the OpenAI API.
+
+        Returns:
+            Tuple[str, str | None]: (api_key, base_url)
+        """
+
+        api_key = api_key or os.getenv("OPENAI_API_KEY")
+        base_url = url or os.getenv("OPENAI_API_BASE_URL")
+        return api_key, base_url
+
     def get_tools(self) -> List[FunctionTool]:
-        r"""Returns a list of FunctionTool objects representing the
-        functions in the toolkit.
+        r"""Returns a list of FunctionTool objects representing the functions
+            in the toolkit.
 
         Returns:
-            List[FunctionTool]: A list of FunctionTool objects
-                representing the functions in the toolkit.
+            List[FunctionTool]: A list of FunctionTool objects representing the
+                functions in the toolkit.
         """
         return [
             FunctionTool(self.generate_image),
-            # could add edit_image function later
         ]
+
+
+# Backward compatibility alias
+OpenAIImageToolkit = ImageGenToolkit