def write_state_chunk()

in databao/executors/frontend/text_frontend.py [0:0]

• 22 lines of code
• 16 McCabe index (conditional complexity)

    def write_state_chunk(self, state_chunk: dict[str, Any]) -> None:
        """The state chunk is assumed to contain a "messages" key."""
        if self._is_tool_calling:
            self.write("\n```\n\n")  # Close code block
            self._is_tool_calling = False

        # Loop through new messages only.
        # We could either force the caller of the frontend to provide new messages only,
        # but for ease of use we assume the state contains a list of messages and do it here.
        messages: list[BaseMessage] = state_chunk.get("messages", [])
        new_messages = messages[self._message_count :]
        self._message_count += len(new_messages)

        for message in new_messages:
            if isinstance(message, ToolMessage):
                tool_call = get_tool_call(messages, message)
                tool_name = tool_call["name"] if tool_call is not None else "unknown"
                self.write(f"\n[tool_call_output: '{tool_name}']")
                self.write(f"\n```\n{message.text.strip()}\n```\n\n")
                if message.artifact is not None and isinstance(message.artifact, dict):
                    for art_name, art_value in message.artifact.items():
                        if isinstance(art_value, pd.DataFrame):
                            self.write_dataframe(art_value, name=art_name)
            elif self._pretty_sql and isinstance(message, AIMessage):
                # During tool calling we show raw JSON chunks, but for SQL we also want pretty formatting.
                for tool_call in message.tool_calls:
                    sql = get_tool_call_sql(tool_call)
                    if sql is not None:
                        self.write(f"\n```sql\n{sql.strip()}\n```\n\n")