[Part 2] 2시간 안에 업무 자동화 AI Agent 만들기: LLM 체인 직접 써보기

Part 1에 이이서 오늘은 Part 2를 올립니다.

Part 1: 전체 프로그램 셋업, OpenAI API 연결, Vector Store 연결

Part 2: 커스텀 LLM 체인 직접 만들고나서 실행하기

Part 3: Agent가 나 대신 실제 구글 검색하는 기능 만들기


LLM 체인 만들기

Part 1에서 기본적인 것들이 연결됐으니 이제 실제 LLM 체인을 만들어볼거에요.

일단 저희가 만들 AI agent가 어떻게 작동되냐면, 먼저 유저가 “페이스북 시장 조사 좀 해줘”라고 하면 이 업무를 하기 위해 해야할 태스크들을 리스트업합니다. 그 다음 태스크별로 실행을 하고 결과값을 생성합니다.

이를 하기 위해서는 3개의 체인이 활용됩니다.

  1. 태스크를 생성하고 리스트에 추가하는 체인

  2. 태스크 중 중요도 높은 태스크부터 처리하게끔 관리하는 체인

  3. 각 태스크를 실제로 실행하는 체인


First LLM Chain

유저 프롬트에 맞게 태스크를 생성하고 리스트에 추가하는 첫번째 체인 코드입니다.

class TaskCreationChain(LLMChain):
    """Chain to generates tasks."""

    @classmethod
    def from_llm(cls, llm: BaseLLM, verbose: bool = True) -> LLMChain:
        """Get the response parser."""
        task_creation_template = (
            "You are a task creation AI that uses the result of an execution agent"
            " to create new tasks with the following objective: {objective},"
            " The last completed task has the result: {result}."
            " This result was based on this task description: {task_description}."
            " These are incomplete tasks: {incomplete_tasks}."
            " Based on the result, create new tasks to be completed"
            " by the AI system that do not overlap with incomplete tasks."
            " Return the tasks as an array."
        )
        prompt = PromptTemplate(
            template=task_creation_template,
            input_variables=[
                "result",
                "task_description",
                "incomplete_tasks",
                "objective",
            ],
        )
        return cls(prompt=prompt, llm=llm, verbose=verbose)

여기서 코드를 일일이 이해할 필요는 없고, 크게 어떻게 작동되는지만 알면 됩니다.

일단 맨 처음에 TaskCreationChain이라는 클래스를 생성했고, 이 클래스를 호출했을때 업무 자동화를 하기 위한 명령 템플릿이 task_creation_template에 작성돼있습니다.

말그대로 LLM 체인에게 이렇게 해라 라고 명시하는 겁니다.

task_creation_template = ( … ) 에 작성돼있는 텍스트를 보면, “너는 태스크를 생성하는 AI이고 태스크 별로 실행까지 하는 에이전트야. 내가 입력하는 {objective}에 맞게 태스크를 생성하고, 가장 최근에 실행된 태스크의 결과값은 {result}에 저장하고, 태스크에 관한 설명은 {task_description}에 저장될거야. 아직 실행되지 않은 태스크들은 {incomplete_tasks}에 저장해. 결과값을 통해 아직 실행되지 않은 태스크들을 실행하고 결과값들은 배열 형태로 정리해줘”라고 써져 있습니다.

이런 프롬트 템플릿은 제가 원하는대로 언제든지 바꿀 수 있습니다. 직접 코드로 개발을 하지 않아도 정말 텍스트로 AI에게 명령하는거죠.

다음 코드 줄에서는 이 task_creation_template을 prompt라는 곳에 저장하고, 인풋은 “result”, “task_description”, “incomplete_tasks”, “objective”이라는 것을 확인할 수 있습니다.

코드를 복붙하고 나면 이렇게 생겼습니다.

Second LLM Chain

다음은 태스크 중 중요도 높은 태스크들을 먼저 실행하게끔 관리하는 두번째 체인 코드입니다.

class TaskPrioritizationChain(LLMChain):
    """Chain to prioritize tasks."""

    @classmethod
    def from_llm(cls, llm: BaseLLM, verbose: bool = True) -> LLMChain:
        """Get the response parser."""
        task_prioritization_template = (
            "You are a task prioritization AI tasked with cleaning the formatting of and reprioritizing"
            " the following tasks: {task_names}."
            " Consider the ultimate objective of your team: {objective}."
            " Do not remove any tasks. Return the result as a numbered list, like:"
            " #. First task"
            " #. Second task"
            " Start the task list with number {next_task_id}."
        )
        prompt = PromptTemplate(
            template=task_prioritization_template,
            input_variables=["task_names", "next_task_id", "objective"],
        )
        return cls(prompt=prompt, llm=llm, verbose=verbose)

첫번째 체인과 아주 비슷한 형식이죠?

이번에 task_prioritization_template에는 “너는 {task_names} 태스크들 중에서 중요도 높은 태스크부터 실행되도록 순번을 관리하는 AI야. 최종 {objective} 목표를 염두하면서 순번 관리를 해줘. 태스크들 중 어떤것도 지우지 말고 결과값은 리스트 형태로 보여줘: “#. 첫번째 태스크, #. 두번째 태스크”. 태스크 리스트는 태스크 숫자 아이디로 {next_task_id} 시작해줘.” 라고 써져 있습니다.

첫번째 체인 코드 밑에 또 복붙하세요.

Third LLM Chain

이제 마지막 체인, 실제 태스크를 실행하게끔 하는 체인 코드입니다.

class ExecutionChain(LLMChain):
    """Chain to execute tasks."""

    @classmethod
    def from_llm(cls, llm: BaseLLM, verbose: bool = True) -> LLMChain:
        """Get the response parser."""
        execution_template = (
            "You are an AI who performs one task based on the following objective: {objective}."
            " Take into account these previously completed tasks: {context}."
            " Your task: {task}."
            " Response:"
        )
        prompt = PromptTemplate(
            template=execution_template,
            input_variables=["objective", "context", "task"],
        )
        return cls(prompt=prompt, llm=llm, verbose=verbose)

여기서 템플릿에는 “너는 큰 목표 {objective}에 맞게 한 개의 태스크를 실행하는 AI야. 전에 실행 완료된 태스크들의 {context}를 생각하면서 태스크를 실행하도록 해. 너의 태스크: {task}.”라고 써져있습니다.

마찬가지로 코드를 두번째 체인 코드 밑에 복붙하세요.

BabyAGI Controller

이제 BabyAGI 컨트롤러를 만들 차례입니다. 위에서 만든 세 개의 체인들을 활용해서 유저가 인풋을 주면 그 인풋에 맞게 자동으로 태스크들을 생성하고 처리하게끔 만들어야 합니다.

def get_next_task(
    task_creation_chain: LLMChain,
    result: Dict,
    task_description: str,
    task_list: List[str],
    objective: str,
) -> List[Dict]:
    """Get the next task."""
    incomplete_tasks = ", ".join(task_list)
    response = task_creation_chain.run(
        result=result,
        task_description=task_description,
        incomplete_tasks=incomplete_tasks,
        objective=objective,
    )
    new_tasks = response.split("\\n")
    return [{"task_name": task_name} for task_name in new_tasks if task_name.strip()]

위 코드는 1) 처음에 에이전트가 해야할 태스크 리스트업을 했을때 첫 태스크를 완료하고 난 뒤 아직 실행되지 않은 다음 태스크를 불러오고, 2) 나머지 미완료 태스크들을 다시 생성 + 리스트업하기 위해 첫번째 태스크 생성 체인을 불러오는 코드입니다.

def prioritize_tasks(
    task_prioritization_chain: LLMChain,
    this_task_id: int,
    task_list: List[Dict],
    objective: str,
) -> List[Dict]:
    """Prioritize tasks."""
    task_names = [t["task_name"] for t in task_list]
    next_task_id = int(this_task_id) + 1
    response = task_prioritization_chain.run(
        task_names=task_names, next_task_id=next_task_id, objective=objective
    )
    new_tasks = response.split("\\n")
    prioritized_task_list = []
    for task_string in new_tasks:
        if not task_string.strip():
            continue
        task_parts = task_string.strip().split(".", 1)
        if len(task_parts) == 2:
            task_id = task_parts[0].strip()
            task_name = task_parts[1].strip()
            prioritized_task_list.append({"task_id": task_id, "task_name": task_name})
    return prioritized_task_list

위 코드는 저희가 만든 두번째 태스크 순번 관리 체인을 실행하기 위한 코드입니다.

def _get_top_tasks(vectorstore, query: str, k: int) -> List[str]:
    """Get the top k tasks based on the query."""
    results = vectorstore.similarity_search_with_score(query, k=k)
    if not results:
        return []
    sorted_results, _ = zip(*sorted(results, key=lambda x: x[1], reverse=True))
    return [str(item.metadata["task"]) for item in sorted_results]

