DxEvolve/agents/executor.py at main · RUCAIBox/DxEvolve · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
import pydantic
from langchain.agents.agent import *
AgentStep.__config__.arbitrary_types_allowed = True
AgentExecutor.__config__.arbitrary_types_allowed = True
AgentAction.__config__.arbitrary_types_allowed = True
IS_V2 = pydantic.__version__.startswith("2.")


_original_init = AgentStep.__init__

def _patched_init(self, *args, **kwargs):
    if "action" in kwargs:
        action = kwargs["action"]
        if not isinstance(action, dict) and hasattr(action, "tool"):
            kwargs["action"] = {
                "tool": getattr(action, "tool", ""),
                "tool_input": getattr(action, "tool_input", ""),
                "log": getattr(action, "log", "")
            }
    _original_init(self, *args, **kwargs)

AgentStep.__init__ = _patched_init


class CustomAgentExecutor(AgentExecutor):
    summarize_observation: bool = True

    def _take_next_step(
        self,
        name_to_tool_map: Dict[str, Any],
        color_mapping: Dict[str, str],
        inputs: Dict[str, Any],
        intermediate_steps: List[Tuple[AgentAction, str]],
        run_manager: Optional[Any] = None,
    ) -> Union[AgentFinish, List[Tuple[AgentAction, str]]]:

        result = super()._take_next_step(
            name_to_tool_map=name_to_tool_map,
            color_mapping=color_mapping,
            inputs=inputs,
            intermediate_steps=intermediate_steps,
            run_manager=run_manager,
        )

        if (not self.summarize_observation) or isinstance(result, AgentFinish):
            return result

        if not isinstance(result, list):
            return result

        input_text = inputs.get("input", "")
        new_result = []

        for item in result:
            if isinstance(item, AgentStep):
                if ("not available" in str(item.observation).lower()) or ("no results available" in str(item.observation).lower()):
                    summarized = str(item.observation)
                    new_result.append(AgentStep(action=item.action, observation=summarized))
                    continue
                summarized = self.agent.summarize_tool_observation(
                    inputs=inputs,
                    input_text=input_text,
                    intermediate_steps=intermediate_steps,
                    action=item.action,
                    observation=str(item.observation),
                )
                new_result.append(AgentStep(action=item.action, observation=summarized))
            elif isinstance(item, tuple) and len(item) == 2:
                action, obs = item
                summarized = self.agent.summarize_tool_observation(
                    inputs=inputs,
                    input_text=input_text,
                    intermediate_steps=intermediate_steps,
                    action=action,
                    observation=str(obs),
                )
                new_result.append((action, summarized))
            else:
                new_result.append(item)

        return new_result