camel-ai 0.2.59__py3-none-any.whl → 0.2.61__py3-none-any.whl

camel/toolkits/jina_reranker_toolkit.py

@@ -11,7 +11,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ========= Copyright 2023-2024 @ CAMEL-AI.org. All Rights Reserved. =========
-from typing import List, Optional, Tuple
+import json
+import os
+from typing import Any, Dict, List, Optional
+
+import requests
 
 from camel.toolkits import FunctionTool
 from camel.toolkits.base import BaseToolkit
@@ -30,7 +34,9 @@ class JinaRerankerToolkit(BaseToolkit):
     def __init__(
         self,
         timeout: Optional[float] = None,
+        model_name: Optional[str] = "jinaai/jina-reranker-m0",
         device: Optional[str] = None,
+        use_api: bool = True,
     ) -> None:
         r"""Initializes a new instance of the JinaRerankerToolkit class.
 
@@ -38,31 +44,57 @@ class JinaRerankerToolkit(BaseToolkit):
             timeout (Optional[float]): The timeout value for API requests
                 in seconds. If None, no timeout is applied.
                 (default: :obj:`None`)
+            model_name (Optional[str]): The reranker model name. If None,
+                will use the default model.
+                (default: :obj:`None`)
             device (Optional[str]): Device to load the model on. If None,
                 will use CUDA if available, otherwise CPU.
+                Only effective when use_api=False.
                 (default: :obj:`None`)
+            use_api (bool): A flag to switch between local model and API.
+                (default: :obj:`True`)
         """
-        import torch
-        from transformers import AutoModel
 
         super().__init__(timeout=timeout)
 
-        self.model = AutoModel.from_pretrained(
-            'jinaai/jina-reranker-m0',
-            torch_dtype="auto",
-            trust_remote_code=True,
-        )
-        DEVICE = (
-            device
-            if device is not None
-            else ("cuda" if torch.cuda.is_available() else "cpu")
-        )
-        self.model.to(DEVICE)
-        self.model.eval()
+        self.use_api = use_api
+        self.model_name = model_name
+
+        if self.use_api:
+            self.model = None
+            self._api_key = os.environ.get("JINA_API_KEY", "None")
+            if self._api_key == "None":
+                raise ValueError(
+                    "Missing or empty required API keys in "
+                    "environment variables\n"
+                    "You can obtain the API key from https://jina.ai/reranker/"
+                )
+            self.url = 'https://api.jina.ai/v1/rerank'
+            self.headers = {
+                'Content-Type': 'application/json',
+                'Accept': 'application/json',
+                'Authorization': f'Bearer {self._api_key}',
+            }
+        else:
+            import torch
+            from transformers import AutoModel
+
+            self.model = AutoModel.from_pretrained(
+                self.model_name,
+                torch_dtype="auto",
+                trust_remote_code=True,
+            )
+            self.device = (
+                device
+                if device is not None
+                else ("cuda" if torch.cuda.is_available() else "cpu")
+            )
+            self.model.to(self.device)
+            self.model.eval()
 
     def _sort_documents(
         self, documents: List[str], scores: List[float]
-    ) -> List[Tuple[str, float]]:
+    ) -> List[Dict[str, object]]:
         r"""Sort documents by their scores in descending order.
 
         Args:
@@ -70,7 +102,7 @@ class JinaRerankerToolkit(BaseToolkit):
             scores (List[float]): Corresponding scores for each document.
 
         Returns:
-            List[Tuple[str, float]]: Sorted list of (document, score) pairs.
+            List[Dict[str, object]]: Sorted list of (document, score) pairs.
 
         Raises:
             ValueError: If documents and scores have different lengths.
@@ -80,14 +112,45 @@ class JinaRerankerToolkit(BaseToolkit):
         doc_score_pairs = list(zip(documents, scores))
         doc_score_pairs.sort(key=lambda x: x[1], reverse=True)
 
-        return doc_score_pairs
+        results = [
+            {'document': {'text': doc}, 'relevance_score': score}
+            for doc, score in doc_score_pairs
+        ]
+
+        return results
+
+    def _call_jina_api(self, data: Dict[str, Any]) -> List[Dict[str, object]]:
+        r"""Makes a call to the JINA API for reranking.
+
+        Args:
+            data (Dict[str]): The data to be passed into the api body.
+
+        Returns:
+            List[Dict[str, object]]: A list of dictionary containing
+                the reranked documents and their relevance scores.
+        """
+        try:
+            response = requests.post(
+                self.url,
+                headers=self.headers,
+                data=json.dumps(data),
+                timeout=self.timeout,
+            )
+            response.raise_for_status()
+            results = [
+                {key: value for key, value in _res.items() if key != 'index'}
+                for _res in response.json()['results']
+            ]
+            return results
+        except requests.exceptions.RequestException as e:
+            raise RuntimeError(f"Failed to get response from Jina AI: {e}")
 
     def rerank_text_documents(
         self,
         query: str,
         documents: List[str],
         max_length: int = 1024,
-    ) -> List[Tuple[str, float]]:
+    ) -> List[Dict[str, object]]:
         r"""Reranks text documents based on their relevance to a text query.
 
         Args:
@@ -97,21 +160,34 @@ class JinaRerankerToolkit(BaseToolkit):
                 (default: :obj:`1024`)
 
         Returns:
-            List[Tuple[str, float]]: A list of tuples containing
+            List[Dict[str, object]]: A list of dictionary containing
                 the reranked documents and their relevance scores.
         """
-        import torch
 
-        if self.model is None:
-            raise ValueError(
-                "Model has not been initialized or failed to initialize."
-            )
+        data = {
+            'model': self.model_name,
+            'query': query,
+            'top_n': len(documents),
+            'documents': documents,
+            'return_documents': True,
+        }
 
-        with torch.inference_mode():
-            text_pairs = [[query, doc] for doc in documents]
-            scores = self.model.compute_score(
-                text_pairs, max_length=max_length, doc_type="text"
-            )
+        if self.use_api:
+            return self._call_jina_api(data)
+
+        else:
+            import torch
+
+            if self.model is None:
+                raise ValueError(
+                    "Model has not been initialized or failed to initialize."
+                )
+
+            with torch.inference_mode():
+                text_pairs = [[query, doc] for doc in documents]
+                scores = self.model.compute_score(
+                    text_pairs, max_length=max_length, doc_type="text"
+                )
 
         return self._sort_documents(documents, scores)
 
@@ -120,7 +196,7 @@ class JinaRerankerToolkit(BaseToolkit):
         query: str,
         documents: List[str],
         max_length: int = 2048,
-    ) -> List[Tuple[str, float]]:
+    ) -> List[Dict[str, object]]:
         r"""Reranks image documents based on their relevance to a text query.
 
         Args:
@@ -130,21 +206,33 @@ class JinaRerankerToolkit(BaseToolkit):
                 (default: :obj:`2048`)
 
         Returns:
-            List[Tuple[str, float]]: A list of tuples containing
+            List[Dict[str, object]]: A list of dictionary containing
                 the reranked image URLs/paths and their relevance scores.
         """
-        import torch
-
-        if self.model is None:
-            raise ValueError(
-                "Model has not been initialized or failed to initialize."
-            )
-
-        with torch.inference_mode():
-            image_pairs = [[query, doc] for doc in documents]
-            scores = self.model.compute_score(
-                image_pairs, max_length=max_length, doc_type="image"
-            )
+        data = {
+            'model': self.model_name,
+            'query': query,
+            'top_n': len(documents),
+            'documents': documents,
+            'return_documents': True,
+        }
+
+        if self.use_api:
+            return self._call_jina_api(data)
+
+        else:
+            import torch
+
+            if self.model is None:
+                raise ValueError(
+                    "Model has not been initialized or failed to initialize."
+                )
+
+            with torch.inference_mode():
+                image_pairs = [[query, doc] for doc in documents]
+                scores = self.model.compute_score(
+                    image_pairs, max_length=max_length, doc_type="image"
+                )
 
         return self._sort_documents(documents, scores)
 
@@ -153,7 +241,7 @@ class JinaRerankerToolkit(BaseToolkit):
         image_query: str,
         documents: List[str],
         max_length: int = 2048,
-    ) -> List[Tuple[str, float]]:
+    ) -> List[Dict[str, object]]:
         r"""Reranks text documents based on their relevance to an image query.
 
         Args:
@@ -163,30 +251,45 @@ class JinaRerankerToolkit(BaseToolkit):
                 (default: :obj:`2048`)
 
         Returns:
