langroid 0.1.85__py3-none-any.whl → 0.1.219__py3-none-any.whl

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

@@ -1,10 +1,13 @@
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 from sqlalchemy import inspect, text
 from sqlalchemy.engine import Engine
 
 
-def extract_postgresql_descriptions(engine: Engine) -> Dict[str, Dict[str, Any]]:
+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.
 
@@ -13,6 +16,7 @@ def extract_postgresql_descriptions(engine: Engine) -> Dict[str, Dict[str, Any]]
 
     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
@@ -20,36 +24,53 @@ def extract_postgresql_descriptions(engine: Engine) -> Dict[str, Dict[str, Any]]
         column descriptions.
     """
     inspector = inspect(engine)
-    table_names: List[str] = inspector.get_table_names()
-
     result: Dict[str, Dict[str, Any]] = {}
 
-    with engine.connect() as conn:
-        for table in table_names:
-            table_comment = (
-                conn.execute(
-                    text(f"SELECT obj_description('{table}'::regclass)")
-                ).scalar()
-                or ""
-            )
+    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}"
 
-            columns = {}
-            col_data = inspector.get_columns(table)
-            for idx, col in enumerate(col_data, start=1):
-                col_comment = (
+                table_comment = (
                     conn.execute(
-                        text(f"SELECT col_description('{table}'::regclass, {idx})")
+                        text(f"SELECT obj_description('{table_name}'::regclass)")
                     ).scalar()
                     or ""
                 )
-                columns[col["name"]] = col_comment
 
-            result[table] = {"description": table_comment, "columns": columns}
+                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) -> Dict[str, Dict[str, Any]]:
+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
@@ -57,6 +78,7 @@ def extract_mysql_descriptions(engine: Engine) -> Dict[str, Dict[str, Any]]:
 
     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
@@ -64,31 +86,45 @@ def extract_mysql_descriptions(engine: Engine) -> Dict[str, Dict[str, Any]]:
         column descriptions.
     """
     inspector = inspect(engine)
-    table_names: List[str] = inspector.get_table_names()
-
     result: Dict[str, Dict[str, Any]] = {}
 
-    with engine.connect() as conn:
-        for table in table_names:
-            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}
-            )
-            table_comment = table_result.scalar() or ""
+    def gen_schema_descriptions(schema: Optional[str] = None) -> None:
+        table_names: List[str] = inspector.get_table_names(schema=schema)
 
-            columns = {}
-            for col in inspector.get_columns(table):
-                columns[col["name"]] = col.get("comment", "")
+        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}
 
