src/research_assistant/utils/state_utils.py

from research_assistant.components.state import ResearchSummary


class SummaryStateUtils:

    def get_current_task(self, state: ResearchSummary):
        if results := state.get("results"):
            return None if len(results) == len(state["arguments"]) else len(results) + 1
        return 1

    def get_current_dependencies(self, state: ResearchSummary, step: int):
        return [
            state["results"].get(i, "") for i in state["dependencies"].get(step, [])
        ] or [state["article_text"]]

    def route(self, state: ResearchSummary):
        return "solve" if self.get_current_task(state) is None else "tool"

    def get_plan_results(self, state: ResearchSummary) -> str:
        results = state.get("results", {})
        plan_lines = [
            f"Plan: {plan}\n Answer = {results.get(i+1, '')}"
            for i, plan in enumerate(state["arguments"])
        ]
        return "\n".join(plan_lines)