-            List[Tuple[str, float]]: A list of tuples containing
+            List[Dict[str, object]]: A list of dictionary containing
                 the reranked documents and their relevance scores.
         """
-        import torch
-
-        if self.model is None:
-            raise ValueError("Model has not been initialized.")
-        with torch.inference_mode():
-            image_pairs = [[image_query, doc] for doc in documents]
-            scores = self.model.compute_score(
-                image_pairs,
-                max_length=max_length,
-                query_type="image",
-                doc_type="text",
-            )
-
-        return self._sort_documents(documents, scores)
+        data = {
+            'model': self.model_name,
+            'query': image_query,
+            'top_n': len(documents),
+            'documents': documents,
+            'return_documents': True,
+        }
+
+        if self.use_api:
+            return self._call_jina_api(data)
+
+        else:
+            import torch
+
+            if self.model is None:
+                raise ValueError(
+                    "Model has not been initialized or failed to initialize."
+                )
+
+            with torch.inference_mode():
+                image_pairs = [[image_query, doc] for doc in documents]
+                scores = self.model.compute_score(
+                    image_pairs,
+                    max_length=max_length,
+                    query_type="image",
+                    doc_type="text",
+                )
+
+            return self._sort_documents(documents, scores)
 
     def image_query_image_documents(
         self,
         image_query: str,
         documents: List[str],
         max_length: int = 2048,
-    ) -> List[Tuple[str, float]]:
+    ) -> List[Dict[str, object]]:
         r"""Reranks image documents based on their relevance to an image query.
 
         Args:
@@ -196,24 +299,38 @@ class JinaRerankerToolkit(BaseToolkit):
                 (default: :obj:`2048`)
 
         Returns:
-            List[Tuple[str, float]]: A list of tuples containing
+            List[Dict[str, object]]: A list of dictionary containing
                 the reranked image URLs/paths and their relevance scores.
         """
-        import torch
-
-        if self.model is None:
-            raise ValueError("Model has not been initialized.")
-
-        with torch.inference_mode():
-            image_pairs = [[image_query, doc] for doc in documents]
-            scores = self.model.compute_score(
-                image_pairs,
-                max_length=max_length,
-                query_type="image",
-                doc_type="image",
-            )
-
-        return self._sort_documents(documents, scores)
+        data = {
+            'model': self.model_name,
+            'query': image_query,
+            'top_n': len(documents),
+            'documents': documents,
+            'return_documents': True,
+        }
+
+        if self.use_api:
+            return self._call_jina_api(data)
+
+        else:
+            import torch
+
+            if self.model is None:
+                raise ValueError(
+                    "Model has not been initialized or failed to initialize."
+                )
+
+            with torch.inference_mode():
+                image_pairs = [[image_query, doc] for doc in documents]
+                scores = self.model.compute_score(
+                    image_pairs,
+                    max_length=max_length,
+                    query_type="image",
+                    doc_type="image",
+                )
+
+            return self._sort_documents(documents, scores)
 
     def get_tools(self) -> List[FunctionTool]:
         r"""Returns a list of FunctionTool objects representing the

camel/toolkits/mcp_toolkit.py

@@ -34,8 +34,10 @@ from urllib.parse import urlparse
 if TYPE_CHECKING:
     from mcp import ClientSession, ListToolsResult, Tool
 
+
 from camel.logger import get_logger
 from camel.toolkits import BaseToolkit, FunctionTool
+from camel.utils.commons import run_async
 
 logger = get_logger(__name__)
 
