def get_result()

in databao/executors/lighthouse/graph.py [0:0]

• 45 lines of code
• 7 McCabe index (conditional complexity)

    def get_result(self, state: AgentState) -> ExecutionResult:
        last_ai_message = None
        for m in reversed(state["messages"]):
            if isinstance(m, AIMessage):
                last_ai_message = m
                break
        if last_ai_message is None:
            raise RuntimeError("No AI message found in message log")
        if len(last_ai_message.tool_calls) == 0:
            # Sometimes models don't call the submit_result tool, but we still want to return some dataframe.
            sql = state.get("sql", "")
            df = state.get("df")  # Latest df result (usually from run_sql_query)
            visualization_prompt = state.get("visualization_prompt")
            result = ExecutionResult(
                text=last_ai_message.text(),
                df=df,
                code=sql,
                meta={
                    "visualization_prompt": visualization_prompt,
                    "messages": state["messages"],
                    "submit_called": False,
                },
            )
        elif len(last_ai_message.tool_calls) > 1:
            raise RuntimeError("Expected exactly one tool call in AI message")
        elif last_ai_message.tool_calls[0]["name"] != "submit_result":
            raise RuntimeError(
                f"Expected submit_result tool call in AI message, got {last_ai_message.tool_calls[0]['name']}"
            )
        else:
            sql = state.get("sql", "")
            df = state.get("df")
            tool_call = last_ai_message.tool_calls[0]
            text = tool_call["args"]["result_description"]
            visualization_prompt = state.get("visualization_prompt", "")
            result = ExecutionResult(
                text=text,
                df=df,
                code=sql,
                meta={
                    "visualization_prompt": visualization_prompt,
                    "messages": state["messages"],
                    "submit_called": True,
                },
            )
        return result