langroid 0.33.4__py3-none-any.whl → 0.33.7__py3-none-any.whl

langroid/agent/special/sql/utils/description_extractors.py

@@ -0,0 +1,190 @@
+from typing import Any, Dict, List, Optional
+
+from langroid.exceptions import LangroidImportError
+
+try:
+    from sqlalchemy import inspect, text
+    from sqlalchemy.engine import Engine
+except ImportError as e:
+    raise LangroidImportError(extra="sql", error=str(e))
+
+
+def extract_postgresql_descriptions(
+    engine: Engine,
+    multi_schema: bool = False,
+) -> Dict[str, Dict[str, Any]]:
+    """
+    Extracts descriptions for tables and columns from a PostgreSQL database.
+
+    This method retrieves the descriptions of tables and their columns
+    from a PostgreSQL database using the provided SQLAlchemy engine.
+
+    Args:
+        engine (Engine): SQLAlchemy engine connected to a PostgreSQL database.
+        multi_schema (bool): Generate descriptions for all schemas in the database.
+
+    Returns:
+        Dict[str, Dict[str, Any]]: A dictionary mapping table names to a
+        dictionary containing the table description and a dictionary of
+        column descriptions.
+    """
+    inspector = inspect(engine)
+    result: Dict[str, Dict[str, Any]] = {}
+
+    def gen_schema_descriptions(schema: Optional[str] = None) -> None:
+        table_names: List[str] = inspector.get_table_names(schema=schema)
+        with engine.connect() as conn:
+            for table in table_names:
+                if schema is None:
+                    table_name = table
+                else:
+                    table_name = f"{schema}.{table}"
+
+                table_comment = (
+                    conn.execute(
+                        text(f"SELECT obj_description('{table_name}'::regclass)")
+                    ).scalar()
+                    or ""
+                )
+
+                columns = {}
+                col_data = inspector.get_columns(table, schema=schema)
+                for idx, col in enumerate(col_data, start=1):
+                    col_comment = (
+                        conn.execute(
+                            text(
+                                f"SELECT col_description('{table_name}'::regclass, "
+                                f"{idx})"
+                            )
+                        ).scalar()
+                        or ""
+                    )
+                    columns[col["name"]] = col_comment
+
+                result[table_name] = {"description": table_comment, "columns": columns}
+
+    if multi_schema:
+        for schema in inspector.get_schema_names():
+            gen_schema_descriptions(schema)
+    else:
+        gen_schema_descriptions()
+
+    return result
+
+
+def extract_mysql_descriptions(
+    engine: Engine,
+    multi_schema: bool = False,
+) -> Dict[str, Dict[str, Any]]:
+    """Extracts descriptions for tables and columns from a MySQL database.
+
+    This method retrieves the descriptions of tables and their columns
+    from a MySQL database using the provided SQLAlchemy engine.
+
+    Args:
+        engine (Engine): SQLAlchemy engine connected to a MySQL database.
+        multi_schema (bool): Generate descriptions for all schemas in the database.
+
+    Returns:
+        Dict[str, Dict[str, Any]]: A dictionary mapping table names to a
+        dictionary containing the table description and a dictionary of
+        column descriptions.
+    """
+    inspector = inspect(engine)
+    result: Dict[str, Dict[str, Any]] = {}
+
+    def gen_schema_descriptions(schema: Optional[str] = None) -> None:
+        table_names: List[str] = inspector.get_table_names(schema=schema)
+
+        with engine.connect() as conn:
+            for table in table_names:
+                if schema is None:
+                    table_name = table
+                else:
+                    table_name = f"{schema}.{table}"
+
+                query = text(
+                    "SELECT table_comment FROM information_schema.tables WHERE"
+                    " table_schema = :schema AND table_name = :table"
+                )
+                table_result = conn.execute(
+                    query, {"schema": engine.url.database, "table": table_name}
+                )
+                table_comment = table_result.scalar() or ""
+
+                columns = {}
+                for col in inspector.get_columns(table, schema=schema):
+                    columns[col["name"]] = col.get("comment", "")
+
+                result[table_name] = {"description": table_comment, "columns": columns}
+
+    if multi_schema:
+        for schema in inspector.get_schema_names():
+            gen_schema_descriptions(schema)
+    else:
+        gen_schema_descriptions()
+
+    return result
+
+
+def extract_default_descriptions(
+    engine: Engine, multi_schema: bool = False
+) -> Dict[str, Dict[str, Any]]:
+    """Extracts default descriptions for tables and columns from a database.
+
+    This method retrieves the table and column names from the given database
+    and associates empty descriptions with them.
+
+    Args:
+        engine (Engine): SQLAlchemy engine connected to a database.
+        multi_schema (bool): Generate descriptions for all schemas in the database.
+
+    Returns:
+        Dict[str, Dict[str, Any]]: A dictionary mapping table names to a
+        dictionary containing an empty table description and a dictionary of
+        empty column descriptions.
+    """
+    inspector = inspect(engine)
+    result: Dict[str, Dict[str, Any]] = {}
+
+    def gen_schema_descriptions(schema: Optional[str] = None) -> None:
+        table_names: List[str] = inspector.get_table_names(schema=schema)
+
+        for table in table_names:
+            columns = {}
+            for col in inspector.get_columns(table):
+                columns[col["name"]] = ""
+
+            result[table] = {"description": "", "columns": columns}
+
+    if multi_schema:
+        for schema in inspector.get_schema_names():
+            gen_schema_descriptions(schema)
+    else:
+        gen_schema_descriptions()
+
+    return result
+
+
+def extract_schema_descriptions(
+    engine: Engine, multi_schema: bool = False
+) -> Dict[str, Dict[str, Any]]:
+    """
+    Extracts the schema descriptions from the database connected to by the engine.
+
+    Args:
+        engine (Engine): SQLAlchemy engine instance.
+        multi_schema (bool): Generate descriptions for all schemas in the database.
+
+    Returns:
+        Dict[str, Dict[str, Any]]: A dictionary representation of table and column
+        descriptions.
+    """
+
+    extractors = {
+        "postgresql": extract_postgresql_descriptions,
+        "mysql": extract_mysql_descriptions,
+    }
+    return extractors.get(engine.dialect.name, extract_default_descriptions)(
+        engine, multi_schema=multi_schema
+    )

