ai-data-science-team 0.0.0.9007__py3-none-any.whl → 0.0.0.9009__py3-none-any.whl

ai_data_science_team/templates/agent_templates.py

@@ -1,14 +1,225 @@
 from langchain_core.messages import AIMessage
 from langgraph.graph import StateGraph, END
 from langgraph.types import interrupt, Command
+from langgraph.graph.state import CompiledStateGraph
+
+from langchain_core.runnables import RunnableConfig
+from langgraph.pregel.types import StreamMode
 
 import pandas as pd
 import sqlalchemy as sql
+import json
 
-from typing import Any, Callable, Dict, Type, Optional
+from typing import Any, Callable, Dict, Type, Optional, Union, List
 
 from ai_data_science_team.tools.parsers import PythonOutputParser
-from ai_data_science_team.tools.regex import relocate_imports_inside_function, add_comments_to_top
+from ai_data_science_team.tools.regex import (
+    relocate_imports_inside_function, 
+    add_comments_to_top,
+    remove_consecutive_duplicates
+)
+
+from IPython.display import Image, display
+import pandas as pd
+
+class BaseAgent(CompiledStateGraph):
+    """
+    A generic base class for agents that interact with compiled state graphs.
+
+    Provides shared functionality for handling parameters, responses, and state
+    graph operations.
+    """
+
+    def __init__(self, **params):
+        """
+        Initialize the agent with provided parameters.
+
+        Parameters:
+            **params: Arbitrary keyword arguments representing the agent's parameters.
+        """
+        self._params = params
+        self._compiled_graph = self._make_compiled_graph()
+        self.response = None
+
+    def _make_compiled_graph(self):
+        """
+        Subclasses should override this method to create a specific compiled graph.
+        """
+        raise NotImplementedError("Subclasses must implement the `_make_compiled_graph` method.")
+
+    def update_params(self, **kwargs):
+        """
+        Update one or more parameters and rebuild the compiled graph.
+
+        Parameters:
+            **kwargs: Parameters to update.
+        """
+        self._params.update(kwargs)
+        self._compiled_graph = self._make_compiled_graph()
+
+    def __getattr__(self, name: str):
+        """
+        Delegate attribute access to the compiled graph if the attribute is not found.
+
+        Parameters:
+            name (str): The attribute name.
+
+        Returns:
+            Any: The attribute from the compiled graph.
+        """
+        return getattr(self._compiled_graph, name)
+
+    def invoke(
+        self, 
+        input: Union[dict[str, Any], Any], 
+        config: Optional[RunnableConfig] = None, 
+        **kwargs
+    ):
+        """
+        Wrapper for self._compiled_graph.invoke()
+
+        Parameters:
+            input: The input data for the graph. It can be a dictionary or any other type.
+            config: Optional. The configuration for the graph run.
+            **kwarg: Arguments to pass to self._compiled_graph.invoke()
+
+        Returns:
+            Any: The agent's response.
+        """
+        self.response = self._compiled_graph.invoke(input=input, config=config,**kwargs)
+        
+        if self.response.get("messages"):
+            self.response["messages"] = remove_consecutive_duplicates(self.response["messages"])
+        
+        return self.response
+    
+    def ainvoke(
+        self, 
+        input: Union[dict[str, Any], Any], 
+        config: Optional[RunnableConfig] = None, 
+        **kwargs
+    ):
+        """
+        Wrapper for self._compiled_graph.ainvoke()
+
+        Parameters:
+            input: The input data for the graph. It can be a dictionary or any other type.
+            config: Optional. The configuration for the graph run.
+            **kwarg: Arguments to pass to self._compiled_graph.ainvoke()
+            
+        Returns:
+            Any: The agent's response.
+        """
+        self.response = self._compiled_graph.ainvoke(input=input, config=config,**kwargs)
+        
+        if self.response.get("messages"):
+            self.response["messages"] = remove_consecutive_duplicates(self.response["messages"])
+        
+        return self.response
+    
+    def stream(
+        self,
+        input: dict[str, Any] | Any,
+        config: RunnableConfig | None = None,
+        stream_mode: StreamMode | list[StreamMode] | None = None, 
+        **kwargs
+    ):
+        """
+        Wrapper for self._compiled_graph.stream()
+
+        Parameters:
+            input: The input to the graph.
+            config: The configuration to use for the run.
+            stream_mode: The mode to stream output, defaults to self.stream_mode.
+                Options are 'values', 'updates', and 'debug'.
+                values: Emit the current values of the state for each step.
+                updates: Emit only the updates to the state for each step.
+                    Output is a dict with the node name as key and the updated values as value.
+                debug: Emit debug events for each step.
+            **kwarg: Arguments to pass to self._compiled_graph.stream()
+
+        Returns:
+            Any: The agent's response.
+        """
+        self.response = self._compiled_graph.stream(input=input, config=config, stream_mode=stream_mode, **kwargs)
+        
+        if self.response.get("messages"):
+            self.response["messages"] = remove_consecutive_duplicates(self.response["messages"])        
+        
+        return self.response
+    
+    def astream(
+        self,
+        input: dict[str, Any] | Any,
+        config: RunnableConfig | None = None,
+        stream_mode: StreamMode | list[StreamMode] | None = None, 
+        **kwargs
+    ):
+        """
+        Wrapper for self._compiled_graph.astream()
+
+        Parameters:
+            input: The input to the graph.
+            config: The configuration to use for the run.
+            stream_mode: The mode to stream output, defaults to self.stream_mode.
+                Options are 'values', 'updates', and 'debug'.
+                values: Emit the current values of the state for each step.
+                updates: Emit only the updates to the state for each step.
+                    Output is a dict with the node name as key and the updated values as value.
+                debug: Emit debug events for each step.
+            **kwarg: Arguments to pass to self._compiled_graph.astream()
+
+        Returns:
+            Any: The agent's response.
+        """
+        self.response = self._compiled_graph.astream(input=input, config=config, stream_mode=stream_mode, **kwargs)
+        
+        if self.response.get("messages"):
+            self.response["messages"] = remove_consecutive_duplicates(self.response["messages"])
+        
+        return self.response
+    
+    def get_state_keys(self):
+        """
+        Returns a list of keys that the state graph response contains.
+
+        Returns:
+            list: A list of keys in the response.
+        """
+        return list(self.get_output_jsonschema()['properties'].keys())
+
+    def get_state_properties(self):
+        """
+        Returns detailed properties of the state graph response.
+
+        Returns:
+            dict: The properties of the response.
+        """
+        return self.get_output_jsonschema()['properties']
+    
+    def get_response(self):
+        """
+        Returns the response generated by the agent.
+
+        Returns:
+            Any: The agent's response.
+        """
+        if self.response.get("messages"):
+            self.response["messages"] = remove_consecutive_duplicates(self.response["messages"])  
+        
+        return self.response
+
+    def show(self, xray: int = 0):
+        """
+        Displays the agent's state graph as a Mermaid diagram.
+
+        Parameters:
+            xray (int): If set to 1, displays subgraph levels. Defaults to 0.
+        """
+        display(Image(self.get_graph(xray=xray).draw_mermaid_png()))
+        
+
+
 
 def create_coding_agent_graph(
     GraphState: Type,
@@ -79,35 +290,37 @@ def create_coding_agent_graph(
 
     workflow = StateGraph(GraphState)
     
-    # Conditionally add the recommended-steps node
-    if not bypass_recommended_steps:
-        workflow.add_node(recommended_steps_node_name, node_functions[recommended_steps_node_name])
+    # * NODES
     
     # Always add create, execute, and fix nodes
     workflow.add_node(create_code_node_name, node_functions[create_code_node_name])
     workflow.add_node(execute_code_node_name, node_functions[execute_code_node_name])
     workflow.add_node(fix_code_node_name, node_functions[fix_code_node_name])
     
+    # Conditionally add the recommended-steps node
+    if not bypass_recommended_steps:
+        workflow.add_node(recommended_steps_node_name, node_functions[recommended_steps_node_name])
+    
+    # Conditionally add the human review node
+    if human_in_the_loop:
+        workflow.add_node(human_review_node_name, node_functions[human_review_node_name])
+    
     # Conditionally add the explanation node
     if not bypass_explain_code:
         workflow.add_node(explain_code_node_name, node_functions[explain_code_node_name])
     
+    # * EDGES
+    
     # Set the entry point
     entry_point = create_code_node_name if bypass_recommended_steps else recommended_steps_node_name
+    
     workflow.set_entry_point(entry_point)
     
-    # Add edges for recommended steps
     if not bypass_recommended_steps:
-        if human_in_the_loop:
-            workflow.add_edge(recommended_steps_node_name, human_review_node_name)
-        else:
-            workflow.add_edge(recommended_steps_node_name, create_code_node_name)
-    elif human_in_the_loop:
-        # Skip recommended steps but still include human review
-        workflow.add_edge(create_code_node_name, human_review_node_name)
+        workflow.add_edge(recommended_steps_node_name, create_code_node_name)
     
-    # Create -> Execute
     workflow.add_edge(create_code_node_name, execute_code_node_name)
+    workflow.add_edge(fix_code_node_name, execute_code_node_name)
     
     # Define a helper to check if we have an error & can still retry
     def error_and_can_retry(state):
@@ -117,39 +330,43 @@ def create_coding_agent_graph(
             and state.get(max_retries_key) is not None
             and state[retry_count_key] < state[max_retries_key]
         )
-    
-    # ---- Split into two branches for bypass_explain_code ----
-    if not bypass_explain_code:
-        # If we are NOT bypassing explain, the next node is fix_code if error,
-        # else explain_code. Then we wire explain_code -> END afterward.
+        
+    # If human in the loop, add a branch for human review
+    if human_in_the_loop:
         workflow.add_conditional_edges(
             execute_code_node_name,
-            lambda s: "fix_code" if error_and_can_retry(s) else "explain_code",
+            lambda s: "fix_code" if error_and_can_retry(s) else "human_review",
             {
+                "human_review": human_review_node_name,
                 "fix_code": fix_code_node_name,
-                "explain_code": explain_code_node_name,
             },
         )
-        # Fix code -> Execute again
-        workflow.add_edge(fix_code_node_name, execute_code_node_name)
-        # explain_code -> END
-        workflow.add_edge(explain_code_node_name, END)
     else:
-        # If we ARE bypassing explain_code, the next node is fix_code if error,
-        # else straight to END.
-        workflow.add_conditional_edges(
-            execute_code_node_name,
-            lambda s: "fix_code" if error_and_can_retry(s) else "END",
-            {
-                "fix_code": fix_code_node_name,
-                "END": END,
-            },
-        )
-        # Fix code -> Execute again
-        workflow.add_edge(fix_code_node_name, execute_code_node_name)
+        # If no human review, the next node is fix_code if error, else explain_code.
+        if not bypass_explain_code:
+            workflow.add_conditional_edges(
+                execute_code_node_name,
+                lambda s: "fix_code" if error_and_can_retry(s) else "explain_code",
+                {
+                    "fix_code": fix_code_node_name,
+                    "explain_code": explain_code_node_name,
+                },
+            )
+        else:
+            workflow.add_conditional_edges(
+                execute_code_node_name,
+                lambda s: "fix_code" if error_and_can_retry(s) else "END",
+                {
+                    "fix_code": fix_code_node_name,
+                    "END": END,
+                },
+            )
+            
+    if not bypass_explain_code:
+        workflow.add_edge(explain_code_node_name, END)
     
     # Finally, compile
-    if human_in_the_loop and checkpointer is not None:
+    if human_in_the_loop:
         app = workflow.compile(checkpointer=checkpointer)
     else:
         app = workflow.compile()
@@ -165,6 +382,8 @@ def node_func_human_review(
     no_goto: str,
     user_instructions_key: str = "user_instructions",
     recommended_steps_key: str = "recommended_steps",
+    code_snippet_key: str = "code_snippet",
+    code_type: str = "python"
 ) -> Command[str]:
     """
     A generic function to handle human review steps.
@@ -183,6 +402,10 @@ def node_func_human_review(
         The key in the state to store user instructions.
     recommended_steps_key : str, optional
         The key in the state to store recommended steps.    
+    code_snippet_key : str, optional
+        The key in the state to store the code snippet.
+    code_type : str, optional
+        The type of code snippet to display (e.g., "python").
     
     Returns
     -------
@@ -190,9 +413,11 @@ def node_func_human_review(
         A Command object directing the next state and updates to the state.    
     """
     print("    * HUMAN REVIEW")
+    
+    code_markdown=f"```{code_type}\n" + state.get(code_snippet_key)+"\n```"
 
     # Display instructions and get user response
-    user_input = interrupt(value=prompt_text.format(steps=state.get(recommended_steps_key, '')))
+    user_input = interrupt(value=prompt_text.format(steps=state.get(recommended_steps_key, '') + "\n\n" + code_markdown))
 
     # Decide next steps based on user input
     if user_input.strip().lower() == "yes":
@@ -200,11 +425,11 @@ def node_func_human_review(
         update = {}
     else:
         goto = no_goto
-        modifications = "Modifications: \n" + user_input
+        modifications = "User Has Requested Modifications To Previous Code: \n" + user_input
         if state.get(user_instructions_key) is None:
-            update = {user_instructions_key: modifications}
+            update = {user_instructions_key: modifications + "\n\nPrevious Code:\n" + code_markdown}
         else:
-            update = {user_instructions_key: state.get(user_instructions_key) + modifications}
+            update = {user_instructions_key: state.get(user_instructions_key) + modifications + "\n\nPrevious Code:\n" + code_markdown}
 
     return Command(goto=goto, update=update)
 
@@ -394,6 +619,7 @@ def node_func_fix_agent_code(
     retry_count_key: str = "retry_count",
     log: bool = False,
     file_path: str = "logs/agent_function.py",
+    function_name: str = "agent_function"
 ) -> dict:
     """
     Generic function to fix a given piece of agent code using an LLM and a prompt template.
@@ -420,6 +646,8 @@ def node_func_fix_agent_code(
         Whether to log the returned code to a file.
     file_path : str, optional
         The path to the file where the code will be logged.
+    function_name : str, optional
+        The name of the function in the code snippet that will be fixed.
     
     Returns
     -------
@@ -436,7 +664,8 @@ def node_func_fix_agent_code(
     # Format the prompt with the code snippet and the error
     prompt = prompt_template.format(
         code_snippet=code_snippet,
-        error=error_message
+        error=error_message,
+        function_name=function_name,
     )
     
     # Execute the prompt with the LLM
@@ -524,3 +753,50 @@ def node_func_explain_agent_code(
         # Return an error message if there was a problem with the code
         message = AIMessage(content=error_message)
         return {result_key: [message]}
+
+
+
+def node_func_report_agent_outputs(
+    state: Dict[str, Any],
+    keys_to_include: List[str],
+    result_key: str,
+    role: str,
+    custom_title: str = "Agent Output Summary"
+) -> Dict[str, Any]:
+    """
+    Gathers relevant data directly from the state (filtered by `keys_to_include`) 
+    and returns them as a structured message in `state[result_key]`.
+
+    No LLM is used.
+
+    Parameters
+    ----------
+    state : Dict[str, Any]
+        The current state dictionary holding all agent variables.
+    keys_to_include : List[str]
+        The list of keys in `state` to include in the output.
+    result_key : str
+        The key in `state` under which we'll store the final structured message.
+    role : str
+        The role that will be used in the final AIMessage (e.g., "DataCleaningAgent").
+    custom_title : str, optional
+        A title or heading for your report. Defaults to "Agent Output Summary".
+    """
+    print("    * REPORT AGENT OUTPUTS")
+
+    final_report = {"report_title": custom_title}
+
+    for key in keys_to_include:
+        final_report[key] = state.get(key, f"<{key}_not_found_in_state>")
+
+    # Wrap it in a list of messages (like the current "messages" pattern).
+    # You can serialize this dictionary as JSON or just cast it to string.
+    return {
+        result_key: [
+            AIMessage(
+                content=json.dumps(final_report, indent=2), 
+                role=role
+            )
+        ]
+    }
+

ai_data_science_team/tools/metadata.py

@@ -1,6 +1,7 @@
 import io
 import pandas as pd
 import sqlalchemy as sql
+from sqlalchemy import inspect
 from typing import Union, List, Dict
 
 def get_dataframe_summary(
@@ -139,8 +140,7 @@ def _summarize_dataframe(df: pd.DataFrame, dataset_name: str, n_sample=30, skip_
 
 
 
-def get_database_metadata(connection: Union[sql.engine.base.Connection, sql.engine.base.Engine],
-                          n_samples: int = 10) -> str:
+def get_database_metadata(connection, n_samples=10) -> dict:
     """
     Collects metadata and sample data from a database, with safe identifier quoting and
     basic dialect-aware row limiting. Prevents issues with spaces/reserved words in identifiers.
@@ -154,77 +154,109 @@ def get_database_metadata(connection: Union[sql.engine.base.Connection, sql.engi
 
     Returns
     -------
-    str
-        A formatted string with database metadata, including some sample data from each column.
+    dict
+        A dictionary with database metadata, including some sample data from each column.
     """
-
-    # If a connection is passed, use it; if an engine is passed, connect to it
     is_engine = isinstance(connection, sql.engine.base.Engine)
     conn = connection.connect() if is_engine else connection
 
-    output = []
+    metadata = {
+        "dialect": None,
+        "driver": None,
+        "connection_url": None,
+        "schemas": [],
+    }
+
     try:
-        # Grab the engine off the connection
         sql_engine = conn.engine
         dialect_name = sql_engine.dialect.name.lower()
 
-        output.append(f"Database Dialect: {sql_engine.dialect.name}")
-        output.append(f"Driver: {sql_engine.driver}")
-        output.append(f"Connection URL: {sql_engine.url}")
-
-        # Inspect the database
-        inspector = sql.inspect(sql_engine)
-        tables = inspector.get_table_names()
-        output.append(f"Tables: {tables}")
-        output.append(f"Schemas: {inspector.get_schema_names()}")
-
-        # Helper to build a dialect-specific limit clause
-        def build_query(col_name_quoted: str, table_name_quoted: str, n: int) -> str:
-            """
-            Returns a SQL query string to select N rows from the given column/table
-            across different dialects (SQLite, MySQL, Postgres, MSSQL, Oracle, etc.)
-            """
-            if "sqlite" in dialect_name or "mysql" in dialect_name or "postgres" in dialect_name:
-                # Common dialects supporting LIMIT
-                return f"SELECT {col_name_quoted} FROM {table_name_quoted} LIMIT {n}"
-            elif "mssql" in dialect_name:
-                # Microsoft SQL Server syntax
-                return f"SELECT TOP {n} {col_name_quoted} FROM {table_name_quoted}"
-            elif "oracle" in dialect_name:
-                # Oracle syntax
-                return f"SELECT {col_name_quoted} FROM {table_name_quoted} WHERE ROWNUM <= {n}"
-            else:
-                # Fallback
-                return f"SELECT {col_name_quoted} FROM {table_name_quoted} LIMIT {n}"
-
-        # Prepare for quoting
-        preparer = inspector.bind.dialect.identifier_preparer
-
-        # For each table, get columns and sample data
-        for table_name in tables:
-            output.append(f"\nTable: {table_name}")
-            # Properly quote the table name
-            table_name_quoted = preparer.quote_identifier(table_name)
-
-            for column in inspector.get_columns(table_name):
-                col_name = column["name"]
-                col_type = column["type"]
-                output.append(f"  Column: {col_name} Type: {col_type}")
+        metadata["dialect"] = sql_engine.dialect.name
+        metadata["driver"] = sql_engine.driver
+        metadata["connection_url"] = str(sql_engine.url)
 
-                # Properly quote the column name
-                col_name_quoted = preparer.quote_identifier(col_name)
-
-                # Build a dialect-aware query with safe quoting
-                query = build_query(col_name_quoted, table_name_quoted, n_samples)
-
-                # Read a few sample values
-                df = pd.read_sql(sql.text(query), conn)
-                first_values = df[col_name].tolist()
-                output.append(f"    First {n_samples} Values: {first_values}")
+        inspector = inspect(sql_engine)
+        preparer = inspector.bind.dialect.identifier_preparer
 
+        # For each schema
+        for schema_name in inspector.get_schema_names():
+            schema_obj = {
+                "schema_name": schema_name,
+                "tables": []
+            }
+
+            tables = inspector.get_table_names(schema=schema_name)
+            for table_name in tables:
+                table_info = {
+                    "table_name": table_name,
+                    "columns": [],
+                    "primary_key": [],
+                    "foreign_keys": [],
+                    "indexes": []
+                }
+                # Get columns
+                columns = inspector.get_columns(table_name, schema=schema_name)
+                for col in columns:
+                    col_name = col["name"]
+                    col_type = str(col["type"])
+                    table_name_quoted = f"{preparer.quote_identifier(schema_name)}.{preparer.quote_identifier(table_name)}"
+                    col_name_quoted = preparer.quote_identifier(col_name)
+
+                    # Build query for sample data
+                    query = build_query(col_name_quoted, table_name_quoted, n_samples, dialect_name)
+
+                    # Retrieve sample data
+                    try:
+                        df = pd.read_sql(query, conn)
+                        samples = df[col_name].head(n_samples).tolist()
+                    except Exception as e:
+                        samples = [f"Error retrieving data: {str(e)}"]
+
+                    table_info["columns"].append({
+                        "name": col_name,
+                        "type": col_type,
+                        "sample_values": samples
+                    })
+
+                # Primary keys
+                pk_constraint = inspector.get_pk_constraint(table_name, schema=schema_name)
+                table_info["primary_key"] = pk_constraint.get("constrained_columns", [])
+
+                # Foreign keys
+                fks = inspector.get_foreign_keys(table_name, schema=schema_name)
+                table_info["foreign_keys"] = [
+                    {
+                        "local_cols": fk["constrained_columns"],
+                        "referred_table": fk["referred_table"],
+                        "referred_cols": fk["referred_columns"]
+                    }
+                    for fk in fks
+                ]
+
+                # Indexes
+                idxs = inspector.get_indexes(table_name, schema=schema_name)
+                table_info["indexes"] = idxs
+
+                schema_obj["tables"].append(table_info)
+            
+            metadata["schemas"].append(schema_obj)
+    
     finally:
-        # Close connection if created inside the function
         if is_engine:
             conn.close()
 
-    return "\n".join(output)
+    return metadata
+
+def build_query(col_name_quoted: str, table_name_quoted: str, n: int, dialect_name: str) -> str:
+    # Example: expand your build_query to handle random sampling if possible
+    if "postgres" in dialect_name:
+        return f"SELECT {col_name_quoted} FROM {table_name_quoted} ORDER BY RANDOM() LIMIT {n}"
+    if "mysql" in dialect_name:
+        return f"SELECT {col_name_quoted} FROM {table_name_quoted} ORDER BY RAND() LIMIT {n}"
+    if "sqlite" in dialect_name:
+        return f"SELECT {col_name_quoted} FROM {table_name_quoted} ORDER BY RANDOM() LIMIT {n}"
+    if "mssql" in dialect_name:
+        return f"SELECT TOP {n} {col_name_quoted} FROM {table_name_quoted} ORDER BY NEWID()"
+    # Oracle or fallback
+    return f"SELECT {col_name_quoted} FROM {table_name_quoted} WHERE ROWNUM <= {n}"
+