agno 2.2.13__py3-none-any.whl → 2.4.3__py3-none-any.whl

agno/tools/postgres.py

@@ -14,6 +14,21 @@ from agno.utils.log import log_debug, log_error
 
 
 class PostgresTools(Toolkit):
+    """
+    A toolkit for interacting with PostgreSQL databases.
+
+    Args:
+        connection (Optional[PgConnection[DictRow]]): Existing database connection to reuse.
+        db_name (Optional[str]): Database name to connect to.
+        user (Optional[str]): Username for authentication.
+        password (Optional[str]): Password for authentication.
+        host (Optional[str]): PostgreSQL server hostname.
+        port (Optional[int]): PostgreSQL server port number.
+        table_schema (str): Default schema for table operations. Default is "public".
+    """
+
+    _requires_connect: bool = True
+
     def __init__(
         self,
         connection: Optional[PgConnection[DictRow]] = None,
@@ -44,50 +59,71 @@ class PostgresTools(Toolkit):
 
         super().__init__(name="postgres_tools", tools=tools, **kwargs)
 
-    @property
-    def connection(self) -> PgConnection[DictRow]:
-        """
-        Returns the Postgres psycopg connection.
-        :return psycopg.connection.Connection: psycopg connection
+    def connect(self) -> PgConnection[DictRow]:
         """
-        if self._connection is None or self._connection.closed:
-            log_debug("Establishing new PostgreSQL connection.")
-            connection_kwargs: Dict[str, Any] = {"row_factory": dict_row}
-            if self.db_name:
-                connection_kwargs["dbname"] = self.db_name
-            if self.user:
-                connection_kwargs["user"] = self.user
-            if self.password:
-                connection_kwargs["password"] = self.password
-            if self.host:
-                connection_kwargs["host"] = self.host
-            if self.port:
-                connection_kwargs["port"] = self.port
-
-            connection_kwargs["options"] = f"-c search_path={self.table_schema}"
-
-            self._connection = psycopg.connect(**connection_kwargs)
-            self._connection.read_only = True
+        Establish a connection to the PostgreSQL database.
 
+        Returns:
+            The database connection object.
+        """
+        if self._connection is not None and not self._connection.closed:
+            log_debug("Connection already established, reusing existing connection")
+            return self._connection
+
+        log_debug("Establishing new PostgreSQL connection.")
+        connection_kwargs: Dict[str, Any] = {"row_factory": dict_row}
+        if self.db_name:
+            connection_kwargs["dbname"] = self.db_name
+        if self.user:
+            connection_kwargs["user"] = self.user
+        if self.password:
+            connection_kwargs["password"] = self.password
+        if self.host:
+            connection_kwargs["host"] = self.host
+        if self.port:
+            connection_kwargs["port"] = self.port
+
+        connection_kwargs["options"] = f"-c search_path={self.table_schema}"
+
+        self._connection = psycopg.connect(**connection_kwargs)
+        self._connection.read_only = True
         return self._connection
 
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        self.close()
-
-    def close(self):
+    def close(self) -> None:
         """Closes the database connection if it's open."""
         if self._connection and not self._connection.closed:
             log_debug("Closing PostgreSQL connection.")
             self._connection.close()
             self._connection = None
 
+    @property
+    def is_connected(self) -> bool:
+        """Check if a connection is currently established."""
+        return self._connection is not None and not self._connection.closed
+
+    def _ensure_connection(self) -> PgConnection[DictRow]:
+        """
+        Ensure a connection exists, creating one if necessary.
+
+        Returns:
+            The database connection object.
+        """
+        if not self.is_connected:
+            return self.connect()
+        return self._connection  # type: ignore
+
+    def __enter__(self):
+        return self.connect()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.is_connected:
+            self.close()
+
     def _execute_query(self, query: str, params: Optional[tuple] = None) -> str:
         try:
-            with self.connection.cursor() as cursor:
-                log_debug(f"Running PostgreSQL Query: {query} with Params: {params}")
+            connection = self._ensure_connection()
+            with connection.cursor() as cursor:
+                log_debug("Running PostgreSQL query")
                 cursor.execute(query, params)
 
                 if cursor.description is None:
@@ -105,8 +141,8 @@ class PostgresTools(Toolkit):
 
         except psycopg.Error as e:
             log_error(f"Database error: {e}")
-            if self.connection and not self.connection.closed:
-                self.connection.rollback()
+            if self._connection and not self._connection.closed:
+                self._connection.rollback()
             return f"Error executing query: {e}"
         except Exception as e:
             log_error(f"An unexpected error occurred: {e}")
@@ -146,7 +182,8 @@ class PostgresTools(Toolkit):
             A string containing a summary of the table.
         """
         try:
-            with self.connection.cursor() as cursor:
+            connection = self._ensure_connection()
+            with connection.cursor() as cursor:
                 # First, get column information using a parameterized query
                 schema_query = """
                     SELECT column_name, data_type
@@ -230,7 +267,8 @@ class PostgresTools(Toolkit):
         stmt = sql.SQL("SELECT * FROM {tbl};").format(tbl=table_identifier)
 
         try:
-            with self.connection.cursor() as cursor:
+            connection = self._ensure_connection()
+            with connection.cursor() as cursor:
                 cursor.execute(stmt)
 
                 if cursor.description is None:
@@ -245,6 +283,8 @@ class PostgresTools(Toolkit):
 
             return f"Successfully exported table '{table}' to '{path}'."
         except (psycopg.Error, IOError) as e:
+            if self._connection and not self._connection.closed:
+                self._connection.rollback()
             return f"Error exporting table: {e}"
 
     def run_query(self, query: str) -> str:

agno/tools/python.py

@@ -4,7 +4,7 @@ from pathlib import Path
 from typing import Any, List, Optional
 
 from agno.tools import Toolkit
-from agno.utils.log import log_debug, log_info, logger
+from agno.utils.log import log_debug, log_error, log_info, logger
 
 
 @functools.lru_cache(maxsize=None)
@@ -18,9 +18,11 @@ class PythonTools(Toolkit):
         base_dir: Optional[Path] = None,
         safe_globals: Optional[dict] = None,
         safe_locals: Optional[dict] = None,
+        restrict_to_base_dir: bool = True,
         **kwargs,
     ):
