Method start_sample

src/server/tasks/dbbench/task.py:79–251

(self, index: int, session: Session)

        return list(range(len(self.dataset)))

    async def start_sample(self, index: int, session: Session) -> TaskSampleExecutionResult:
        self.env_controller.loop = asyncio.get_running_loop()
        if not self.env_controller_background_task:
            self.env_controller_background_task = asyncio.create_task(self.env_controller.background_task())
            weakref.finalize(self, self.env_controller_background_task.cancel)

        database: Optional[Database] = None
        try:
            entry = self.dataset[index][0]
            ground_truth = self.dataset[index][1]

            use_sqlite = entry.get("user_sqlite", False)
            if use_sqlite:
                db_dir = entry['create']['database']
                init_file = entry['create']['init']
                sqlite_path = os.path.join(self.db_root_dir, db_dir, init_file)
                database = SQLiteDatabase(sqlite_path)
                await database.initialize()
            else:
                init_sql = self._build_init_sql(entry)
                database = MySQLDatabase(self.env_controller)
                await database.initialize()
                await database.batch_execute(init_sql)

            session.inject(ChatCompletionSystemMessageParam(
                role='system',
                content=SYSTEM_PROMPT
            ))

            user_prompt = ""
            if "evidence" in entry and entry['evidence'] != "":
                user_prompt += "Evidence about the question: " + entry["evidence"] + "\n"
            if "add_description" in entry and entry['add_description'] != "":
                user_prompt += "Additional table information about the question: " + entry["add_description"] + "\n"
            user_prompt += "Question: " + entry["description"] + "\n"
            session.inject(ChatCompletionUserMessageParam(
                role='user',
                content=user_prompt
            ))

            for current_round in range(self.max_round):
                response = await session.action()

                tool_calls = []
                for message in response.messages:
                    tool_calls.extend(message.get('tool_calls', []) or [])

                if not tool_calls:
                    session.inject(ChatCompletionUserMessageParam(
                        role='user',
                        content='Internal error: No tool calls found despite finish reason.'
                    ))
                    continue

                for tool_call in tool_calls:
                    call_id = tool_call.get('id', '')
                    try:
                        function_name = tool_call.get('function', {}).get('name', '')