Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Open
quartermaine opened this issue Jan 3, 2025 · 0 comments

Comments

@quartermaine
Copy link

quartermaine commented Jan 3, 2025

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?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

1 participant