agno 0.1.2__py3-none-any.whl → 2.3.13__py3-none-any.whl

agno/tools/postgres.py

@@ -1,35 +1,46 @@
-from typing import Any, Dict, Optional
+import csv
+from typing import Any, Dict, List, Optional
 
 try:
-    import psycopg2
+    import psycopg
+    from psycopg import sql
+    from psycopg.connection import Connection as PgConnection
+    from psycopg.rows import DictRow, dict_row
 except ImportError:
-    raise ImportError(
-        "`psycopg2` not installed. Please install using `pip install psycopg2`. If you face issues, try `pip install psycopg2-binary`."
-    )
+    raise ImportError("`psycopg` not installed. Please install using `pip install 'psycopg-binary'`.")
 
 from agno.tools import Toolkit
-from agno.utils.log import logger
+from agno.utils.log import log_debug, log_error
 
 
 class PostgresTools(Toolkit):
-    """A basic tool to connect to a PostgreSQL database and perform read-only operations on it."""
+    """
+    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[psycopg2.extensions.connection] = None,
+        connection: Optional[PgConnection[DictRow]] = None,
         db_name: Optional[str] = None,
         user: Optional[str] = None,
         password: Optional[str] = None,
         host: Optional[str] = None,
         port: Optional[int] = None,
-        run_queries: bool = True,
-        inspect_queries: bool = False,
-        summarize_tables: bool = True,
-        export_tables: bool = False,
         table_schema: str = "public",
+        **kwargs,
     ):
-        super().__init__(name="postgres_tools")
-        self._connection: Optional[psycopg2.extensions.connection] = connection
+        self._connection: Optional[PgConnection[DictRow]] = connection
         self.db_name: Optional[str] = db_name
         self.user: Optional[str] = user
         self.password: Optional[str] = password
@@ -37,208 +48,250 @@ class PostgresTools(Toolkit):
         self.port: Optional[int] = port
         self.table_schema: str = table_schema
 
-        self.register(self.show_tables)
-        self.register(self.describe_table)
-        if inspect_queries:
-            self.register(self.inspect_query)
-        if run_queries:
-            self.register(self.run_query)
-        if summarize_tables:
-            self.register(self.summarize_table)
-        if export_tables:
-            self.register(self.export_table_to_path)
+        tools: List[Any] = [
+            self.show_tables,
+            self.describe_table,
+            self.summarize_table,
+            self.inspect_query,
+            self.run_query,
+            self.export_table_to_path,
+        ]
+
+        super().__init__(name="postgres_tools", tools=tools, **kwargs)
+
+    def connect(self) -> PgConnection[DictRow]:
+        """
+        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 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 connection(self) -> psycopg2.extensions.connection:
+    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]:
         """
-        Returns the Postgres psycopg2 connection.
+        Ensure a connection exists, creating one if necessary.
 
-        :return psycopg2.extensions.connection: psycopg2 connection
+        Returns:
+            The database connection object.
         """
-        if self._connection is None:
-            connection_kwargs: Dict[str, Any] = {}
-            if self.db_name is not None:
-                connection_kwargs["database"] = self.db_name
-            if self.user is not None:
-                connection_kwargs["user"] = self.user
-            if self.password is not None:
-                connection_kwargs["password"] = self.password
-            if self.host is not None:
-                connection_kwargs["host"] = self.host
-            if self.port is not None:
-                connection_kwargs["port"] = self.port
-            if self.table_schema is not None:
-                connection_kwargs["options"] = f"-c search_path={self.table_schema}"
-
-            self._connection = psycopg2.connect(**connection_kwargs)
-            self._connection.set_session(readonly=True)
+        if not self.is_connected:
+            return self.connect()
+        return self._connection  # type: ignore
 
-        return self._connection
+    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:
+            connection = self._ensure_connection()
+            with connection.cursor() as cursor:
+                log_debug("Running PostgreSQL query")
+                cursor.execute(query, params)
+
+                if cursor.description is None:
+                    return cursor.statusmessage or "Query executed successfully with no output."
+
+                columns = [desc[0] for desc in cursor.description]
+                rows = cursor.fetchall()
+
+                if not rows:
+                    return f"Query returned no results.\nColumns: {', '.join(columns)}"
+
+                header = ",".join(columns)
+                data_rows = [",".join(map(str, row.values())) for row in rows]
+                return f"{header}\n" + "\n".join(data_rows)
+
+        except psycopg.Error as e:
+            log_error(f"Database error: {e}")
+            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}")
+            return f"An unexpected error occurred: {e}"
 
     def show_tables(self) -> str:
-        """Function to show tables in the database
+        """Lists all tables in the configured schema."""
 
-        :return: List of tables in the database
-        """
-        stmt = f"SELECT table_name FROM information_schema.tables WHERE table_schema = '{self.table_schema}';"
-        tables = self.run_query(stmt)
-        logger.debug(f"Tables: {tables}")
-        return tables
+        stmt = "SELECT table_name FROM information_schema.tables WHERE table_schema = %s;"
+        return self._execute_query(stmt, (self.table_schema,))
 
     def describe_table(self, table: str) -> str:
-        """Function to describe a table
-
-        :param table: Table to describe
-        :return: Description of the table
         """
-        stmt = f"SELECT column_name, data_type, character_maximum_length FROM information_schema.columns WHERE table_name = '{table}' AND table_schema = '{self.table_schema}';"
-        table_description = self.run_query(stmt)
+        Provides the schema (column name, data type, is nullable) for a given table.
 
-        logger.debug(f"Table description: {table_description}")
-        return f"{table}\n{table_description}"
+        Args:
+            table: The name of the table to describe.
 
-    def summarize_table(self, table: str) -> str:
-        """Function to compute a number of aggregates over a table.
-        The function launches a query that computes a number of aggregates over all columns,
-        including min, max, avg, std and approx_unique.
-
-        :param table: Table to summarize
-        :return: Summary of the table
+        Returns:
+            A string describing the table's columns and data types.
         """
-        stmt = f"""WITH column_stats AS (
-                SELECT
-                    column_name,
-                    data_type
-                FROM
-                    information_schema.columns
-                WHERE
-                    table_name = '{table}'
-                    AND table_schema = '{self.table_schema}'
-            )
-            SELECT
-                column_name,
-                data_type,
-                COUNT(COALESCE(column_name::text, '')) AS non_null_count,
-                COUNT(*) - COUNT(COALESCE(column_name::text, '')) AS null_count,
-                SUM(COALESCE(column_name::numeric, 0)) AS sum,
-                AVG(COALESCE(column_name::numeric, 0)) AS mean,
-                MIN(column_name::numeric) AS min,
-                MAX(column_name::numeric) AS max,
-                STDDEV(COALESCE(column_name::numeric, 0)) AS stddev
-            FROM
-                column_stats,
-                LATERAL (
-                    SELECT
-                        *
-                    FROM
-                        {table}
-                ) AS tbl
-            WHERE
-                data_type IN ('integer', 'numeric', 'real', 'double precision')
-            GROUP BY
-                column_name, data_type
-            UNION ALL
-            SELECT
-                column_name,
-                data_type,
-                COUNT(COALESCE(column_name::text, '')) AS non_null_count,
-                COUNT(*) - COUNT(COALESCE(column_name::text, '')) AS null_count,
-                NULL AS sum,
-                NULL AS mean,
-                NULL AS min,
-                NULL AS max,
-                NULL AS stddev
-            FROM
-                column_stats,
-                LATERAL (
-                    SELECT
-                        *
-                    FROM
-                        {table}
-                ) AS tbl
-            WHERE
-                data_type NOT IN ('integer', 'numeric', 'real', 'double precision')
-            GROUP BY
-                column_name, data_type;
+        stmt = """
+            SELECT column_name, data_type, is_nullable
+            FROM information_schema.columns
+            WHERE table_schema = %s AND table_name = %s;
         """
-        table_summary = self.run_query(stmt)
+        return self._execute_query(stmt, (self.table_schema, table))
 
-        logger.debug(f"Table summary: {table_summary}")
-        return table_summary
+    def summarize_table(self, table: str) -> str:
+        """
+        Computes and returns key summary statistics for a table's columns.
 
-    def inspect_query(self, query: str) -> str:
-        """Function to inspect a query and return the query plan. Always inspect your query before running them.
+        Args:
+            table: The name of the table to summarize.
 
-        :param query: Query to inspect
-        :return: Query plan
+        Returns:
+            A string containing a summary of the table.
         """
