def __call__()

in nl2sql_library/nl2sql/tasks/sql_generation/react.py [0:0]

• 56 lines of code
• 6 McCabe index (conditional complexity)

    def __call__(self, db: Database, question: str) -> ReactSqlGenratorResult:
        """
        Runs the SQL Generation pipeline
        """
        intermediate_steps: list[Any] = []
        logger.info(f"Running {self.tasktype} ...")
        agent = create_sql_agent(
            llm=self.llm,
            toolkit=SQLDatabaseToolkit(db=db.db, llm=self.llm),
            agent_type=self.agent_type,
            top_k=self.max_rows_limit,
            max_iterations=self.max_iterations,
            verbose=False,
            early_stopping_method="generate",
        )
        agent.return_intermediate_steps = True
        agent.handle_parsing_errors = True
        try:
            result = agent(question)
        except Exception as exc:  # pylint: disable=broad-exception-caught
            intermediate_steps.append(f"Exception in Agent.run : {exc}")
            query = None
        else:
            isteps = result.get("intermediate_steps")
            intermediate_steps.extend(
                [
                    {"input": step[0].to_json(), "output": step[1]}
                    if isinstance(step, tuple)
                    else step
                    for step in isteps
                ]
                if isteps
                else []
            )

            try:
                query = next(
                    map(
                        lambda x: x[0]
                        .tool_input.replace(";", "")
                        .replace("sql```", "")
                        .replace("```sql", "")
                        .replace("```", ""),
                        filter(
                            lambda x: x[0].tool
                            in ["sql_db_query", "sql_db_query_checker"],
                            reversed(result.get("intermediate_steps", [])),
                        ),
                    ),
                    None,
                )
            except Exception as exc:  # pylint: disable=broad-exception-caught
                query = None
                intermediate_steps.append(f"Exception in parsing query : {exc}")

        return ReactSqlGenratorResult(
            db_name=db.name,
            question=question,
            generated_query=query,
            intermediate_steps=intermediate_steps,
        )