Gemma 4のFunction Callingでローカルエージェントを構築する

import json
from transformers import AutoTokenizer, AutoModelForCausalLM
 
# モデルをロード（26B MoEの場合）
model_name = "google/gemma-4-26b-moe"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    torch_dtype="auto"
)
 
# 利用可能なツール定義
tools = [
    {
        "name": "search_web",
        "description": "ウェブから情報を検索します",
        "parameters": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "検索クエリ"
                }
            },
            "required": ["query"]
        }
    },
    {
        "name": "fetch_document",
        "description": "URLからドキュメントを取得します",
        "parameters": {
            "type": "object",
            "properties": {
                "url": {
                    "type": "string",
                    "description": "ドキュメントのURL"
                }
            },
            "required": ["url"]
        }
    }
]
 
def build_system_prompt(tools):
    """ツール情報を含むシステムプロンプトを構築"""
    prompt = "あなたは有能なアシスタントです。以下のツールを使用してタスクを完了します。\n\n"
    prompt += "利用可能なツール:\n"
    for tool in tools:
        prompt += f"- {tool['name']}: {tool['description']}\n"
    prompt += "\n関数を呼び出す際は、JSONフォーマットで以下のように記述してください:\n"
    prompt += '{"tool": "tool_name", "params": {...}}'
    return prompt
 
def run_agent_loop(user_query, max_iterations=5):
    """エージェントループを実行"""
    system_prompt = build_system_prompt(tools)
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_query}
    ]
    
    for i in range(max_iterations):
        # モデルの応答を生成
        input_ids = tokenizer.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=True,
            return_tensors="pt"
        ).to(model.device)
        
        outputs = model.generate(
            input_ids,
            max_new_tokens=512,
            temperature=0.7,
            top_p=0.9
        )
        
        response = tokenizer.decode(
            outputs[0][input_ids.shape[-1]:],
            skip_special_tokens=True
        )
        
        messages.append({"role": "assistant", "content": response})
        
        # レスポンスからツール呼び出しを抽出（簡略版）
        if '{"tool":' in response:
            try:
                # JSONを抽出（実際はより堅牢なパーサが必要）
                tool_call = json.loads(
                    response[response.find('{"tool":'):response.rfind('}')+1]
                )
                
                # ツールを実行（ここでは文字列を返す）
                tool_result = f"ツール '{tool_call['tool']}' の結果"
                messages.append({"role": "user", "content": tool_result})
                
            except json.JSONDecodeError:
                break
        else:
            # ツール呼び出しなし、終了
            return response
    
    return response
 
# 使用例
result = run_agent_loop("2024年の主要なAI動向を調べて、簡潔にまとめてください")
print(result)