@@ -84,7 +86,6 @@ class MCPClient(BaseToolkit):
            await client.disconnect()
            ```
 
-
     Attributes:
         command_or_url (str): URL for SSE mode or command executable for stdio
             mode. (default: :obj:`None`)
@@ -96,6 +97,10 @@ class MCPClient(BaseToolkit):
             (default: :obj:`None`)
         headers (Dict[str, str]): Headers for the HTTP request.
             (default: :obj:`None`)
+        mode (Optional[str]): Connection mode. Can be "sse" for Server-Sent
+            Events, "streamable-http" for streaming HTTP,
+            or None for stdio mode.
+            (default: :obj:`None`)
         strict (Optional[bool]): Whether to enforce strict mode for the
             function call. (default: :obj:`False`)
     """
@@ -107,6 +112,7 @@ class MCPClient(BaseToolkit):
         env: Optional[Dict[str, str]] = None,
         timeout: Optional[float] = None,
         headers: Optional[Dict[str, str]] = None,
+        mode: Optional[str] = None,
         strict: Optional[bool] = False,
     ):
         from mcp import Tool
@@ -118,6 +124,7 @@ class MCPClient(BaseToolkit):
         self.env = env or {}
         self.headers = headers or {}
         self.strict = strict
+        self.mode = mode
 
         self._mcp_tools: List[Tool] = []
         self._session: Optional['ClientSession'] = None
@@ -133,6 +140,7 @@ class MCPClient(BaseToolkit):
         from mcp.client.session import ClientSession
         from mcp.client.sse import sse_client
         from mcp.client.stdio import StdioServerParameters, stdio_client
+        from mcp.client.streamable_http import streamablehttp_client
 
         if self._is_connected:
             logger.warning("Server is already connected")
@@ -140,16 +148,37 @@ class MCPClient(BaseToolkit):
 
         try:
             if urlparse(self.command_or_url).scheme in ("http", "https"):
-                (
-                    read_stream,
-                    write_stream,
-                ) = await self._exit_stack.enter_async_context(
-                    sse_client(
-                        self.command_or_url,
-                        headers=self.headers,
-                        timeout=self.timeout,
+                if self.mode == "sse" or self.mode is None:
+                    (
+                        read_stream,
+                        write_stream,
+                    ) = await self._exit_stack.enter_async_context(
+                        sse_client(
+                            self.command_or_url,
+                            headers=self.headers,
+                            timeout=self.timeout,
+                        )
+                    )
+                elif self.mode == "streamable-http":
+                    try:
+                        (
+                            read_stream,
+                            write_stream,
+                            _,
+                        ) = await self._exit_stack.enter_async_context(
+                            streamablehttp_client(
+                                self.command_or_url,
+                                headers=self.headers,
+                                timeout=timedelta(seconds=self.timeout),
+                            )
+                        )
+                    except Exception as e:
+                        # Handle anyio task group errors
+                        logger.error(f"Streamable HTTP client error: {e}")
+                else:
+                    raise ValueError(
+                        f"Invalid mode '{self.mode}' for HTTP URL"
                     )
-                )
             else:
                 command = self.command_or_url
                 arguments = self.args
@@ -192,16 +221,28 @@ class MCPClient(BaseToolkit):
             logger.error(f"Failed to connect to MCP server: {e}")
             raise e
 
+    def connect_sync(self):
+        r"""Synchronously connect to the MCP server."""
+        return run_async(self.connect)()
+
     async def disconnect(self):
         r"""Explicitly disconnect from the MCP server."""
         # If the server is not connected, do nothing
         if not self._is_connected:
             return
         self._is_connected = False
-        await self._exit_stack.aclose()
-        # Reset the exit stack and session for future reuse purposes
-        self._exit_stack = AsyncExitStack()
-        self._session = None
+
+        try:
+            await self._exit_stack.aclose()
+        except Exception as e:
+            logger.warning(f"{e}")
+        finally:
+            self._exit_stack = AsyncExitStack()
+            self._session = None
+
+    def disconnect_sync(self):
+        r"""Synchronously disconnect from the MCP server."""
+        return run_async(self.disconnect)()
 
     @asynccontextmanager
     async def connection(self):
@@ -217,7 +258,14 @@ class MCPClient(BaseToolkit):
             await self.connect()
             yield self
         finally:
-            await self.disconnect()
+            try:
+                await self.disconnect()
+            except Exception as e:
+                logger.warning(f"Error: {e}")
+
+    def connection_sync(self):
+        r"""Synchronously connect to the MCP server."""
+        return run_async(self.connection)()
 
     async def list_mcp_tools(self) -> Union[str, "ListToolsResult"]:
         r"""Retrieves the list of available tools from the connected MCP
@@ -234,6 +282,11 @@ class MCPClient(BaseToolkit):
             logger.exception("Failed to list MCP tools")
             raise e
 
+    def list_mcp_tools_sync(self) -> Union[str, "ListToolsResult"]:
+        r"""Synchronously list the available tools from the connected MCP
+        server."""
+        return run_async(self.list_mcp_tools)()
+
     def generate_function_from_mcp_tool(self, mcp_tool: "Tool") -> Callable:
         r"""Dynamically generates a Python callable function corresponding to
         a given MCP tool.
@@ -355,6 +408,10 @@ class MCPClient(BaseToolkit):
 
         return dynamic_function
 
+    def generate_function_from_mcp_tool_sync(self, mcp_tool: "Tool") -> Any:
+        r"""Synchronously generate a function from an MCP tool."""
+        return run_async(self.generate_function_from_mcp_tool)(mcp_tool)
+
     def _build_tool_schema(self, mcp_tool: "Tool") -> Dict[str, Any]:
         input_schema = mcp_tool.inputSchema
         properties = input_schema.get("properties", {})
@@ -428,6 +485,10 @@ class MCPClient(BaseToolkit):
 
         return await self._session.call_tool(tool_name, tool_args)
 
+    def call_tool_sync(self, tool_name: str, tool_args: Dict[str, Any]) -> Any:
+        r"""Synchronously call a tool."""
+        return run_async(self.call_tool)(tool_name, tool_args)
+
     @property
     def session(self) -> Optional["ClientSession"]:
         return self._session
@@ -440,6 +501,7 @@ class MCPClient(BaseToolkit):
         env: Optional[Dict[str, str]] = None,
         timeout: Optional[float] = None,
         headers: Optional[Dict[str, str]] = None,
+        mode: Optional[str] = None,
     ) -> "MCPClient":
         r"""Factory method that creates and connects to the MCP server.
 
@@ -457,6 +519,10 @@ class MCPClient(BaseToolkit):
                 (default: :obj:`None`)
             headers (Optional[Dict[str, str]]): Headers for the HTTP request.
                 (default: :obj:`None`)
+            mode (Optional[str]): Connection mode. Can be "sse" for
+                Server-Sent Events, "streamable-http" for
+                streaming HTTP, or None for stdio mode.
+                (default: :obj:`None`)
 
         Returns:
             MCPClient: A fully initialized and connected MCPClient instance.
@@ -470,6 +536,7 @@ class MCPClient(BaseToolkit):
             env=env,
             timeout=timeout,
             headers=headers,
+            mode=mode,
         )
         try:
             await client.connect()
@@ -480,6 +547,21 @@ class MCPClient(BaseToolkit):
             logger.error(f"Failed to initialize MCPClient: {e}")
             raise RuntimeError(f"Failed to initialize MCPClient: {e}") from e
 
+    @classmethod
+    def create_sync(
+        self,
+        command_or_url: str,
+        args: Optional[List[str]] = None,
+        env: Optional[Dict[str, str]] = None,
+        timeout: Optional[float] = None,
+        headers: Optional[Dict[str, str]] = None,
+        mode: Optional[str] = None,
+    ) -> "MCPClient":
+        r"""Synchronously create and connect to the MCP server."""
+        return run_async(self.create)(
+            command_or_url, args, env, timeout, headers, mode
+        )
+
     async def __aenter__(self) -> "MCPClient":
         r"""Async context manager entry point. Automatically connects to the
         MCP server when used in an async with statement.
@@ -490,12 +572,35 @@ class MCPClient(BaseToolkit):
         await self.connect()
         return self
 
-    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+    def __enter__(self) -> "MCPClient":
+        r"""Synchronously enter the async context manager."""
+        return run_async(self.__aenter__)()
+
+    async def __aexit__(self) -> None:
         r"""Async context manager exit point. Automatically disconnects from
         the MCP server when exiting an async with statement.
+
+        Returns:
+            None
         """
         await self.disconnect()
 
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        r"""Synchronously exit the async context manager.
+
+        Args:
+            exc_type (Optional[Type[Exception]]): The type of exception that
+                occurred during the execution of the with statement.
+            exc_val (Optional[Exception]): The exception that occurred during
+                the execution of the with statement.
+            exc_tb (Optional[TracebackType]): The traceback of the exception
+                that occurred during the execution of the with statement.
+
+        Returns:
+            None
+        """
+        return run_async(self.__aexit__)()
+
 
 class MCPToolkit(BaseToolkit):
     r"""MCPToolkit provides a unified interface for managing multiple
@@ -679,6 +784,7 @@ class MCPToolkit(BaseToolkit):
                 env={**os.environ, **cfg.get("env", {})},
                 timeout=cfg.get("timeout", None),
                 headers=headers,
+                mode=cfg.get("mode", None),
                 strict=strict,
             )
             all_servers.append(server)
@@ -707,6 +813,10 @@ class MCPToolkit(BaseToolkit):
             logger.error(f"Failed to connect to one or more MCP servers: {e}")
             raise e
 
+    def connect_sync(self):
+        r"""Synchronously connect to all MCP servers."""
+        return run_async(self.connect)()
+
     async def disconnect(self):
         r"""Explicitly disconnect from all MCP servers."""
         if not self._connected:
@@ -716,6 +826,10 @@ class MCPToolkit(BaseToolkit):
             await server.disconnect()
         self._connected = False
 
+    def disconnect_sync(self):
+        r"""Synchronously disconnect from all MCP servers."""
+        return run_async(self.disconnect)()
+
     @asynccontextmanager
     async def connection(self) -> AsyncGenerator["MCPToolkit", None]:
         r"""Async context manager that simultaneously establishes connections
@@ -730,6 +844,10 @@ class MCPToolkit(BaseToolkit):
         finally:
             await self.disconnect()
 
+    def connection_sync(self):
+        r"""Synchronously connect to all MCP servers."""
+        return run_async(self.connection)()
+
     def is_connected(self) -> bool:
         r"""Checks if all the managed servers are connected.