langroid/agent/special/sql/utils/populate_metadata.py

@@ -0,0 +1,85 @@
+from typing import Dict, List, Union
+
+from langroid.exceptions import LangroidImportError
+
+try:
+    from sqlalchemy import MetaData
+except ImportError as e:
+    raise LangroidImportError(extra="sql", error=str(e))
+
+
+def populate_metadata_with_schema_tools(
+    metadata: MetaData | List[MetaData],
+    info: Dict[str, Dict[str, Union[str, Dict[str, str]]]],
+) -> Dict[str, Dict[str, Union[str, Dict[str, str]]]]:
+    """
+    Extracts information from an SQLAlchemy database's metadata and combines it
+    with another dictionary with context descriptions.
+
+    Args:
+        metadata (MetaData): SQLAlchemy metadata object of the database.
+        info (Dict[str, Dict[str, Any]]): A dictionary with table and column
+                                             descriptions.
+
+    Returns:
+        Dict[str, Dict[str, Any]]: A dictionary with table and context information.
+    """
+    db_info: Dict[str, Dict[str, Union[str, Dict[str, str]]]] = {}
+
+    def populate_metadata(md: MetaData) -> None:
+        # Create empty metadata dictionary with column datatypes
+        for table_name, table in md.tables.items():
+            # Populate tables with empty descriptions
+            db_info[table_name] = {
+                "description": info[table_name]["description"] or "",
+                "columns": {},
+            }
+
+            for column in table.columns:
+                # Populate columns with datatype
+                db_info[table_name]["columns"][str(column.name)] = (  # type: ignore
+                    str(column.type)
+                )
+
+    if isinstance(metadata, list):
+        for md in metadata:
+            populate_metadata(md)
+    else:
+        populate_metadata(metadata)
+
+    return db_info
+
+
+def populate_metadata(
+    metadata: MetaData | List[MetaData],
+    info: Dict[str, Dict[str, Union[str, Dict[str, str]]]],
+) -> Dict[str, Dict[str, Union[str, Dict[str, str]]]]:
+    """
+    Populate metadata based on the provided database metadata and additional info.
+
+    Args:
+        metadata (MetaData): Metadata object from SQLAlchemy.
+        info (Dict): Additional information for database tables and columns.
+
+    Returns:
+        Dict: A dictionary containing populated metadata information.
+    """
+    # Fetch basic metadata info using available tools
+    db_info: Dict[str, Dict[str, Union[str, Dict[str, str]]]] = (
+        populate_metadata_with_schema_tools(metadata=metadata, info=info)
+    )
+
+    # Iterate over tables to update column metadata
+    for table_name in db_info.keys():
+        # Update only if additional info for the table exists
+        if table_name in info:
+            for column_name in db_info[table_name]["columns"]:
+                # Merge and update column description if available
+                if column_name in info[table_name]["columns"]:
+                    db_info[table_name]["columns"][column_name] = (  # type: ignore
+                        db_info[table_name]["columns"][column_name]  # type: ignore
+                        + "; "
+                        + info[table_name]["columns"][column_name]  # type: ignore
+                    )
+
+    return db_info

