camel-ai 0.2.72a8__py3-none-any.whl → 0.2.73__py3-none-any.whl

camel/storages/vectordb_storages/surreal.py

@@ -68,7 +68,9 @@ class SurrealStorage(BaseVectorStorage):
         url: str = "ws://localhost:8000/rpc",
         table: str = "vector_store",
         vector_dim: int = 786,
+        vector_type: str = "F64",
         distance: VectorDistance = VectorDistance.COSINE,
+        hnsw_effort: int = 40,
         namespace: str = "default",
         database: str = "demo",
         user: str = "root",
@@ -96,6 +98,8 @@ class SurrealStorage(BaseVectorStorage):
                 (default: :obj:`"root"`)
         """
 
+        from surrealdb import Surreal
+
         self.url = url
         self.table = table
         self.ns = namespace
@@ -103,9 +107,14 @@ class SurrealStorage(BaseVectorStorage):
         self.user = user
         self.password = password
         self.vector_dim = vector_dim
+        self.vector_type = vector_type
         self.distance = distance
+        self._hnsw_effort = hnsw_effort
+        self._surreal_client = Surreal(self.url)
+        self._surreal_client.signin({"username": user, "password": password})
+        self._surreal_client.use(namespace, database)
+
         self._check_and_create_table()
-        self._surreal_client = None
 
     def _table_exists(self) -> bool:
         r"""Check whether the target table exists in the database.
@@ -113,74 +122,66 @@ class SurrealStorage(BaseVectorStorage):
         Returns:
             bool: True if the table exists, False otherwise.
         """
-        from surrealdb import Surreal  # type: ignore[import-not-found]
+        res = self._surreal_client.query("INFO FOR DB;")
+        tables = res.get('tables', {})
+        logger.debug(f"_table_exists: {res}")
+        return self.table in tables
 
-        with Surreal(self.url) as db:
-            db.signin({"username": self.user, "password": self.password})
-            db.use(self.ns, self.db)
-            res = db.query_raw("INFO FOR DB;")
-            tables = res['result'][0]['result'].get('tables', {})
-            return self.table in tables
-
-    def _get_table_info(self) -> Dict[str, int]:
+    def _get_table_info(self) -> dict[str, int | None]:
         r"""Retrieve dimension and record count from the table metadata.
 
         Returns:
             Dict[str, int]: A dictionary with 'dim' and 'count' keys.
         """
-        from surrealdb import Surreal  # type: ignore[import-not-found]
-
         if not self._table_exists():
             return {"dim": self.vector_dim, "count": 0}
-        with Surreal(self.url) as db:
-            db.signin({"username": self.user, "password": self.password})
-            db.use(self.ns, self.db)
-            res = db.query_raw(f"INFO FOR TABLE {self.table};")
-
-            indexes = res['result'][0]['result'].get("indexes", {})
-
-            dim = self.vector_dim
-            idx_def = indexes.get("hnsw_idx")
-            if idx_def and isinstance(idx_def, str):
-                m = re.search(r"DIMENSION\s+(\d+)", idx_def)
-                if m:
-                    dim = int(m.group(1))
-            cnt = db.query_raw(f"SELECT COUNT() AS count FROM {self.table};")
-            try:
-                count = cnt['result'][0]['result'][0]['count']
-            except (KeyError, IndexError, TypeError):
-                logger.warning(
-                    "Unexpected result format when counting records: %s", cnt
-                )
-                count = 0
-
-            return {"dim": dim, "count": count}
+        res = self._surreal_client.query(f"INFO FOR TABLE {self.table};")
+        logger.debug(f"_get_table_info: {res}")
+        indexes = res.get("indexes", {})
+
+        dim = self.vector_dim
+        idx_def = indexes.get("hnsw_idx")
+        if idx_def and isinstance(idx_def, str):
+            m = re.search(r"DIMENSION\s+(\d+)", idx_def)
+            if m:
+                dim = int(m.group(1))
+        cnt = self._surreal_client.query(
+            f"SELECT COUNT() FROM ONLY {self.table} GROUP ALL LIMIT 1;"
+        )
+        count = cnt.get("count", 0)
+        return {"dim": dim, "count": count}
 
     def _create_table(self):
-        r"""Define and create the vector storage table with HNSW index."""
-        from surrealdb import Surreal  # type: ignore[import-not-found]
-
-        with Surreal(self.url) as db:
-            db.signin({"username": self.user, "password": self.password})
-            db.use(self.ns, self.db)
-            db.query_raw(
-                f"""DEFINE TABLE {self.table} SCHEMALESS;
-                DEFINE FIELD payload    ON {self.table} TYPE object;
-                DEFINE FIELD embedding  ON {self.table} TYPE array;
-                DEFINE INDEX hnsw_idx   ON {self.table}
-                            FIELDS embedding HNSW DIMENSION {self.vector_dim};
-                """
+        r"""Define and create the vector storage table with HNSW index.
+
+        Documentation: https://surrealdb.com/docs/surrealdb/reference-guide/
+        vector-search#vector-search-cheat-sheet
+        """
+        if self.distance.value not in ["cosine", "euclidean", "manhattan"]:
+            raise ValueError(
+                f"Unsupported distance metric: {self.distance.value}"
             )
+        surql_query = f"""
+        DEFINE TABLE {self.table} SCHEMALESS;
+        DEFINE FIELD payload    ON {self.table} FLEXIBLE TYPE object;
+        DEFINE FIELD embedding  ON {self.table} TYPE array<float>;
+        DEFINE INDEX hnsw_idx   ON {self.table}
+                    FIELDS embedding
+                    HNSW DIMENSION {self.vector_dim}
+                    DIST {self.distance.value}
+                    TYPE {self.vector_type}
+                    EFC 150 M 12 M0 24;
+        """
+        logger.debug(f"_create_table query: {surql_query}")
+        res = self._surreal_client.query_raw(surql_query)
+        logger.debug(f"_create_table response: {res}")
+        if "error" in res:
+            raise ValueError(f"Failed to create table: {res['error']}")
         logger.info(f"Table '{self.table}' created successfully.")
 
     def _drop_table(self):
         r"""Drop the vector storage table if it exists."""
-        from surrealdb import Surreal  # type: ignore[import-not-found]
-
-        with Surreal(self.url) as db:
-            db.signin({"username": self.user, "password": self.password})
-            db.use(self.ns, self.db)
-            db.query_raw(f"REMOVE TABLE IF EXISTS {self.table};")
+        self._surreal_client.query_raw(f"REMOVE TABLE IF EXISTS {self.table};")
         logger.info(f"Table '{self.table}' deleted successfully.")
 
     def _check_and_create_table(self):
@@ -240,62 +241,34 @@ class SurrealStorage(BaseVectorStorage):
             List[VectorDBQueryResult]: Ranked list of matching records
                 with similarity scores.
         """