-            result[table] = {"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) -> Dict[str, Dict[str, Any]]:
+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
@@ -96,6 +132,7 @@ def extract_default_descriptions(engine: Engine) -> Dict[str, Dict[str, Any]]:
 
     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
@@ -103,26 +140,36 @@ def extract_default_descriptions(engine: Engine) -> Dict[str, Dict[str, Any]]:
         empty column descriptions.
     """
     inspector = inspect(engine)
-    table_names: List[str] = inspector.get_table_names()
-
     result: Dict[str, Dict[str, Any]] = {}
 
-    for table in table_names:
-        columns = {}
-        for col in inspector.get_columns(table):
-            columns[col["name"]] = ""
+    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}
 
-        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) -> Dict[str, Dict[str, Any]]:
+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
@@ -133,4 +180,6 @@ def extract_schema_descriptions(engine: Engine) -> Dict[str, Dict[str, Any]]:
         "postgresql": extract_postgresql_descriptions,
         "mysql": extract_mysql_descriptions,
     }
-    return extractors.get(engine.dialect.name, extract_default_descriptions)(engine)
+    return extractors.get(engine.dialect.name, extract_default_descriptions)(
+        engine, multi_schema=multi_schema
+    )

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

@@ -1,10 +1,11 @@
-from typing import Dict, Union
+from typing import Dict, List, Union
 
 from sqlalchemy import MetaData
 
 
 def populate_metadata_with_schema_tools(
-    metadata: MetaData, info: Dict[str, Dict[str, Union[str, Dict[str, str]]]]
+    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
@@ -18,28 +19,35 @@ def populate_metadata_with_schema_tools(
     Returns:
         Dict[str, Dict[str, Any]]: A dictionary with table and context information.
     """
-
     db_info: Dict[str, Dict[str, Union[str, Dict[str, str]]]] = {}
 
-    # Create empty metadata dictionary with column datatypes
-    for table_name, table in metadata.tables.items():
-        # Populate tables with empty descriptions
-        db_info[str(table_name)] = {
-            "description": info[table_name]["description"] or "",
-            "columns": {},
-        }
-
-        for column in table.columns:
-            # Populate columns with datatype
-            db_info[str(table_name)]["columns"][str(column.name)] = (  # type: ignore
-                str(column.type)
-            )
+    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, info: Dict[str, Dict[str, Union[str, Dict[str, str]]]]
+    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.
@@ -51,11 +59,10 @@ def populate_metadata(
     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)
+    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():

langroid/agent/special/table_chat_agent.py

@@ -7,33 +7,37 @@ code to answer the query. The code is passed via the `run_code` tool/function-ca
 which is handled by the Agent's `run_code` method. This method executes/evaluates
 the code and returns the result as a string.
 """
+
 import io
 import logging
 import sys
-from typing import List, no_type_check
+from typing import List, Optional, 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
 from langroid.vector_store.base import VectorStoreConfig
 
 logger = logging.getLogger(__name__)
 
 console = Console()
 
-DEFAULT_TABLE_CHAT_SYSTEM_MESSAGE = """
+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 `run_code` tool/function-call to answer the question.
+will need to use the `run_code` tool/function-call to answer my questions.
 Here is a summary of the dataframe:
-{summary}
+{{summary}}
 Do not assume any columns other than those shown.
 In the code you submit to the `run_code` tool/function, 
 do not forget to include any necessary imports, such as `import pandas as pd`.
@@ -45,10 +49,12 @@ If you receive a null or other unexpected result, see if you have made an assump
 in your code, and try another way, or use `run_code` to explore the dataframe 
 before submitting your final code. 
 
-Once you have the answer to the question, say DONE and show me the answer.
-If you receive an error message, try using the `run_code` tool/function 
-again with the corrected code. 
+Once you have the answer to the question, possibly after a few steps,
+say {DONE} and show me the answer. If you receive an error message, 
+try using the `run_code` tool/function again with the corrected code. 
 
+VERY IMPORTANT: When using the `run_code` 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.
 """
 
@@ -72,7 +78,7 @@ def dataframe_summary(df: pd.DataFrame) -> str:
     )
 
     # Numerical data summary
-    num_summary = df.describe().applymap(lambda x: "{:.2f}".format(x))
+    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
@@ -101,7 +107,6 @@ def dataframe_summary(df: pd.DataFrame) -> str:
 class TableChatAgentConfig(ChatAgentConfig):
     system_message: str = DEFAULT_TABLE_CHAT_SYSTEM_MESSAGE
     user_message: None | str = None
-    max_context_tokens: int = 1000
     cache: bool = True  # cache results
     debug: bool = False
     stream: bool = True  # allow streaming where needed
@@ -125,6 +130,7 @@ class RunCodeTool(ToolMessage):
     purpose: str = """
             To run <code> on the dataframe 'df' and 
             return the results to answer a question.
+            IMPORTANT: ALL the code should be in the <code> field.
             """
     code: str
 
@@ -141,6 +147,8 @@ class TableChatAgent(ChatAgent):
     Agent for chatting with a collection of documents.
     """
 
+    sent_code: bool = False
+
     def __init__(self, config: TableChatAgentConfig):
         if isinstance(config.data, pd.DataFrame):
             df = config.data
@@ -165,6 +173,15 @@ class TableChatAgent(ChatAgent):
         # enable the agent to use and handle the RunCodeTool
         self.enable_message(RunCodeTool)
 
+    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_code = False
+        return response
+
     def run_code(self, msg: RunCodeTool) -> str:
         """
         Handle a RunCodeTool message by running the code and returning the result.
@@ -174,6 +191,7 @@ class TableChatAgent(ChatAgent):
         Returns:
             str: The result of running the code along with any print output.
         """
+        self.sent_code = True
         code = msg.code
         # Create a dictionary that maps 'df' to the actual DataFrame
         local_vars = {"df": self.df}
@@ -220,3 +238,19 @@ class TableChatAgent(ChatAgent):
             result = "No result"
         # Return the result
         return result
+
+    def handle_message_fallback(
+        self, msg: str | ChatDocument
+    ) -> str | ChatDocument | None:
+        """Handle scenario where LLM forgets to say DONE or forgets to use run_code"""
+        if isinstance(msg, ChatDocument) and msg.metadata.sender == lr.Entity.LLM:
+            if self.sent_code:
+                return DONE
+            else:
+                return """
+                    You forgot to use the `run_code` tool/function to find the answer.
+                    Try again using the `run_code` tool/function.
+                    Remember that ALL your code, including imports, 
+                    should be in the `code` field.
+                    """
+        return None