File size: 4,886 Bytes
4962437
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
from __future__ import annotations

from typing import List, Optional

from langchain.chains.llm import LLMChain

from swarms.agents.utils.Agent import AgentOutputParser
from swarms.agents.utils.human_input import HumanInputRun
from swarms.memory.base_memory import BaseChatMessageHistory, ChatMessageHistory
from swarms.memory.document import Document
from swarms.models.base import AbstractModel
from swarms.models.prompts.agent_prompt_auto import (
    MessageFormatter,
    PromptConstructor,
)
from swarms.models.prompts.agent_prompt_generator import FINISH_NAME
from swarms.models.prompts.base import (
    AIMessage,
    HumanMessage,
    SystemMessage,
)
from swarms.tools.base import BaseTool


class Agent:
    """Base Agent class"""
    def __init__(
        self,
        ai_name: str,
        chain: LLMChain,
        memory,
        output_parser: AgentOutputParser,
        tools: List[BaseTool],
        feedback_tool: Optional[HumanInputRun] = None,
        chat_history_memory: Optional[BaseChatMessageHistory] = None,
    ):
        self.ai_name = ai_name
        self.chain = chain
        self.memory = memory
        self.next_action_count = 0
        self.output_parser = output_parser
        self.tools = tools
        self.feedback_tool = feedback_tool
        self.chat_history_memory = chat_history_memory or ChatMessageHistory()

    @classmethod
    def integrate(
        cls,
        ai_name: str,
        ai_role: str,
        memory,
        tools: List[BaseTool],
        llm: AbstractModel,
        human_in_the_loop: bool = False,
        output_parser: Optional[AgentOutputParser] = None,
        chat_history_memory: Optional[BaseChatMessageHistory] = None,
    ) -> Agent:
        prompt_constructor = PromptConstructor(ai_name=ai_name,
                                               ai_role=ai_role,
                                               tools=tools)
        message_formatter = MessageFormatter()
        human_feedback_tool = HumanInputRun() if human_in_the_loop else None
        chain = LLMChain(llm=llm, prompt_constructor=prompt_constructor, message_formatter=message_formatter)
        return cls(
            ai_name,
            memory,
            chain,
            output_parser or AgentOutputParser(),
            tools,
            feedback_tool=human_feedback_tool,
            chat_history_memory=chat_history_memory,
        )

    def run(self, goals: List[str]) -> str:
        user_input = (
            "Determine which next command to use, and respond using the format specified above:"
        )
        loop_count = 0
        while True:
            loop_count += 1

            # Send message to AI, get response
            assistant_reply = self.chain.run(
                goals=goals,
                messages=self.chat_history_memory.messages,
                memory=self.memory,
                user_input=user_input,
            )

            print(assistant_reply)
            self.chat_history_memory.add_message(HumanMessage(content=user_input))
            self.chat_history_memory.add_message(AIMessage(content=assistant_reply))

            # Get command name and arguments
            action = self.output_parser.parse(assistant_reply)
            tools = {t.name: t for t in self.tools}
            if action.name == FINISH_NAME:
                return action.args["response"]
            if action.name in tools:
                tool = tools[action.name]
                try:
                    observation = tool.run(action.args)
                except Exception as error:
                    observation = (
                        f"Validation Error in args: {str(error)}, args: {action.args}"
                    )
                except Exception as e:
                    observation = (
                        f"Error: {str(e)}, {type(e).__name__}, args: {action.args}"
                    )
                result = f"Command {tool.name} returned: {observation}"
            elif action.name == "ERROR":
                result = f"Error: {action.args}. "
            else:
                result = (
                    f"""Unknown command '{action.name}'. 
                    Please refer to the 'COMMANDS' list for available
                    commands and only respond in the specified JSON format."""
                )
            memory_to_add = (
                f"Assistant Reply: {assistant_reply} " f"\nResult: {result} "
            )
            if self.feedback_tool is not None:
                feedback = f"\n{self.feedback_tool.run('Input: ')}"
                if feedback in {"q", "stop"}:
                    print("EXITING")
                    return "EXITING"
                memory_to_add += feedback

            self.memory.add_documents([Document(page_content=memory_to_add)])
            self.chat_history_memory.add_message(SystemMessage(content=result))