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

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