-        from surrealdb import Surreal  # type: ignore[import-not-found]
-
-        metric = {
-            VectorDistance.COSINE: "cosine",
-            VectorDistance.EUCLIDEAN: "euclidean",
-            VectorDistance.DOT: "dot",
-        }[self.distance]
-
-        metric_func = {
-            VectorDistance.COSINE: "vector::similarity::cosine",
-            VectorDistance.EUCLIDEAN: "vector::distance::euclidean",
-            VectorDistance.DOT: "vector::dot",
-        }.get(self.distance)
-
-        if not metric_func:
-            raise ValueError(f"Unsupported distance metric: {self.distance}")
-
-        with Surreal(self.url) as db:
-            db.signin({"username": self.user, "password": self.password})
-            db.use(self.ns, self.db)
-
-            # Use parameterized query to prevent SQL injection
-            sql_query = f"""SELECT payload, embedding,
-                    {metric_func}(embedding, $query_vec) AS score
-                FROM {self.table}
-                WHERE embedding <|{query.top_k},{metric}|> $query_vec
-                ORDER BY score;
-            """
-
-            response = db.query_raw(
-                sql_query, {"query_vec": query.query_vector}
-            )
+        surql_query = f"""
+            SELECT id, embedding, payload, vector::distance::knn() AS dist
+            FROM {self.table}
+            WHERE embedding <|{query.top_k},{self._hnsw_effort}|> $vector
+            ORDER BY dist;
+        """
+        logger.debug(
+            f"query surql: {surql_query} with $vector = {query.query_vector}"
+        )
 
