概要

LLM(大規模言語モデル)を利用したアプリ開発のフレームワークであるLangchainを使ってネットワーク機器からログを取得してみます

www.langchain.com

(2024/09/13) ライブラリを最新のものに修正
(2024/09/17) Langchain v0.3.0に対応

環境

AlmaLinux 9.3
Python 3.9.18

Pythonライブラリ

langchain-cohere==0.3.0
langchain-community==0.3.0
langchain-core==0.3.0
langchain==0.3.0
pydantic==2.9.1
python-dotenv==1.0.1
netmiko==4.4.0

やれたこと

CohereのCommand R+モデルをつかってエージェントを作成し自作したツールを動かして質問に答えさせることができました。

Command R +はトライアルAPIキーを使っているため現時点(2024/4/29)では無料です cohere.com

コード

早速完成したコードを載せます。
StreamlitでWebApp化しているためGUI上で質問の入力と回答の確認をできるようにしました。

import streamlit as st

from langchain.agents import AgentExecutor
from langchain_cohere import ChatCohere
from langchain_cohere.react_multi_hop.agent import create_cohere_react_agent
from langchain_core.tools import BaseTool
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.runnables import RunnablePassthrough

from dotenv import load_dotenv
load_dotenv()

from netmiko import ConnectHandler
from netmiko.exceptions import NetmikoTimeoutException
from typing import Type, List, Optional
from pydantic import BaseModel

# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-

@st.dialog("機器の認証情報")
def device_auth_form():
    st.write("機器の認証情報を入れるよ")
    host = st.text_input("ホスト", value="dummy")
    username = st.text_input("ユーザ", value="admin")
    password = st.text_input("パスワード",value="admin")
    secret = st.text_input("特権パスワード")
    device_type=st.text_input("device_type",value="cisco_ios")
    if st.button("Submit"):
        st.session_state.network_device_info = NetworkDeviceInfo(
            host=host,
            username=username,
            password=password,
            secret=secret,
            device_type=device_type
        )
        st.rerun()

# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
# データ構造定義
class NetworkDeviceInfo(BaseModel):
    device_type: str
    host: str
    username: str
    password: str
    secret: Optional[str]

class Command(BaseModel):
    command: str

class CommandList(BaseModel):
    commands: List[Command]
    class Config:
        from_attributes = True

class NetworkDeviceInput(BaseModel):
    device: NetworkDeviceInfo
    commands: List[Command]

# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
# NetmikoでSSH接続してログを取得するツール
class NetworkDevice(BaseTool):
    name: str = "get_command_nw_device"
    description: str = "Access NW devices and execute commands with Netmiko"
    args_schema: Type[NetworkDeviceInput] = NetworkDeviceInput

    def _run(self, **param) -> List[str]:
        # 辞書からNetworkDeviceInputを作成
        input_data = NetworkDeviceInput(**param)
        output = []
        device_data = input_data.device.dict()  # 辞書に変換
        command_list = input_data.commands

        try:
            with ConnectHandler(**device_data) as net_connect:
                if device_data.get('secret'):
                    net_connect.enable()
                for command in command_list:
                    output.append(net_connect.send_command(command.command))
        except NetmikoTimeoutException:
            pass
        return output

# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
# プロンプトテンプレート
command_template = """
・質問を解決するために実施すべきshowコマンドを出力します
・回答はJSON形式で出力してください
    ・複数コマンドの変数サンプル:
        {{"commands": [{{"command": "show run"}}, {{"command": "show ip route"}}]}}
    ・単一コマンドの変数サンプル:
        {{"commands": [{{"command": "show ip route"}}]}}

質問: {question}
"""

template = """
前提:
・日本語で回答してください

利用可能なツール:
  - ツール名: get_command_nw_device
    説明: ネットワークデバイスにアクセスし、Netmikoを使用してコマンドを実行します。
    使用方法:
      引数:
        - `device`: デバイスへの接続情報を含むJSONオブジェクト。
             - フィールド:
                - `device_type`: デバイスの種類（例: `"cisco_ios"`）
                - `host`: デバイスのIPアドレス（例: `"192.168.1.1"`）
                - `username`: ユーザー名（例: `"admin"`）
                - `password`: パスワード（例: `"password123"`）
                - `secret`: （オプション）特権モードへのパスワード（例: `"enable_secret"`）
        - `commands`: 実行したいコマンドのリスト。各コマンドは辞書形式で、以下のフィールドを持ちます。
                - `command`: 実行するコマンドの文字列（例: `"show running-config"`）
   - 出力: コマンドの実行結果のリスト。

質問: {question}

あなたの役割:
- ユーザーの質問に答えるために、必要に応じてツールを使用してください。

"""

# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
st.set_page_config(layout="wide")
st.title("Cohere Agent Custom Tool")

# ツール定義
tools = [NetworkDevice()]  # ツールをリストにまとめる

# プロンプト定義
prompt = ChatPromptTemplate.from_template(template)
command_prompt = ChatPromptTemplate.from_template(command_template)

command_parser = JsonOutputParser(pydantic_object=CommandList)

# LLM定義
command_r = ChatCohere(model="command-r-plus-08-2024", temperature=0)
command_gen = command_prompt | command_r | command_parser

# エージェント定義
agent = create_cohere_react_agent(
    llm=command_r,
    tools=tools,  # ツールをエージェントに渡す
    prompt=prompt  # プロンプトもエージェントに渡す
)

# AgentExecutorを定義
agent_executor = AgentExecutor(
    agent=agent,  # 作成したエージェントを指定
    tools=tools,  # ツールも再度渡す
    verbose=True  # 実行時の詳細出力を有効にする
)

# フォーム
if "network_device_info" not in st.session_state:
    st.session_state.network_device_info = NetworkDeviceInfo(host="",
            username="",
            password="",
            secret=None,
            device_type=""
        )


device = st.session_state.network_device_info.dict()

# 認証フォーム
if st.button("認証情報"):
    device_auth_form()

# デバイス情報の生成確認
question = st.text_area("質問", value="ルートの統計情報教えて")

# チェーンにデータを渡して実行
chain = (
    {
        "device": RunnablePassthrough(),  # 生成されたデバイス情報を渡す
        "commands": command_gen,  # 生成されたコマンドを渡す
        "question": RunnablePassthrough()  # 質問を文字列として渡す
    }
    | agent_executor  # ここでエージェントを使ってコマンド実行
)

if st.button("実行"):
    with st.spinner("生成中...."):
        result = chain.invoke({"question": question, "device": device})
        st.write(result['output'])

Streamlit起動

python -m streamlit run cohere_tool.py

質問実行

応答

エージェントがツールを実行している様子

コードの詳細

1. 各種ライブラリのインポート

初めに、必要なライブラリをインポートします
- langchainのライブラリ群 - 認証情報を扱うためのdotenv - SSHしてログをとるためのNetmiko

import streamlit as st

from langchain.agents import AgentExecutor
from langchain_cohere import ChatCohere
from langchain_cohere.react_multi_hop.agent import create_cohere_react_agent
from langchain_core.tools import BaseTool
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.runnables import RunnablePassthrough

from dotenv import load_dotenv
load_dotenv()

from netmiko import ConnectHandler
from netmiko.exceptions import NetmikoTimeoutException
from typing import Type, List, Optional
from pydantic import BaseModel

2.自作ツールで利用するデータ構造の定義

自作ツールを動かすために、LLMにどんな変数構造で情報を引き渡す必要があるかを伝える必要があるようです、そのため引数はどんな型であるかpydanticを使い定義します。

# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
# データ構造定義
class NetworkDeviceInfo(BaseModel):
    device_type: str
    host: str
    username: str
    password: str
    secret: Optional[str]

class Command(BaseModel):
    command: str

class CommandList(BaseModel):
    commands: List[Command]
    class Config:
        from_attributes = True

class NetworkDeviceInput(BaseModel):
    device: NetworkDeviceInfo
    commands: List[Command]

3.自作ツール部分

Netmikoを使い、ログを取得する処理部分を定義します。
リストを作成しコマンドの応答を要素として追加して戻すシンプルなものにしました。

# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
# NetmikoでSSH接続してログを取得するツール
class NetworkDevice(BaseTool):
    name: str = "get_command_nw_device"
    description: str = "Access NW devices and execute commands with Netmiko"
    args_schema: Type[NetworkDeviceInput] = NetworkDeviceInput

    def _run(self, **param) -> List[str]:
        # 辞書からNetworkDeviceInputを作成
        input_data = NetworkDeviceInput(**param)
        output = []
        device_data = input_data.device.dict()  # 辞書に変換
        command_list = input_data.commands

        try:
            with ConnectHandler(**device_data) as net_connect:
                if device_data.get('secret'):
                    net_connect.enable()
                for command in command_list:
                    output.append(net_connect.send_command(command.command))
        except NetmikoTimeoutException:
            pass
        return output

4.プロンプトの定義

LLMに質問をする際のプロンプトのテンプレートを定義しています。
今回は、以下3つの質問をすることでツールを動かす内容を決めるようなイメージにしています

実行コマンドの生成
認証情報の生成(今回はテキストのインプットから解析する形にした)
ツール実行結果から質問の回答の生成

# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
# プロンプトテンプレート
command_template = """
・質問を解決するために実施すべきshowコマンドを出力します
・回答はJSON形式で出力してください
    ・複数コマンドの変数サンプル:
        {{"commands": [{{"command": "show run"}}, {{"command": "show ip route"}}]}}
    ・単一コマンドの変数サンプル:
        {{"commands": [{{"command": "show ip route"}}]}}

質問: {question}
"""

template = """
前提:
・日本語で回答してください

利用可能なツール:
  - ツール名: get_command_nw_device
    説明: ネットワークデバイスにアクセスし、Netmikoを使用してコマンドを実行します。
    使用方法:
      引数:
        - `device`: デバイスへの接続情報を含むJSONオブジェクト。
             - フィールド:
                - `device_type`: デバイスの種類（例: `"cisco_ios"`）
                - `host`: デバイスのIPアドレス（例: `"192.168.1.1"`）
                - `username`: ユーザー名（例: `"admin"`）
                - `password`: パスワード（例: `"password123"`）
                - `secret`: （オプション）特権モードへのパスワード（例: `"enable_secret"`）
        - `commands`: 実行したいコマンドのリスト。各コマンドは辞書形式で、以下のフィールドを持ちます。
                - `command`: 実行するコマンドの文字列（例: `"show running-config"`）
   - 出力: コマンドの実行結果のリスト。

質問: {question}

あなたの役割:
- ユーザーの質問に答えるために、必要に応じてツールを使用してください。

"""

5.メイン処理部分

ツールやプロンプトなどを定義していきます
認証情報とコマンドは質問に対する回答をJSON形式にするためJsonOutputParserを利用しています。

chainで処理定義しており、以下のような順序で動作します
1. 実行コマンドの生成
2. エージェント起動

エージェント内で利用するためのツールはtoolにリストで定義します。
エージェントはモデルごとに専用のものがあるようで、Cohereではcreate_cohere_react_agent を利用します(create_cohere_tools_agentといったものもあるようですが、うまく動かずでした。。)

認証情報は、Streamlitのフォームから入力するようにしています。

# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
st.set_page_config(layout="wide")
st.title("Cohere Agent Custom Tool")

# ツール定義
tools = [NetworkDevice()]  # ツールをリストにまとめる

# プロンプト定義
prompt = ChatPromptTemplate.from_template(template)
command_prompt = ChatPromptTemplate.from_template(command_template)

command_parser = JsonOutputParser(pydantic_object=CommandList)

# LLM定義
command_r = ChatCohere(model="command-r-plus-08-2024", temperature=0)
command_gen = command_prompt | command_r | command_parser

# エージェント定義
agent = create_cohere_react_agent(
    llm=command_r,
    tools=tools,  # ツールをエージェントに渡す
    prompt=prompt  # プロンプトもエージェントに渡す
)

# AgentExecutorを定義
agent_executor = AgentExecutor(
    agent=agent,  # 作成したエージェントを指定
    tools=tools,  # ツールも再度渡す
    verbose=True  # 実行時の詳細出力を有効にする
)

# フォーム
if "network_device_info" not in st.session_state:
    st.session_state.network_device_info = NetworkDeviceInfo(host="",
            username="",
            password="",
            secret=None,
            device_type=""
        )


device = st.session_state.network_device_info.dict()

# 認証フォーム
if st.button("認証情報"):
    device_auth_form()

# デバイス情報の生成確認
question = st.text_area("質問", value="ルートの統計情報教えて")

# チェーンにデータを渡して実行
chain = (
    {
        "device": RunnablePassthrough(),  # 生成されたデバイス情報を渡す
        "commands": command_gen,  # 生成されたコマンドを渡す
        "question": RunnablePassthrough()  # 質問を文字列として渡す
    }
    | agent_executor  # ここでエージェントを使ってコマンド実行
)

if st.button("実行"):
    with st.spinner("生成中...."):
        result = chain.invoke({"question": question, "device": device})
        st.write(result['output'])

認証情報

認証情報はdotenvを利用したためコードには登場しませんが、スクリプトと同じ階層に.envファイルを作成しAPI KEY情報などを定義します

.envファイル

COHERE_API_KEY = "<API KEY>"

感想

何とか自作ツールを動かすことができました、ツール実行のために定義した pydantic周りの理解が足りたいように感じました。。

いろいろなツールを作って、高度な運用補助ツールみたいなもの作りたいなーってモチベーションで頑張ります

NetBoxやServiceNowから機器情報を引っ張ってきてSSHして原因報告する。とかとか想像すると楽しいですね

ただ、エージェント難しい・・・なんもわからん。。。

理解が誤ってる箇所などあれば(優しく)ご指摘いただけたら嬉しいです。

うさラボ

お勉強と備忘録

【Langchain】エージェント経由でNetmikoを使ってネットワーク機器からログを取得する

概要

環境