-        self.base_dir: Path = base_dir or Path.cwd()
+        self.base_dir: Path = (base_dir or Path.cwd()).resolve()
+        self.restrict_to_base_dir = restrict_to_base_dir
 
         # Restricted global and local scope
         self.safe_globals: dict = safe_globals or globals()
@@ -55,7 +57,9 @@ class PythonTools(Toolkit):
         """
         try:
             warn()
-            file_path = self.base_dir.joinpath(file_name)
+            safe, file_path = self._check_path(file_name, self.base_dir, self.restrict_to_base_dir)
+            if not safe:
+                return f"Error: Path '{file_name}' is outside the allowed base directory"
             log_debug(f"Saving code to {file_path}")
             if not file_path.parent.exists():
                 file_path.parent.mkdir(parents=True, exist_ok=True)
@@ -89,8 +93,9 @@ class PythonTools(Toolkit):
         """
         try:
             warn()
-            file_path = self.base_dir.joinpath(file_name)
-
+            safe, file_path = self._check_path(file_name, self.base_dir, self.restrict_to_base_dir)
+            if not safe:
+                return f"Error: Path '{file_name}' is outside the allowed base directory"
             log_info(f"Running {file_path}")
             globals_after_run = runpy.run_path(str(file_path), init_globals=self.safe_globals, run_name="__main__")
             if variable_to_return:
