ローカルLLMでPlanAndExecuteAgent

概要

今回はPlanAndExecuteという構造のAgentを実装してみました。
PlanAndExecuteは、ユーザーのリクエストに対してPlanモジュールによって行動計画を作成し、Agentがその計画を実行する。
さらに、その出力をReplanモジュールで評価、再計画します。
実装してみての結論を先に書いてしまいますが、あまり期待通りの動作をしてくれませんでした。。。
参考にしたnotebookは以下。

実装

LLMの準備

vLLMでモデルをデプロイします。

model_id = 'elyza/Llama-3-ELYZA-JP-8B'
vllm_model = LLM(
    model=model_id, dtype="float16",
    quantization="bitsandbytes", load_format="bitsandbytes",
    gpu_memory_utilization=1.0, max_model_len=2048,
)

ツール定義

LLMが使用可能なツールは簡単な計算を行うだけの関数です。

def calculator(expression: str) -> str:
    try:
        return f'{expression}の計算結果は、{str(eval(expression))}です。'
    except Exception as e:
        return f"計算エラー: {str(e)}"

tools = [
    Tool(
        name="calculator",
        func=calculator,
        description="数式の計算を実行します。入力は数式の文字列です。"
    )
]

def format_tools(tools) -> str:
    """ツール情報のフォーマット"""
    return "\n".join([
        f"- {tool.name}: {tool.description}"
        for tool in tools
    ])

Planモジュールの実装

実行計画を持つデータモデルを返却するようなモジュールを定義します。
Planモジュールは、ユーザーの入力に対して最初に１度だけ呼び出されることになります。

class Plan(BaseModel):
    steps: List[str]

class PlanModule:
    def __init__(self, model, prompt_template, data_model, tools: List[Tool] = None):
        self.vllm_model = model
        self.tools = tools or []
        self.tool_map = {tool.name: tool for tool in self.tools}        
        self.schema = data_model.model_json_schema()
        self.data_model = data_model
        json_logits_processor = JSONLogitsProcessor(self.schema, self.vllm_model)
        self.sampling_params = SamplingParams(
            temperature=0.3, top_p=0.5, max_tokens=1024,
            repetition_penalty=1.2,
            logits_processors=[json_logits_processor]
        )
        self.prompt_template = prompt_template

    @observe(as_type="generation", name='Plan')
    def run(self, question):
        current_prompt = self.prompt_template.format(
            tools=format_tools(self.tools),
            objective=question,
            schema=self.schema
        )

        outputs = self.vllm_model.generate(
            [current_prompt],
            self.sampling_params,
        )
        response_text = outputs[0].outputs[0].text
        step_data = json.loads(response_text)
        res = self.data_model(**step_data)
        return res

planner_content = """
<|start_header_id|>system<|end_header_id|>
与えられたユーザーの入力に対して、簡単な step-by-step の計画するAIです。
この計画には、正しく実行すれば正しい答えが得られるような、個々の作業を含むようにします。
余計なステップを追加しないでください。
最終ステップの結果が最終的な答えであるべきです。
各ステップに必要な情報がすべて含まれていることを確認してください。

以下はあなたが使用できるツールです。
ツールが必要な場合は使用を検討してください。
利用可能なツール:
{tools}
<|eot_id|>
<|start_header_id|>user<|end_header_id|>
与えられた以下の入力に対して、
正確な情報をユーザーに提供するために必要な行動の
簡単な step-by-step の作業計画を立ててください。

入力：
{objective}
<|eot_id|>
<|start_header_id|>assistant<|end_header_id|>
"""
planner = PlanModule(vllm_model, planner_content, Plan, tools=tools)

Replanモジュールの実装