Agent-s knowledge.py

import

import json
import os
from typing import Dict, Tuple

import numpy as np
from sklearn.metrics.pairwise import cosine_similarity

from gui_agents.s2.core.module import BaseModule
from gui_agents.s2.memory.procedural_memory import PROCEDURAL_MEMORY
from gui_agents.s2.core.engine import OpenAIEmbeddingEngine
from gui_agents.s2.utils.common_utils import (
    call_llm_safe,
    load_embeddings,
    load_knowledge_base,
    save_embeddings,
)
from gui_agents.s2.utils.query_perplexica import query_to_perplexica

클래스 초기화 -> 지식 베이스 파일 저장 경로 설정 및 engine.py의 임베딩 엔진 초기화

class KnowledgeBase(BaseModule):
    def __init__(
        self,
        local_kb_path: str,
        platform: str,
        engine_params: Dict,
        save_knowledge: bool = True,
    ):
        super().__init__(engine_params, platform)

        self.local_kb_path = local_kb_path

        # 임베딩 엔진 초기화
        # TODO: 다른 임베딩 엔진도 지원하도록 개선
        self.embedding_engine = OpenAIEmbeddingEngine()

        # 다양한 메모리 타입을 위한 경로 초기화
        self.episodic_memory_path = os.path.join(
            self.local_kb_path, "episodic_memory.json"
        )
        self.narrative_memory_path = os.path.join(
            self.local_kb_path, "narrative_memory.json"
        )
        self.embeddings_path = os.path.join(self.local_kb_path, "embeddings.pkl")

        # 작업 궤적 추적 초기화
        self.task_trajectory = ""
        self.current_subtask_trajectory = ""
        self.current_search_query = ""

•  query_formulator -> 검색 쿼리를 생성하는 에이전트
• llm_search_agent -> 검색을 수행하는 에이전트
• knowledge_fusion_agent -> 여러 정보를 융합하는 에이전트
• narrative_summarization_agent -> 전체 작업을 요약하는 에이전트
• episode_summarization_agent -> 하위 작업을 요약하는 에이전트

        self.rag_module_system_prompt = PROCEDURAL_MEMORY.RAG_AGENT.replace(
            "CURRENT_OS", self.platform
        )

        # 세 에이전트 모두 현재 OS에서 UI 자동화를 위한 정보를 제공하는 일반적인 RAG 프롬프트 공유
        self.query_formulator = self._create_agent(self.rag_module_system_prompt)
        self.llm_search_agent = self._create_agent(self.rag_module_system_prompt)
        self.knowledge_fusion_agent = self._create_agent(self.rag_module_system_prompt)

        self.narrative_summarization_agent = self._create_agent(
            PROCEDURAL_MEMORY.TASK_SUMMARIZATION_PROMPT
        )
        self.episode_summarization_agent = self._create_agent(
            PROCEDURAL_MEMORY.SUBTASK_SUMMARIZATION_PROMPT
        )

        self.save_knowledge = save_knowledge

지식 검색 메서드

def retrieve_knowledge(
    self, instruction: str, search_query: str, search_engine: str = "llm"
) -> Tuple[str, str]:
    """검색 엔진을 사용하여 지식 검색
    인자:
        instruction (str): 작업 지시
        observation (Dict): 현재 관찰
        search_engine (str): 사용할 검색 엔진"""

    # 검색 엔진을 사용하여 생성된 쿼리를 기반으로 지식 검색
    search_results = self._search(instruction, search_query, search_engine)

    return search_query, search_results

쿼리 생성 메서드

• 캐시된 쿼리가 있는지 확인하고 있으면 재사용
• 없으면 AI 에이전트에 현재 작업과 화면 스크린샷을 제공하여 적절한 검색 쿼리 생성 요청
• 생성된 쿼리를 캐시에 저장하고 반환

def formulate_query(self, instruction: str, observation: Dict) -> str:
    """지시와 현재 상태를 기반으로 검색 쿼리 생성"""
    query_path = os.path.join(self.local_kb_path, "formulate_query.json")
    try:
        with open(query_path, "r") as f:
            formulate_query = json.load(f)
    except:
        formulate_query = {}

    if instruction in formulate_query:
        return formulate_query[instruction]

    self.query_formulator.reset()

    self.query_formulator.add_message(
        f"The task is: {instruction}\n"
        "To use google search to get some useful information, first carefully analyze "
        "the screenshot of the current desktop UI state, then given the task "
        "instruction, formulate a question that can be used to search on the Internet "
        "for information in helping with the task execution.\n"
        "The question should not be too general or too specific. Please ONLY provide "
        "the question.\nQuestion:",
        image_content=(
            observation["screenshot"] if "screenshot" in observation else None
        ),
        role="user",
    )

    search_query = self.query_formulator.get_response().strip().replace('"', "")
    print("search query: ", search_query)
    formulate_query[instruction] = search_query
    with open(query_path, "w") as f:
        json.dump(formulate_query, f, indent=2)

    return search_query

검색 수행 메서드

• 캐시된 검색 결과가 있는지 확인하고 있으면 재사용
• 검색 엔진에 따라 다른 방식으로 검색:
  • llm: AI 모델의 내부 지식을 활용
  • perplexica: 외부 검색 엔진을 활용
• 검색 결과를 캐시에 저장하고 반환

def _search(self, instruction: str, search_query: str, search_engine: str) -> str:
    """지정된 엔진을 사용하여 검색 실행"""

    # 기본적으로 perplexica rag 지식을 사용하여 쿼리가 존재하는지 확인
    file = os.path.join(self.local_kb_path, f"{search_engine}_rag_knowledge.json")

    try:
        with open(file, "r") as f:
            exist_search_results = json.load(f)
    except:
        exist_search_results = {}

    if instruction in exist_search_results:
        return exist_search_results[instruction]
    if search_engine.lower() == "llm":
        self.llm_search_agent.reset()
        # LLM의 내부 지식을 검색 엔진처럼 사용
        self.llm_search_agent.add_message(search_query, role="user")
        search_results = self.llm_search_agent.get_response()
    elif search_engine.lower() == "perplexica":
        # perplexica를 사용하여 쿼리 검색
        search_results = query_to_perplexica(search_query)
    else:
        raise ValueError(f"지원되지 않는 검색 엔진: {search_engine}")

경험 검색 메서드

• 서술적 기억(전체 작업 수준)에서 지식 베이스와 임베딩 로드
• 현재 작업 지시에 대한 임베딩 생성 또는 검색
• 모든 저장된 경험에 대한 임베딩 생성 또는 검색
• 코사인 유사도를 계산하여 가장 유사한 경험 찾기
• 유사한 경험의 키와 내용 반환

def retrieve_narrative_experience(self, instruction: str) -> Tuple[str, str]:
    """임베딩을 사용하여 서술적 경험 검색"""

    knowledge_base = load_knowledge_base(self.narrative_memory_path)
    if not knowledge_base:
        return "None", "None"

    embeddings = load_embeddings(self.embeddings_path)

    # 지시에 대한 임베딩 가져오거나 생성
    instruction_embedding = embeddings.get(instruction)

    if instruction_embedding is None:
        instruction_embedding = self.embedding_engine.get_embeddings(instruction)
        embeddings[instruction] = instruction_embedding

    # 지식 베이스 항목에 대한 임베딩 가져오거나 생성
    candidate_embeddings = []
    for key in knowledge_base:
        candidate_embedding = embeddings.get(key)
        if candidate_embedding is None:
            candidate_embedding = self.embedding_engine.get_embeddings(key)
            embeddings[key] = candidate_embedding

        candidate_embeddings.append(candidate_embedding)

    save_embeddings(self.embeddings_path, embeddings)

    similarities = cosine_similarity(
        instruction_embedding, np.vstack(candidate_embeddings)
    )[0]
    sorted_indices = np.argsort(similarities)[::-1]

    keys = list(knowledge_base.keys())
    idx = 1 if keys[sorted_indices[0]] == instruction else 0
    return keys[sorted_indices[idx]], knowledge_base[keys[sorted_indices[idx]]]

지식 융합 메서드

• 작업 지시, 웹 검색 결과, 유사 작업 경험을 AI 에이전트에 제공
• AI가 두 정보를 결합하여 현재 작업에 가장 유용한 지식 생성
• 통합된 지식 반환

def knowledge_fusion(
    self,
    observation: Dict,
    instruction: str,
    web_knowledge: str,
    similar_task: str,
    experience: str,
) -> str:
    """웹 지식과 유사한 작업 경험 결합"""

    self.knowledge_fusion_agent.reset()

    self.knowledge_fusion_agent.add_message(
        f"Task: {instruction}\n"
        f"**Web search result**:\n{web_knowledge}\n\n"
        f"**Retrieved similar task experience**:\n"
        f"Similar task:{similar_task}\n{experience}\n\n"
        f"Based on the web search result and the retrieved similar task experience, "
        f"if you think the similar task experience is indeed useful to the main task, "
        f"integrate it with the web search result. Provide the final knowledge in a numbered list.",
        image_content=(
            observation["screenshot"] if "screenshot" in observation else None
        ),
        role="user",
    )
    return self.knowledge_fusion_agent.get_response()

기억 저장 메서드

• 지식 저장이 비활성화되어 있으면 아무것도 하지 않음
• 기존 일화적 기억 로드 또는 새로 생성
• 하위 작업 경험을 요약하고 저장

def save_episodic_memory(self, subtask_key: str, subtask_traj: str) -> None:
    """일화적 기억(하위 작업 수준 지식) 저장.

    인자:
        subtask_key (str): 하위 작업을 식별하는 키
        subtask_traj (str): 하위 작업의 궤적/경험
    """
    if not self.save_knowledge:
        return

    try:
        kb = load_knowledge_base(self.episodic_memory_path)
    except:
        kb = {}

    if subtask_key not in kb:
        subtask_summarization = self.summarize_episode(subtask_traj)
        kb[subtask_key] = subtask_summarization

작업 궤적 관리 메서드 -> 새 작업의 진행 과정을 초기화

def initialize_task_trajectory(self, instruction: str) -> None:
    """새 작업 궤적 초기화.

    인자:
        instruction (str): 작업 지시
    """
    self.task_trajectory = f"Task:\n{instruction}"
    self.current_search_query = ""
    self.current_subtask_trajectory = ""

에이전트가 작업을 수행하면서 생성한 메타데이터로 진행 과정을 업데이트

def update_task_trajectory(self, meta_data: Dict) -> None:
    """에이전트의 예측에서 나온 메타데이터로 작업 궤적 업데이트.

    인자:
        meta_data (Dict): 에이전트의 예측에서 나온 메타데이터
    """
    if not self.current_search_query and "search_query" in meta_data:
        self.current_search_query = meta_data["search_query"]

    self.task_trajectory += (
        "\n\nReflection:\n"
        + str(meta_data["reflection"])
        + "\n\n----------------------\n\nPlan:\n"
        + meta_data["executor_plan"]
    )

 

요약 기능 메서드

• 경험 궤적을 AI 에이전트에 제공
• AI가 경험을 요약하여 핵심 내용만 추출
• 요약된 내용 반환

def summarize_episode(self, trajectory):
    """평생 학습 반영을 위한 에피소드 경험 요약
    인자:
        trajectory: str: 요약할 에피소드 경험
    """

    # 다음 라운드 시도를 위해 전체 궤적에 대한 반영 생성, 이전 메시지를 예시로 유지
    self.episode_summarization_agent.add_message(trajectory)
    subtask_summarization = call_llm_safe(self.episode_summarization_agent)
    self.episode_summarization_agent.add_message(subtask_summarization)

    return subtask_summarization

LLM -> Llama 2 모델로 대체 가능 -> Ollama로 대체 가능