-        stmt = f"EXPLAIN {query};"
-        explain_plan = self.run_query(stmt)
+        try:
+            connection = self._ensure_connection()
+            with connection.cursor() as cursor:
+                # First, get column information using a parameterized query
+                schema_query = """
+                    SELECT column_name, data_type
+                    FROM information_schema.columns
+                    WHERE table_schema = %s AND table_name = %s;
+                """
+                cursor.execute(schema_query, (self.table_schema, table))
+                columns = cursor.fetchall()
+                if not columns:
+                    return f"Error: Table '{table}' not found in schema '{self.table_schema}'."
+
+                summary_parts = [f"Summary for table: {table}\n"]
+                table_identifier = sql.Identifier(self.table_schema, table)
+
+                for col in columns:
+                    col_name, data_type = col["column_name"], col["data_type"]
+                    col_identifier = sql.Identifier(col_name)
+
+                    query = None
+                    if any(
+                        t in data_type for t in ["integer", "numeric", "real", "double precision", "bigint", "smallint"]
+                    ):
+                        query = sql.SQL("""
+                            SELECT
+                                COUNT(*) AS total_rows,
+                                COUNT({col}) AS non_null_rows,
+                                MIN({col}) AS min,
+                                MAX({col}) AS max,
+                                AVG({col}) AS average,
+                                STDDEV({col}) AS std_deviation
+                            FROM {tbl};
+                        """).format(col=col_identifier, tbl=table_identifier)
+                    elif any(t in data_type for t in ["char", "text", "uuid"]):
+                        query = sql.SQL("""
+                            SELECT
+                                COUNT(*) AS total_rows,
+                                COUNT({col}) AS non_null_rows,
+                                COUNT(DISTINCT {col}) AS unique_values,
+                                AVG(LENGTH({col}::text)) as avg_length
+                            FROM {tbl};
+                        """).format(col=col_identifier, tbl=table_identifier)
+
+                    if query:
+                        cursor.execute(query)
+                        stats = cursor.fetchone()
+                        summary_parts.append(f"\n--- Column: {col_name} (Type: {data_type}) ---")
+                        if stats is not None:
+                            for key, value in stats.items():
+                                val_str = (
+                                    f"{value:.2f}" if isinstance(value, float) and value is not None else str(value)
+                                )
+                                summary_parts.append(f"  {key}: {val_str}")
+                        else:
+                            summary_parts.append("  No statistics available")
 
-        logger.debug(f"Explain plan: {explain_plan}")
-        return explain_plan
+                return "\n".join(summary_parts)
 
-    def export_table_to_path(self, table: str, path: Optional[str] = None) -> str:
-        """Save a table in CSV format.
-        If the path is provided, the table will be saved under that path.
-            Eg: If path is /tmp, the table will be saved as /tmp/table.csv
-        Otherwise it will be saved in the current directory
+        except psycopg.Error as e:
+            return f"Error summarizing table: {e}"
 
-        :param table: Table to export
-        :param path: Path to export to
-        :return: None
+    def inspect_query(self, query: str) -> str:
         """
+        Shows the execution plan for a SQL query (using EXPLAIN).
 
-        logger.debug(f"Exporting Table {table} as CSV to path {path}")
-        if path is None:
-            path = f"{table}.csv"
-        else:
-            path = f"{path}/{table}.csv"
-
-        export_statement = f"COPY {self.table_schema}.{table} TO '{path}' DELIMITER ',' CSV HEADER;"
-        result = self.run_query(export_statement)
-        logger.debug(f"Exported {table} to {path}/{table}")
-
-        return result
+        :param query: The SQL query to inspect.
+        :return: The query's execution plan.
+        """
+        return self._execute_query(f"EXPLAIN {query}")
 
-    def run_query(self, query: str) -> str:
-        """Function that runs a query and returns the result.
+    def export_table_to_path(self, table: str, path: str) -> str:
+        """
+        Exports a table's data to a local CSV file.
 
-        :param query: SQL query to run
-        :return: Result of the query
+        :param table: The name of the table to export.
+        :param path: The local file path to save the file.
+        :return: A confirmation message with the file path.
         """
+        log_debug(f"Exporting Table {table} as CSV to local path {path}")
 
-        # -*- Format the SQL Query
-        # Remove backticks
-        formatted_sql = query.replace("`", "")
-        # If there are multiple statements, only run the first one
-        formatted_sql = formatted_sql.split(";")[0]
+        table_identifier = sql.Identifier(self.table_schema, table)
+        stmt = sql.SQL("SELECT * FROM {tbl};").format(tbl=table_identifier)
 
         try:
-            logger.info(f"Running: {formatted_sql}")
-
-            cursor = self.connection.cursor()
-            cursor.execute(query)
-            query_result = cursor.fetchall()
-
-            result_output = "No output"
-            if query_result is not None:
-                try:
-                    results_as_python_objects = query_result
-                    result_rows = []
-                    for row in results_as_python_objects:
-                        if len(row) == 1:
-                            result_rows.append(str(row[0]))
-                        else:
-                            result_rows.append(",".join(str(x) for x in row))
+            connection = self._ensure_connection()
+            with connection.cursor() as cursor:
+                cursor.execute(stmt)
 
-                    result_data = "\n".join(result_rows)
-                    result_output = ",".join(query_result.columns) + "\n" + result_data
-                except AttributeError:
-                    result_output = str(query_result)
+                if cursor.description is None:
+                    return f"Error: Query returned no description for table '{table}'."
 
-            logger.debug(f"Query result: {result_output}")
+                columns = [desc[0] for desc in cursor.description]
 
-            return result_output
-        except Exception as e:
-            return str(e)
+                with open(path, "w", newline="", encoding="utf-8") as f:
+                    writer = csv.writer(f)
+                    writer.writerow(columns)
+                    writer.writerows(row.values() for row in cursor)
+
+            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:
+        """
+        Runs a read-only SQL query and returns the result.
+
+        :param query: The SQL query to run.
+        :return: The query result as a formatted string.
+        """
+        return self._execute_query(query)

agno/tools/pubmed.py

@@ -5,6 +5,7 @@ from xml.etree import ElementTree
 import httpx
 
 from agno.tools import Toolkit
+from agno.utils.log import log_debug
 
 
 class PubmedTools(Toolkit):
@@ -12,12 +13,20 @@ class PubmedTools(Toolkit):
         self,
         email: str = "your_email@example.com",
         max_results: Optional[int] = None,
+        results_expanded: bool = False,
+        enable_search_pubmed: bool = True,
+        all: bool = False,
+        **kwargs,
     ):
-        super().__init__(name="pubmed")
         self.max_results: Optional[int] = max_results
         self.email: str = email
+        self.results_expanded: bool = results_expanded
 
-        self.register(self.search_pubmed)
+        tools: List[Any] = []
+        if enable_search_pubmed or all:
+            tools.append(self.search_pubmed)
+
+        super().__init__(name="pubmed", tools=tools, **kwargs)
 
     def fetch_pubmed_ids(self, query: str, max_results: int, email: str) -> List[str]:
         url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi"
@@ -41,36 +50,139 @@ class PubmedTools(Toolkit):
     def parse_details(self, xml_root: ElementTree.Element) -> List[Dict[str, Any]]:
         articles = []
         for article in xml_root.findall(".//PubmedArticle"):
+            # Get existing fields
             pub_date = article.find(".//PubDate/Year")
             title = article.find(".//ArticleTitle")