-            if not response.get("result") or not response["result"]:
-                return []
-
-            results = response["result"][0]
-
-            if "result" not in results:
-                return []
-
-            return [
-                VectorDBQueryResult(
-                    record=VectorRecord(
-                        vector=row["embedding"], payload=row.get("payload", {})
-                    ),
-                    similarity=(
-                        1.0 - row["score"]
-                        if self.distance == VectorDistance.COSINE
-                        else 1.0 / (1.0 + row["score"])
-                        if self.distance == VectorDistance.EUCLIDEAN
-                        else row["score"]
-                    ),
-                )
-                for row in results["result"]
-            ]
+        response = self._surreal_client.query(
+            surql_query, {"vector": query.query_vector}
+        )
+        logger.debug(f"query response: {response}")
+
+        return [
+            VectorDBQueryResult(
+                record=VectorRecord(
+                    id=row["id"].id,
+                    vector=row["embedding"],
+                    payload=row["payload"],
+                ),
+                similarity=1.0 - row["dist"]
+                if self.distance == VectorDistance.COSINE
+                else -row["score"],
+            )
+            for row in response
+        ]
 
     def add(self, records: List[VectorRecord], **kwargs) -> None:
         r"""Insert validated vector records into the SurrealDB table.
@@ -306,16 +279,10 @@ class SurrealStorage(BaseVectorStorage):
         logger.info(
             "Adding %d records to table '%s'.", len(records), self.table
         )
-        from surrealdb import Surreal  # type: ignore[import-not-found]
-
         try:
-            with Surreal(self.url) as db:
-                db.signin({"username": self.user, "password": self.password})
-                db.use(self.ns, self.db)
-
-                validated_records = self._validate_and_convert_records(records)
-                for record in validated_records:
-                    db.create(self.table, record)
+            validated_records = self._validate_and_convert_records(records)
+            for record in validated_records:
+                self._surreal_client.create(self.table, record)
 
             logger.info(
                 "Successfully added %d records to table '%s'.",
@@ -340,32 +307,23 @@ class SurrealStorage(BaseVectorStorage):
             ids (Optional[List[str]]): List of record IDs to delete.
             if_all (bool): Whether to delete all records in the table.
         """
-        from surrealdb import Surreal  # type: ignore[import-not-found]
-        from surrealdb.data.types.record_id import (  # type: ignore[import-not-found]
-            RecordID,
-        )
+        from surrealdb.data.types.record_id import RecordID
 
         try:
-            with Surreal(self.url) as db:
-                db.signin({"username": self.user, "password": self.password})
-                db.use(self.ns, self.db)
-
-                if if_all:
-                    db.delete(self.table, **kwargs)
-                    logger.info(
-                        f"Deleted all records from table '{self.table}'"
-                    )
-                    return
-
-                if not ids:
-                    raise ValueError(
-                        "Either `ids` must be provided or `if_all=True`"
-                    )
-
-                for id_str in ids:
-                    rec = RecordID(self.table, id_str)
-                    db.delete(rec, **kwargs)
-                    logger.info(f"Deleted record {rec}")
+            if if_all:
+                self._surreal_client.delete(self.table, **kwargs)
+                logger.info(f"Deleted all records from table '{self.table}'")
+                return
+
+            if not ids:
+                raise ValueError(
+                    "Either `ids` must be provided or `if_all=True`"
+                )
+
+            for id_str in ids:
+                rec = RecordID(self.table, id_str)
+                self._surreal_client.delete(rec, **kwargs)
+                logger.info(f"Deleted record {rec}")
 
         except Exception as e:
             logger.exception("Error deleting records from SurrealDB")
@@ -404,12 +362,4 @@ class SurrealStorage(BaseVectorStorage):
     @property
     def client(self) -> "Surreal":
         r"""Provides access to the underlying SurrealDB client."""
-        if self._surreal_client is None:
-            from surrealdb import Surreal  # type: ignore[import-not-found]
-
-            self._surreal_client = Surreal(self.url)
-            self._surreal_client.signin(  # type: ignore[attr-defined]
-                {"username": self.user, "password": self.password}
-            )
-            self._surreal_client.use(self.ns, self.db)  # type: ignore[attr-defined]
         return self._surreal_client

camel/toolkits/__init__.py

@@ -31,7 +31,7 @@ from .meshy_toolkit import MeshyToolkit
 from .openbb_toolkit import OpenBBToolkit
 from .bohrium_toolkit import BohriumToolkit
 
-from .base import BaseToolkit
+from .base import BaseToolkit, RegisteredAgentToolkit
 from .google_maps_toolkit import GoogleMapsToolkit
 from .code_execution import CodeExecutionToolkit
 from .github_toolkit import GithubToolkit
@@ -87,6 +87,8 @@ from .note_taking_toolkit import NoteTakingToolkit
 from .message_agent_toolkit import AgentCommunicationToolkit
 from .web_deploy_toolkit import WebDeployToolkit
 from .screenshot_toolkit import ScreenshotToolkit
+from .message_integration import ToolkitMessageIntegration
+from .notion_mcp_toolkit import NotionMCPToolkit
 
 __all__ = [
     'BaseToolkit',
@@ -162,4 +164,7 @@ __all__ = [
     'AgentCommunicationToolkit',
     'WebDeployToolkit',
     'ScreenshotToolkit',
+    'RegisteredAgentToolkit',
+    'ToolkitMessageIntegration',
+    'NotionMCPToolkit',
 ]

camel/toolkits/base.py

@@ -12,11 +12,18 @@
 # limitations under the License.
 # ========= Copyright 2023-2024 @ CAMEL-AI.org. All Rights Reserved. =========
 
-from typing import List, Literal, Optional
+from typing import TYPE_CHECKING, List, Literal, Optional
 
+from camel.logger import get_logger
 from camel.toolkits import FunctionTool
 from camel.utils import AgentOpsMeta, with_timeout
 
+if TYPE_CHECKING:
+    from camel.agents import ChatAgent
+
+
+logger = get_logger(__name__)
+
 
 class BaseToolkit(metaclass=AgentOpsMeta):
     r"""Base class for toolkits.
@@ -41,7 +48,9 @@ class BaseToolkit(metaclass=AgentOpsMeta):
         super().__init_subclass__(**kwargs)
         for attr_name, attr_value in cls.__dict__.items():
             if callable(attr_value) and not attr_name.startswith("__"):
-                setattr(cls, attr_name, with_timeout(attr_value))
+                # Skip methods that have manual timeout management
+                if not getattr(attr_value, '_manual_timeout', False):
+                    setattr(cls, attr_name, with_timeout(attr_value))
 
     def get_tools(self) -> List[FunctionTool]:
         r"""Returns a list of FunctionTool objects representing the
@@ -63,3 +72,52 @@ class BaseToolkit(metaclass=AgentOpsMeta):
                 the MCP server in.
         """
         self.mcp.run(mode)
+
+
+class RegisteredAgentToolkit:
+    r"""Mixin class for toolkits that need to register a ChatAgent.
+
+    This mixin provides a standard interface for toolkits that require
+    a reference to a ChatAgent instance. The ChatAgent will check if a
+    toolkit has this mixin and automatically register itself.
+    """
+
+    def __init__(self):
+        self._agent: Optional["ChatAgent"] = None
+
+    @property
+    def agent(self) -> Optional["ChatAgent"]:
+        r"""Get the registered ChatAgent instance.
+
+        Returns:
+            Optional[ChatAgent]: The registered agent, or None if not
+                registered.
+
+        Note:
+            If None is returned, it means the toolkit has not been registered
+            with a ChatAgent yet. Make sure to pass this toolkit to a ChatAgent
+            via the toolkits parameter during initialization.
+        """
+        if self._agent is None:
+            logger.warning(
+                f"{self.__class__.__name__} does not have a "
+                f"registered ChatAgent. "
+                f"Please ensure this toolkit is passed to a ChatAgent via the "
+                f"'toolkits_to_register_agent' parameter during ChatAgent "
+                f"initialization if you want to use the tools that require a "
+                f"registered agent."
+            )
+        return self._agent
+
+    def register_agent(self, agent: "ChatAgent") -> None:
+        r"""Register a ChatAgent with this toolkit.
+
+        This method allows registering an agent after initialization. The
+        ChatAgent will automatically call this method if the toolkit to
+        register inherits from RegisteredAgentToolkit.
+
+        Args:
+            agent (ChatAgent): The ChatAgent instance to register.
+        """
+        self._agent = agent
+        logger.info(f"Agent registered with {self.__class__.__name__}")