langroid/agent/special/sql/utils/system_message.py

@@ -0,0 +1,35 @@
+DEFAULT_SYS_MSG = """You are a savvy data scientist/database administrator, 
+with expertise in answering questions by querying a {dialect} database.
+You do not have access to the database 'db' directly, so you will need to use the 
+`run_query` tool/function-call to answer questions.
+
+The below JSON schema maps the SQL database structure. It outlines tables, each 
+with a description and columns. Each table is identified by a key, 
+and holds a description and a dictionary of columns, 
+with column names as keys and their descriptions as values. 
+{schema_dict}
+
+ONLY the tables and column names and tables specified above should be used in
+the generated queries. 
+You must be smart about using the right tables and columns based on the 
+english description. If you are thinking of using a table or column that 
+does not exist, you are probably on the wrong track, so you should try
+your best to answer based on an existing table or column.
+DO NOT assume any tables or columns other than those above."""
+
+SCHEMA_TOOLS_SYS_MSG = """You are a savvy data scientist/database administrator, 
+with expertise in answering questions by interacting with a SQL database.
+
+You will have to follow these steps to complete your job:
+1) Use the `get_table_names` tool/function-call to get a list of all possibly 
+relevant table names.
+2) Use the `get_table_schema` tool/function-call to get the schema of all 
+possibly relevant tables to identify possibly relevant columns. Only 
+call this method on potentially relevant tables.
+3) Use the `get_column_descriptions` tool/function-call to get more information 
+about any relevant columns.
+4) Write a {dialect} query and use `run_query` tool the Execute the SQL query 
+on the database to obtain the results.
+
+Do not make assumptions about the database schema before using the tools.
+Use the tool/functions to learn more about the database schema."""

langroid/agent/special/sql/utils/tools.py

@@ -0,0 +1,64 @@
+from typing import List, Tuple
+
+from langroid.agent.tool_message import ToolMessage
+
+
+class RunQueryTool(ToolMessage):
+    request: str = "run_query"
+    purpose: str = """
+            To run <query> on the database 'db' and 
+            return the results to answer a question.
+            """
+    query: str
+
+    @classmethod
+    def examples(cls) -> List["ToolMessage" | Tuple[str, "ToolMessage"]]:
+        return [
+            cls(
+                query="SELECT * FROM movies WHERE genre = 'comedy'",
+            ),
+            (
+                "Find all movies with a rating of 5",
+                cls(
+                    query="SELECT * FROM movies WHERE rating = 5",
+                ),
+            ),
+        ]
+
+
+class GetTableNamesTool(ToolMessage):
+    request: str = "get_table_names"
+    purpose: str = """
+            To retrieve the names of all <tables> in the database 'db'.
+            """
+
+
+class GetTableSchemaTool(ToolMessage):
+    request: str = "get_table_schema"
+    purpose: str = """
+            To retrieve the schema of all provided <tables> in the database 'db'.
+            """
+    tables: List[str]
+
+    @classmethod
+    def example(cls) -> "GetTableSchemaTool":
+        return cls(
+            tables=["employees", "departments", "sales"],
+        )
+
+
+class GetColumnDescriptionsTool(ToolMessage):
+    request: str = "get_column_descriptions"
+    purpose: str = """
+            To retrieve the description of one or more <columns> from the respective 
+            <table> in the database 'db'.
+            """
+    table: str
+    columns: str
+
+    @classmethod
+    def example(cls) -> "GetColumnDescriptionsTool":
+        return cls(
+            table="employees",
+            columns="name, department_id",
+        )

langroid/agent/special/table_chat_agent.py