@@ -113,7 +118,10 @@ class PythonTools(Toolkit):
         """
         try:
             log_info(f"Reading file: {file_name}")
-            file_path = self.base_dir.joinpath(file_name)
+            safe, file_path = self._check_path(file_name, self.base_dir, self.restrict_to_base_dir)
+            if not safe:
+                log_error(f"Attempted to read file outside base directory: {file_name}")
+                return "Error reading file: path outside allowed directory"
             contents = file_path.read_text(encoding="utf-8")
             return str(contents)
         except Exception as e:

agno/tools/reasoning.py

@@ -1,7 +1,8 @@
 from textwrap import dedent
-from typing import Any, Dict, List, Optional
+from typing import Any, List, Optional
 
 from agno.reasoning.step import NextAction, ReasoningStep
+from agno.run import RunContext
 from agno.tools import Toolkit
 from agno.utils.log import log_debug, log_error
 
@@ -49,7 +50,7 @@ class ReasoningTools(Toolkit):
 
     def think(
         self,
-        session_state: Dict[str, Any],
+        run_context: RunContext,
         title: str,
         thought: str,
         action: Optional[str] = None,
@@ -80,21 +81,24 @@ class ReasoningTools(Toolkit):
                 confidence=confidence,
             )
 
-            current_run_id = session_state.get("current_run_id", None)
+            current_run_id = run_context.run_id
 
             # Add this step to the Agent's session state
-            if session_state is None:
-                session_state = {}
-            if "reasoning_steps" not in session_state:
-                session_state["reasoning_steps"] = {}
-            if current_run_id not in session_state["reasoning_steps"]:
-                session_state["reasoning_steps"][current_run_id] = []
-            session_state["reasoning_steps"][current_run_id].append(reasoning_step.model_dump_json())
+            if run_context.session_state is None:
+                run_context.session_state = {}
+            if "reasoning_steps" not in run_context.session_state:
+                run_context.session_state["reasoning_steps"] = {}
+            if current_run_id not in run_context.session_state["reasoning_steps"]:
+                run_context.session_state["reasoning_steps"][current_run_id] = []
+            run_context.session_state["reasoning_steps"][current_run_id].append(reasoning_step.model_dump_json())
 
             # Return all previous reasoning_steps and the new reasoning_step
-            if "reasoning_steps" in session_state and current_run_id in session_state["reasoning_steps"]:
+            if (
+                "reasoning_steps" in run_context.session_state
+                and current_run_id in run_context.session_state["reasoning_steps"]
+            ):
                 formatted_reasoning_steps = ""
-                for i, step in enumerate(session_state["reasoning_steps"][current_run_id], 1):
+                for i, step in enumerate(run_context.session_state["reasoning_steps"][current_run_id], 1):
                     step_parsed = ReasoningStep.model_validate_json(step)
                     step_str = dedent(f"""\
 Step {i}:
@@ -112,7 +116,7 @@ Confidence: {step_parsed.confidence}
 
     def analyze(
         self,
-        session_state: Dict[str, Any],
+        run_context: RunContext,
         title: str,
         result: str,
         analysis: str,
@@ -150,20 +154,23 @@ Confidence: {step_parsed.confidence}
                 confidence=confidence,
             )
 
-            current_run_id = session_state.get("current_run_id", None)
+            current_run_id = run_context.run_id
             # Add this step to the Agent's session state
-            if session_state is None:
-                session_state = {}
-            if "reasoning_steps" not in session_state:
-                session_state["reasoning_steps"] = {}
-            if current_run_id not in session_state["reasoning_steps"]:
-                session_state["reasoning_steps"][current_run_id] = []
-            session_state["reasoning_steps"][current_run_id].append(reasoning_step.model_dump_json())
+            if run_context.session_state is None:
+                run_context.session_state = {}
+            if "reasoning_steps" not in run_context.session_state:
+                run_context.session_state["reasoning_steps"] = {}
+            if current_run_id not in run_context.session_state["reasoning_steps"]:
+                run_context.session_state["reasoning_steps"][current_run_id] = []
+            run_context.session_state["reasoning_steps"][current_run_id].append(reasoning_step.model_dump_json())
 
             # Return all previous reasoning_steps and the new reasoning_step
-            if "reasoning_steps" in session_state and current_run_id in session_state["reasoning_steps"]:
+            if (
+                "reasoning_steps" in run_context.session_state
+                and current_run_id in run_context.session_state["reasoning_steps"]
+            ):
                 formatted_reasoning_steps = ""
-                for i, step in enumerate(session_state["reasoning_steps"][current_run_id], 1):
+                for i, step in enumerate(run_context.session_state["reasoning_steps"][current_run_id], 1):
                     step_parsed = ReasoningStep.model_validate_json(step)
                     step_str = dedent(f"""\
 Step {i}: