def __call__()

in nl2sql/tasks/eval_fix/core.py [0:0]

• 108 lines of code
• 9 McCabe index (conditional complexity)

    def __call__(self,
                 db: Database,
                 question: str,
                 query: str
                ) -> CoreEvalFixResult:
        """
        Runs the Core Eval and Fix Pipeline

        Args:
            db (Database): Name of the database.
            question (str): Natural language query
            query (str): Generated SQL query that throws error.

        Returns:
            CoreEvalFixResult: Fixed Result.
        """
        logger.info(f"Running {self.tasktype} ...")
        original_query = query
        modified_query = query
        intermediate_steps = []
        trials: List[str]= []
        trials.append(modified_query)

        @retry(
                stop=stop_after_attempt(self.num_retries)
        )
        def evaluate():
            trial_id = len(trials)
            sql = trials[-1]
            logger.info(f"Trial Id: {trial_id}")
            logger.info(f"Evaluating Generated Query: {sql}")
            try:
                _ = db.execute(sql) # type: ignore
            except DatabaseError as db_error:
                error_message = db_error.args[0].splitlines()[0]
                logger.warning(f"Evaluation Failed: "
                             f"{error_message}")
                logger.debug("Trying to fix the query ...")
                prompt_params = {
                    "question": question,
                    "query": question,
                    "input": question,
                    "generated_query": trials[-1],
                    "error_message": error_message,
                    "thoughts": [],
                    "answer": None,
                    "dialect": db.db.dialect,
                    "top_k": self.max_rows_limit,
                    "table_info": db.db.table_info,
                    "db_descriptor": {db.name: db.descriptor},
                    "table_name": ", ".join(db.db._usable_tables),
                    "table_names": list(db.db._usable_tables),
                }

                prompt_template = self.prompt.dialect_prompt_template_map.get(
                        db.db.dialect,
                        self.prompt.dialect_prompt_template_map.get("default"),
                    )
                if prompt_template is None:
                    raise ValueError(
                        f"No suitable / default prompt template found for {db.db.dialect}"
                    ) from db_error
                prepared_prompt = prompt_template.format(
                    **{
                        k: v
                        for k, v in prompt_params.items()
                        if k in prompt_template.input_variables
                    }
                )

                llm_response = self.llm.generate([prepared_prompt])
                logger.debug(
                    f"[{self.tasktype}] : Received LLM Response : {llm_response.json()}"
                )
                try:
                    raw_response = llm_response.generations[0][0].text.strip()
                except IndexError as exc:
                    raise ValueError(
                        f"Empty / Invalid Response received from LLM : {llm_response.json()}"
                    ) from exc

                parsed_response = (
                    self.prompt.parser.parse(raw_response)
                    if self.prompt.parser
                    else raw_response
                )
                processed_response = self.prompt.post_processor(parsed_response)
                intermediate_steps.append(
                    {
                        f"trial_{trial_id}": {
                                    "tasktype": self.tasktype,
                                    "prepared_prompt": prepared_prompt,
                                    "llm_response": llm_response.dict(),
                                    "raw_response": raw_response,
                                    "parsed_response": parsed_response,
                                    "processed_response": processed_response,
                                }
                    }
                )
                logger.info(f"New generated query: {processed_response}")
                trials.append(processed_response)
                trial_id += 1
                raise RuntimeError("Retry Evaluation...") from db_error
            else:
                logger.success("Generated Query successfully evaluated ...")
                return sql

        # Eval and Fix
        try:
            output =  evaluate()
            logger.success("EvalFix Successful.")
        except Exception as exc:
            logger.error(f"EvalFix Failed: {exc}")
            output = trials[-1]

        evalfixresult = CoreEvalFixResult(
            db_name=db.name,
            question=question,
            original_query=original_query,
            modified_query=output,
            intermediate_steps=intermediate_steps
        )
        logger.debug(
                    f"[{evalfixresult.resulttype}] : Final Query: {output}"
                )
        return evalfixresult