DataBeak

Instructions

# Remove multiple columns
remove_columns(ctx, ["col1", "col2", "col3"])

# Clean up after analysis
remove_columns(ctx, ["_temp", "_backup", "old_value"])

Implementation Reference

• src/databeak/servers/column_server.py:327-367 (handler)

  The core handler function implementing the remove_columns tool. Validates column existence, drops the specified columns from the session's dataframe, and returns a ColumnOperationResult with operation details.

  async def remove_columns(
      ctx: Annotated[Context, Field(description="FastMCP context for session access")],
      columns: Annotated[
          list[str],
          Field(description="List of column names to remove from the dataframe"),
      ],
  ) -> ColumnOperationResult:
      """Remove columns from the dataframe.
  
      Returns:
          ColumnOperationResult with removal details
  
      Examples:
          # Remove single column
          remove_columns(ctx, ["temp_column"])
  
          # Remove multiple columns
          remove_columns(ctx, ["col1", "col2", "col3"])
  
          # Clean up after analysis
          remove_columns(ctx, ["_temp", "_backup", "old_value"])
  
      """
      # Get session_id from FastMCP context
      session_id = ctx.session_id
      session, df = get_session_data(session_id)
  
      # Validate columns exist
      missing_cols = [col for col in columns if col not in df.columns]
      if missing_cols:
          raise ColumnNotFoundError(str(missing_cols[0]), df.columns.tolist())
  
      session.df = df.drop(columns=columns)
      # No longer recording operations (simplified MCP architecture)
  
      return ColumnOperationResult(
          operation="remove",
          rows_affected=len(df),
          columns_affected=columns,
      )

• src/databeak/servers/column_server.py:656-656 (registration)

  Registers the remove_columns handler as an MCP tool on the column_server FastMCP instance, specifying the tool name.

  column_server.tool(name="remove_columns")(remove_columns)