Tool Methods in Python Class Missing self Argument When Used in Custom Agents

Question

I am working on a custom implementation of AI agents using Python, where the agents rely on langchain tools (methods) for interacting with a PostgreSQL database. The ToolWithConnection class initializes a database connection and defines tools like list_tables, tables_schema, check_sql, and execute_sql_to_df. These tools are later used by CrewAI agents to perform various tasks.

Here’s the relevant part of the code:

Connection Method:

def connection(x):
    postgres_schm = x["database_schm"]
    postgres_host = x["database_host"]
    postgres_port = x["database_port"]
    postgres_user = x["database_user"]
    postgres_pass = x["database_pass"]
    postgres_dryr = x["database_dryr"]

    if postgres_dryr == "true":
        postgres_pass_dec = postgres_pass        
    else:
        postgres_pass_dec = decrypt_string(postgres_pass)
        
    CONNECTION_STRING = (
        f"postgresql+psycopg2://{postgres_user}:{postgres_pass_dec}@{postgres_host}:{postgres_port}/{postgres_schm}"
    )
    
    db = SQLDatabase.from_uri(
        CONNECTION_STRING,
        schema=postgres_schm,
        view_support=True
    )
    return db

Tool Class:

class ToolWithConnection:
    def __init__(self, x):
        """Initialize with a connection based on the input 'x'."""
        self.db = connection(x)
        print('>> Initialize db connection ✔\n')

    @tool("list_tables")
    def list_tables(self) -> str:
        """List the available tables in the database."""
        return ListSQLDatabaseTool(db=self.db).invoke("")
        
    @tool("tables_schema")
    def tables_schema(self, tables: str) -> str:
        """Get the schema of specified tables."""
        tool = InfoSQLDatabaseTool(db=self.db)
        return tool.invoke(tables)

CrewAI Wrapper:

class CrewAIWrapper:
    def __init__(self, x):
        self.x = x
        self.llm_crew = LLM(model='azure/GPT4o')
        tool_instance = ToolWithConnection(self.x)

        self.sql_dev = Agent(
            role="Senior Database Developer",
            goal="Construct and execute PostgreSQL queries based on a query {query}",
            tools=[
                tool_instance.list_tables,  
                tool_instance.tables_schema,
            ],
            allow_delegation=False,
            memory=False,
            verbose=True,
        )

    def kickoff_crew(self):
        extract_data_task = Task(
            description="Extract data that is required for the query {query}.",
            agent=self.sql_dev,
        )

        my_crew = Crew(
            agents=[self.sql_dev],
            tasks=[extract_data_task],
            verbose=True,
        )

        try:
            crew_result = my_crew.kickoff(
                inputs={"query": self.x['question']}
            )
            return crew_result.raw
        except Exception as e:
            print(f"Error during task execution: {str(e)}")

Chain setup:

class CustomCrewAILLM(LLM):
    model: str
        
    def __init__(self, model='azure/GPT4o'):
        super().__init__(model=model)
        
    @property
    def _llm_type(self):
        return 'custom_crew_ai'
    
    def _call(self, obj: dict,) -> str:
        crew_ai_wrapper = CrewAIWrapper(obj, )
        result = crew_ai_wrapper.kickoff_crew()
        print(type(result))
        # print(result)
        return result
class InputType(BaseModel):
    crewai_input: Optional[str]
    db: Optional[SQLDatabase] = Field(default=None, exclude=True)  # Exclude from schema validation
    database_schm: str
    database_host: str
    database_port: str
    database_user: str
    database_pass: str
    database_name: str
    database_dryr: str
    class Config:
        arbitrary_types_allowed = True  # Allow non-Pydantic types like SQLDatabase

    # @field_validator("crewai_input")
    # def validate_input(cls, v):
    #     if not isinstance(v, str):
    #         raise ValidationError(f"Expected crewai_input to be a str, got {type(v)} instead.")
    #     return v

class OutputType(BaseModel):
    crewai_output: Optional[str]

    # @field_validator("crewai_output")
    # def validate_output(cls, v):
    #     if not isinstance(v, str):
    #         raise ValidationError(f"Expected crewai_output to be a str, got {type(v)} instead.")
    #     return v

chain = (
    RunnablePassthrough.assign(
        crewai_output=RunnableLambda(lambda x: CustomCrewAILLM()._call(x))
        )
        ).with_types(input_type=InputType,
                     output_type=OutputType
                     )

The Problem:

When I run the CrewAIWrapper with the agents and tools, I get the following error for the list_tables tool:

I encountered an error while trying to use the tool. This was the error: ToolWithConnection.list_tables() missing 1 required positional argument: 'self'.

Question:

• Why is the self argument missing when calling list_tables?

• How can I properly pass instance methods (like list_tables) from a class to agents without encountering this error?

Answer 1

I managed to find the solution by making each tool a separate class. Below is an example of how I structured the tools:

Tool Classes

Each tool is implemented as a standalone class, with attributes for its name, functionality, and description. This encapsulation provides clarity and modularity.

class ListTablesTool:
    def __init__(self, db_connection):
        self.name = "list_tables"
        self.func = self.list_tables
        self.description = "Lists the available tables in the database"
        self.db_connection = db_connection

    def list_tables(self) -> str:
        """List the available tables in the database"""
        return ListSQLDatabaseTool(db=self.db_connection).invoke("")


class TablesSchemaTool:
    def __init__(self, db_connection):
        self.name = "tables_schema"
        self.func = self.tables_schema
        self.description = "Gets schema and sample rows for specified tables"
        self.db_connection = db_connection

    def tables_schema(self, tables: str) -> str:
        """Fetch schema for the provided tables"""
        tool = InfoSQLDatabaseTool(db=self.db_connection)
        return tool.invoke(tables)

Initializing the Tools

In the main workflow or kickoff process, the tools are initialized like this:

self.list_tables = ListTablesTool(self.db)
self.tables_schema = TablesSchemaTool(self.db)

I hope this helps!