-            abstract = article.find(".//AbstractText")
+
+            # Handle abstract sections with labels (methods, results, etc.)
+            abstract_sections = article.findall(".//AbstractText")
+            abstract_text = ""
+            if abstract_sections:
+                for section in abstract_sections:
+                    label = section.get("Label", "")
+                    if label:
+                        abstract_text += f"{label}: {section.text}\n\n"
+                    else:
+                        abstract_text += f"{section.text}\n\n"
+                abstract_text = abstract_text.strip()
+            else:
+                abstract_text = "No abstract available"
+
+            # Get first author
+            first_author_elem = article.find(".//AuthorList/Author[1]")
+            first_author = "Unknown"
+            if first_author_elem is not None:
+                last_name = first_author_elem.find("LastName")
+                fore_name = first_author_elem.find("ForeName")
+                if last_name is not None and fore_name is not None:
+                    first_author = f"{last_name.text}, {fore_name.text}"
+                elif last_name is not None and last_name.text:
+                    first_author = last_name.text
+
+            # Get DOI
+            doi_elem = article.find(".//ArticleIdList/ArticleId[@IdType='doi']")
+            doi = doi_elem.text if doi_elem is not None else "No DOI available"
+
+            # Get PMID for URL construction
+            pmid_elem = article.find(".//PMID")
+            pmid = pmid_elem.text if pmid_elem is not None else ""
+            pubmed_url = f"https://pubmed.ncbi.nlm.nih.gov/{pmid}/" if pmid else "No URL available"
+
+            # Check if full text is available via PMC
+            pmc_elem = article.find(".//ArticleIdList/ArticleId[@IdType='pmc']")
+            full_text_url = "Not available"
+            if pmc_elem is not None:
+                full_text_url = f"https://www.ncbi.nlm.nih.gov/pmc/articles/{pmc_elem.text}/"
+            elif doi_elem is not None:
+                full_text_url = f"https://doi.org/{doi}"
+
+            # Get keywords
+            keywords = []
+            for keyword in article.findall(".//KeywordList/Keyword"):
+                if keyword.text:
+                    keywords.append(keyword.text)
+
+            # Get MeSH terms (useful for understanding medical context)
+            mesh_terms = []
+            for mesh in article.findall(".//MeshHeading/DescriptorName"):
+                if mesh.text:
+                    mesh_terms.append(mesh.text)
+
+            # Get journal info
+            journal_elem = article.find(".//Journal/Title")
+            journal = journal_elem.text if journal_elem is not None else "Unknown Journal"
+
+            # Publication type (research article, review, etc.)
+            pub_types = []
+            for pub_type in article.findall(".//PublicationTypeList/PublicationType"):
+                if pub_type.text:
+                    pub_types.append(pub_type.text)
+
             articles.append(
                 {
-                    "Published": (pub_date.text if pub_date is not None else "No date available"),
+                    "Published": pub_date.text if pub_date is not None else "No date available",
                     "Title": title.text if title is not None else "No title available",
-                    "Summary": (abstract.text if abstract is not None else "No abstract available"),
+                    "Summary": abstract_text,
+                    "First_Author": first_author,
+                    "DOI": doi,
+                    "PubMed_URL": pubmed_url,
+                    "Full_Text_URL": full_text_url,
+                    "Keywords": ", ".join(keywords) if keywords else "No keywords available",
+                    "MeSH_Terms": ", ".join(mesh_terms) if mesh_terms else "No MeSH terms available",
+                    "Journal": journal,
+                    "Publication_Type": ", ".join(pub_types) if pub_types else "Not specified",
                 }
             )
+
         return articles
 
-    def search_pubmed(self, query: str, max_results: int = 10) -> str:
+    def search_pubmed(self, query: str, max_results: Optional[int] = 10) -> str:
         """Use this function to search PubMed for articles.
 
         Args:
             query (str): The search query.
-            max_results (int): The maximum number of results to return.
+            max_results (int): The maximum number of results to return (default 10).
 
         Returns:
             str: A JSON string containing the search results.
         """
         try:
-            ids = self.fetch_pubmed_ids(query, self.max_results or max_results, self.email)
+            log_debug(f"Searching PubMed for: {query}")
+            max_results = max_results or self.max_results or 10
+            ids = self.fetch_pubmed_ids(query, max_results, self.email)
             details_root = self.fetch_details(ids)
             articles = self.parse_details(details_root)
-            results = [
-                f"Published: {article.get('Published')}\nTitle: {article.get('Title')}\nSummary:\n{article.get('Summary')}"
-                for article in articles
-            ]
+
+            # Create result strings based on configured detail level
+            results = []
+            for article in articles:
+                if self.results_expanded:
+                    # Comprehensive format with all metadata
+                    article_text = (
+                        f"Published: {article.get('Published')}\n"
+                        f"Title: {article.get('Title')}\n"
+                        f"First Author: {article.get('First_Author')}\n"
+                        f"Journal: {article.get('Journal')}\n"
+                        f"Publication Type: {article.get('Publication_Type')}\n"
+                        f"DOI: {article.get('DOI')}\n"
+                        f"PubMed URL: {article.get('PubMed_URL')}\n"
+                        f"Full Text URL: {article.get('Full_Text_URL')}\n"
+                        f"Keywords: {article.get('Keywords')}\n"
+                        f"MeSH Terms: {article.get('MeSH_Terms')}\n"
+                        f"Summary:\n{article.get('Summary')}"
+                    )
+                else:
+                    # Concise format with just essential information
+                    summary = article.get("Summary", "")
+                    article_text = (
+                        f"Title: {article.get('Title')}\n"
+                        f"Published: {article.get('Published')}\n"
+                        f"Summary: {summary[:200]}..."
+                        if len(summary) > 200
+                        else f"Summary: {summary}"
+                    )
+                results.append(article_text)
+
             return json.dumps(results)
         except Exception as e:
-            return f"Cound not fetch articles. Error: {e}"
+            return f"Could not fetch articles. Error: {e}"