def execute_task(
    vectorstore, execution_chain: LLMChain, objective: str, task: str, k: int = 5
) -> str:
    """Execute a task."""
    context = _get_top_tasks(vectorstore, query=objective, k=k)
    return execution_chain.run(objective=objective, context=context, task=task)

마지막으로 _get_top_tasks는 가장 첫번째 순서인 태스크를 불러오는 거고, execute_task는 저희가 세번째로 만든 태스크 실행하는 체인을 불러와서 첫번째 순서 태스크부터 실행하게끔 하는 코드입니다.

이 모든 코드를 순서대로 복붙하세요. 그럼 아래 이미지처럼 생겼습니다.

이제 정말 거의 다왔습니다.

여태까지 만든 모든 코드 블록들을 저희 AI Agent가 적절하게 쓸 수 있도록 짜집기 해줄거에요.

여기서 class BabyAGI 코드 블록이 실질적으로 AI Agent의 역할을 한다고 생각하시면 됩니다.

class BabyAGI(Chain, BaseModel):
    """Controller model for the BabyAGI agent."""

    task_list: deque = Field(default_factory=deque)
    task_creation_chain: TaskCreationChain = Field(...)
    task_prioritization_chain: TaskPrioritizationChain = Field(...)
    execution_chain: ExecutionChain = Field(...)
    task_id_counter: int = Field(1)
    vectorstore: VectorStore = Field(init=False)
    max_iterations: Optional[int] = None

    class Config:
        """Configuration for this pydantic object."""

        arbitrary_types_allowed = True

    def add_task(self, task: Dict):
        self.task_list.append(task)

    def print_task_list(self):
        print("\033[95m\033[1m" + "\n*****TASK LIST*****\n" + "\033[0m\033[0m")
        for t in self.task_list:
            print(str(t["task_id"]) + ": " + t["task_name"])

    def print_next_task(self, task: Dict):
        print("\033[92m\033[1m" + "\n*****NEXT TASK*****\n" + "\033[0m\033[0m")
        print(str(task["task_id"]) + ": " + task["task_name"])

    def print_task_result(self, result: str):
        print("\033[93m\033[1m" + "\n*****TASK RESULT*****\n" + "\033[0m\033[0m")
        print(result)

    @property
    def input_keys(self) -> List[str]:
        return ["objective"]

    @property
    def output_keys(self) -> List[str]:
        return []

    def _call(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        """Run the agent."""
        objective = inputs["objective"]
        first_task = inputs.get("first_task", "Make a todo list")
        self.add_task({"task_id": 1, "task_name": first_task})
        num_iters = 0
        while True:
            if self.task_list:
                self.print_task_list()

                # Step 1: Pull the first task
                task = self.task_list.popleft()
                self.print_next_task(task)

                # Step 2: Execute the task
                result = execute_task(
                    self.vectorstore, self.execution_chain, objective, task["task_name"]
                )
                this_task_id = int(task["task_id"])
                self.print_task_result(result)

                # Step 3: Store the result in Pinecone
                result_id = f"result_{task['task_id']}_{num_iters}"
                self.vectorstore.add_texts(
                    texts=[result],
                    metadatas=[{"task": task["task_name"]}],
                    ids=[result_id],
                )

                # Step 4: Create new tasks and reprioritize task list
                new_tasks = get_next_task(
                    self.task_creation_chain,
                    result,
                    task["task_name"],
                    [t["task_name"] for t in self.task_list],
                    objective,
                )
                for new_task in new_tasks:
                    self.task_id_counter += 1
                    new_task.update({"task_id": self.task_id_counter})
                    self.add_task(new_task)
                self.task_list = deque(
                    prioritize_tasks(
                        self.task_prioritization_chain,
                        this_task_id,
                        list(self.task_list),
                        objective,
                    )
                )
            num_iters += 1
            if self.max_iterations is not None and num_iters == self.max_iterations:
                print(
                    "\033[91m\033[1m" + "\n*****TASK ENDING*****\n" + "\033[0m\033[0m"
                )
                break
        return {}

    @classmethod
    def from_llm(
        cls, llm: BaseLLM, vectorstore: VectorStore, verbose: bool = False, **kwargs
    ) -> "BabyAGI":
        """Initialize the BabyAGI Controller."""
        task_creation_chain = TaskCreationChain.from_llm(llm, verbose=verbose)
        task_prioritization_chain = TaskPrioritizationChain.from_llm(
            llm, verbose=verbose
        )
        execution_chain = ExecutionChain.from_llm(llm, verbose=verbose)
        return cls(
            task_creation_chain=task_creation_chain,
            task_prioritization_chain=task_prioritization_chain,
            execution_chain=execution_chain,
            vectorstore=vectorstore,
            **kwargs,
        )

이 코드는 굉장히 긴데 겁먹을 필요 없구요, 그저 단계별로 결과값이 어떻게 프린트될지 예쁘게 formatting해주는 코드입니다. 코드를 자세히 보시면 Step 1, Step 2, Step 3, Step 4가 있습니다.

  • Step 1: 첫번째 태스크 불러오기

  • Step 2: 태스크 실행하기 (저희가 만들었던 세번째 체인, 태스크 실행하는 체인이 여기서 쓰입니다)

  • Step 3: 결과값을 Vector Store에 저장하기 (처음에 만들었던 vector store 기억하시죠? 여기에 저장되는 겁니다)

  • Step 4: 새로운 태스크들을 생성하고 중요도 높은 순으로 순번 관리하기 (저희가 만든 첫번째 체인 (태스크 생성 체인)과 두번째 체인 (순번 관리 체인)이 여기서 쓰입니다)

코드를 그대로 복붙하면 아래 이미지처럼 생겼습니다.

테스트해보기!!

OBJECTIVE = "Write a weather report for Seoul today"

llm = OpenAI(temperature=0)

# Logging of LLMChains
verbose = False
# If None, will keep on going forever
max_iterations: Optional[int] = 3
baby_agi = BabyAGI.from_llm(
    llm=llm, vectorstore=vectorstore, verbose=verbose, max_iterations=max_iterations
)

baby_agi({"objective": OBJECTIVE})

마지막으로 위 코드를 복붙합니다. OBJECTIVE는 유저 인풋입니다. AI Agent에게 어떤 업무를 시키고 싶은지에 따라 텍스트로 작성하면 됩니다.

저는 “서울의 날씨 리포트를 작성해줘”라는 인풋을 주었습니다.

여기서 max_iterations는 몇번의 루프를 돌건지 제한을 두는 것입니다. 지금은 3번의 루프만 돌게끔 해놨고, 만약 지정을 하지 않는다면 무한대로 루프가 돌아갈 것입니다.

복붙했다면 터미널에 python agent.py 명령어를 실행합니다. 이제 AI Agent가 자동으로 일을 하기 시작할겁니다!!

첫번째 태스크인 "todo list 만들기"를 실행한 결과입니다. 제가 "서울의 날씨 리포트를 써줘"라고 해서 이 업무를 하기 위해 만든 태스크 리스트입니다. 리포트를 작성하기 위해 1. 현재 서울의 온도, 2. 오늘 날씨 변화, 3. 습도, 4. 바람 속도와 방향, 5. UV index, 6. 미세먼지 인덱스 등을 리스트업 했습니다.

첫번째 태스크를 성공적으로 실행했으면 리스트에서 제외하고 두번째 태스크인 "서울의 현재 온도"를 실행합니다. 결과값으로 "National Weather Service Website"에서 불러온 정보를 토대로 현재 서울의 날씨는 24도라고 하네요.

다음으로 세번째 태스크를 실행합니다. 마찬가지로 두번째 태스크를 제외시킨 리스트가 업데이트됩니다.

결과값으로는 오늘 서울의 기온은 8도에서 13도 사이에서 왔다갔다 할거라고 하네요. (실제 구글링을 하지 않은 가상의 결과값들이기 때문에 정보가 정확하지 않습니다)

Part 2 마무리

max_iterations: Optional[int] = 3 숫자를 바꾸면서 몇번의 루프를 돌릴건지 테스트해봐도 되고, 직접 prompt template들을 수정하면서 특정 업무에 집중된 AI agent를 만들 수 있습니다.

예를 들어서 나는 자동으로 시장 조사를 해주는 에이전트를 만들고 싶다면 처음에 저희가 만든 세개의 체인들의 prompt template을 수정할 수 있겠죠?

혹시나 만드시면서 어려움을 겪거나 질문이 있으면 포스트에 댓글로 남겨주시거나 커피챗 요청 보내주세요! 10X AI Club의 두번째 글, 정말 길었는데 읽어주셔서 감사합니다.


다음 Part 3에서는 저희가 만든 AI Agent가 실제로 구글을 검색하면서 정확한 자료 조사를 할 수 있도록 만들 예정입니다.

6
3개의 답글

(채용) 콘텐츠 마케터, AI 엔지니어, 백엔드 개발자

지피터스의 수 천개 AI 활용 사례 데이터를 AI로 재가공 할 인재를 찾습니다

👉 이 게시글도 읽어보세요