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

Crews
, ,
1

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?