@@ -0,0 +1,263 @@
+"""
+Agent that supports asking queries about a tabular dataset, internally
+represented as a Pandas dataframe. The `TableChatAgent` is configured with a
+dataset, which can be a Pandas df, file or URL. The delimiter/separator
+is auto-detected. In response to a user query, the Agent's LLM generates a Pandas
+expression (involving a dataframe `df`) to answer the query.
+The expression is passed via the `pandas_eval` tool/function-call,
+which is handled by the Agent's `pandas_eval` method. This method evaluates
+the expression and returns the result as a string.
+"""
+
+import io
+import logging
+import sys
+from typing import List, Optional, Tuple, no_type_check
+
+import numpy as np
+import pandas as pd
+from rich.console import Console
+
+import langroid as lr
+from langroid.agent import ChatDocument
+from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
+from langroid.agent.tool_message import ToolMessage
+from langroid.language_models.openai_gpt import OpenAIChatModel, OpenAIGPTConfig
+from langroid.parsing.table_loader import read_tabular_data
+from langroid.prompts.prompts_config import PromptsConfig
+from langroid.utils.constants import DONE, PASS
+from langroid.vector_store.base import VectorStoreConfig
+
+logger = logging.getLogger(__name__)
+
+console = Console()
+
+DEFAULT_TABLE_CHAT_SYSTEM_MESSAGE = f"""
+You are a savvy data scientist, with expertise in analyzing tabular datasets,
+using Python and the Pandas library for dataframe manipulation.
+Since you do not have access to the dataframe 'df', you
+will need to use the `pandas_eval` tool/function-call to answer my questions.
+Here is a summary of the dataframe:
+{{summary}}
+Do not assume any columns other than those shown.
+In the expression you submit to the `pandas_eval` tool/function, 
+you are allowed to use the variable 'df' to refer to the dataframe.
+
+Sometimes you may not be able to answer the question in a single call to `pandas_eval`,
+so you can use a series of calls to `pandas_eval` to build up the answer. 
+For example you may first want to know something about the possible values in a column.
+
+If you receive a null or other unexpected result, see if you have made an assumption
+in your code, and try another way, or use `pandas_eval` to explore the dataframe 
+before submitting your final code. 
+
+Once you have the answer to the question, possibly after a few steps,
+say {DONE} and PRESENT THE ANSWER TO ME; do not just say {DONE}.
+If you receive an error message, 
+try using the `pandas_eval` tool/function again with the corrected code. 
+
+VERY IMPORTANT: When using the `pandas_eval` tool/function, DO NOT EXPLAIN ANYTHING,
+   SIMPLY USE THE TOOL, with the CODE.
+Start by asking me what I want to know about the data.
+"""
+
+
+@no_type_check
+def dataframe_summary(df: pd.DataFrame) -> str:
+    """
+    Generate a structured summary for a pandas DataFrame containing numerical
+    and categorical values.
+
+    Args:
+        df (pd.DataFrame): The input DataFrame to summarize.
+
+    Returns:
+        str: A nicely structured and formatted summary string.
+    """
+
+    # Column names display
+    col_names_str = (
+        "COLUMN NAMES:\n" + " ".join([f"'{col}'" for col in df.columns]) + "\n\n"
+    )
+
+    # Numerical data summary
+    num_summary = df.describe().map(lambda x: "{:.2f}".format(x))
+    num_str = "Numerical Column Summary:\n" + num_summary.to_string() + "\n\n"
+
+    # Categorical data summary
+    cat_columns = df.select_dtypes(include=[np.object_]).columns
+    cat_summary_list = []
+
+    for col in cat_columns:
+        unique_values = df[col].unique()
+        if len(unique_values) < 10:
+            cat_summary_list.append(f"'{col}': {', '.join(map(str, unique_values))}")
+        else:
+            cat_summary_list.append(f"'{col}': {df[col].nunique()} unique values")
+
+    cat_str = "Categorical Column Summary:\n" + "\n".join(cat_summary_list) + "\n\n"
+
+    # Missing values summary
+    nan_summary = df.isnull().sum().rename("missing_values").to_frame()
+    nan_str = "Missing Values Column Summary:\n" + nan_summary.to_string() + "\n"
+
+    # Combine the summaries into one structured string
+    summary_str = col_names_str + num_str + cat_str + nan_str
+
+    return summary_str
+
+
+class TableChatAgentConfig(ChatAgentConfig):
+    system_message: str = DEFAULT_TABLE_CHAT_SYSTEM_MESSAGE
+    user_message: None | str = None
+    cache: bool = True  # cache results
+    debug: bool = False
+    stream: bool = True  # allow streaming where needed
+    data: str | pd.DataFrame  # data file, URL, or DataFrame
+    separator: None | str = None  # separator for data file
+    vecdb: None | VectorStoreConfig = None
+    llm: OpenAIGPTConfig = OpenAIGPTConfig(
+        type="openai",
+        chat_model=OpenAIChatModel.GPT4,
+        completion_model=OpenAIChatModel.GPT4,
+    )
+    prompts: PromptsConfig = PromptsConfig(
+        max_tokens=1000,
+    )
+
+
+class PandasEvalTool(ToolMessage):
+    """Tool/function to evaluate a pandas expression involving a dataframe `df`"""
+
+    request: str = "pandas_eval"
+    purpose: str = """
+            To eval a pandas <expression> on the dataframe 'df' and 
+            return the results to answer a question.
+            IMPORTANT: the <expression> field should be a valid pandas expression.
+            """
+    expression: str
+
+    @classmethod
+    def examples(cls) -> List["ToolMessage" | Tuple[str, "ToolMessage"]]:
+        return [
+            cls(expression="df.head()"),
+            cls(expression="df[(df['gender'] == 'Male')]['income'].mean()"),
+        ]
+
+    @classmethod
+    def instructions(cls) -> str:
+        return """
+            Use the `pandas_eval` tool/function to evaluate a pandas expression
+            involving the dataframe 'df' to answer the user's question.
+            """
+
+
+class TableChatAgent(ChatAgent):
+    """
+    Agent for chatting with a collection of documents.
+    """
+
+    sent_expression: bool = False
+
+    def __init__(self, config: TableChatAgentConfig):
+        if isinstance(config.data, pd.DataFrame):
+            df = config.data
+        else:
+            df = read_tabular_data(config.data, config.separator)
+
+        df.columns = df.columns.str.strip().str.replace(" +", "_", regex=True)
+
+        self.df = df
+        summary = dataframe_summary(df)
+        config.system_message = config.system_message.format(summary=summary)
+
+        super().__init__(config)
+        self.config: TableChatAgentConfig = config
+
+        logger.info(
+            f"""TableChatAgent initialized with dataframe of shape {self.df.shape}
+            and columns: 
+            {self.df.columns}
+            """
+        )
+        # enable the agent to use and handle the PandasEvalTool
+        self.enable_message(PandasEvalTool)
+
+    def user_response(
+        self,
+        msg: Optional[str | ChatDocument] = None,
+    ) -> Optional[ChatDocument]:
+        response = super().user_response(msg)
+        if response is not None and response.content != "":
+            self.sent_expression = False
+        return response
+
+    def pandas_eval(self, msg: PandasEvalTool) -> str:
+        """
+        Handle a PandasEvalTool message by evaluating the `expression` field
+            and returning the result.
+        Args:
+            msg (PandasEvalTool): The tool-message to handle.
+
+        Returns:
+            str: The result of running the code along with any print output.
+        """
+        self.sent_expression = True
+        exprn = msg.expression
+        local_vars = {"df": self.df}
+        # Create a string-based I/O stream
+        code_out = io.StringIO()
+
+        # Temporarily redirect standard output to our string-based I/O stream
+        sys.stdout = code_out
+
+        # Evaluate the last line and get the result
+        try:
+            eval_result = pd.eval(exprn, local_dict=local_vars)
+        except Exception as e:
+            eval_result = f"ERROR: {type(e)}: {e}"
+
+        if eval_result is None:
+            eval_result = ""
+
+        # Always restore the original standard output
+        sys.stdout = sys.__stdout__
+
+        # If df has been modified in-place, save the changes back to self.df
+        self.df = local_vars["df"]
+
+        # Get the resulting string from the I/O stream
+        print_result = code_out.getvalue() or ""
+        sep = "\n" if print_result else ""
+        # Combine the print and eval results
+        result = f"{print_result}{sep}{eval_result}"
+        if result == "":
+            result = "No result"
+        # Return the result
+        return result
+
+    def handle_message_fallback(
+        self, msg: str | ChatDocument
+    ) -> str | ChatDocument | None:
+        """Handle various LLM deviations"""
+        if isinstance(msg, ChatDocument) and msg.metadata.sender == lr.Entity.LLM:
+            if msg.content.strip() == DONE and self.sent_expression:
+                # LLM sent an expression (i.e. used the `pandas_eval` tool)
+                # but upon receiving the results, simply said DONE without
+                # narrating the result as instructed.
+                return """
+                    You forgot to PRESENT the answer to the user's query
+                    based on the results from `pandas_eval` tool.
+                """
+            if self.sent_expression:
+                # LLM forgot to say DONE
+                self.sent_expression = False
+                return DONE + " " + PASS
+            else:
+                # LLM forgot to use the `pandas_eval` tool
+                return """
+                    You forgot to use the `pandas_eval` tool/function 
+                    to find the answer.
+                    Try again using the `pandas_eval` tool/function.
+                    """
+        return None