【LangGraph入門】LangGraph Studioの簡単な使い方

from langgraph.graph import StateGraph, END
from my_agent.utils.nodes import selection_node, search_perplexity, database_search_employee, result_summary
from my_agent.utils.state import AgentState

# 新しいグラフを定義
workflow = StateGraph(AgentState)

# ノードを定義
workflow.add_node("selection", selection_node)
workflow.add_node("search_perplexity", search_perplexity)
workflow.add_node("database_search_employee", database_search_employee)
workflow.add_node("result_summary", result_summary)

# エントリーポイントを selection に設定
workflow.set_entry_point("selection")

# エッジの定義
workflow.add_conditional_edges(
    "selection",
    lambda state: state.current_role, # リスト
    {
        "1": "search_perplexity",
        "2": "database_search_employee"
    },
)
workflow.add_edge("search_perplexity", "result_summary")
workflow.add_edge("database_search_employee", "result_summary")
workflow.add_edge("result_summary", END)

# コンパイルの実施
graph = workflow.compile()

from pydantic import BaseModel, Field
from typing import Annotated
import operator

class AgentState(BaseModel):
    query: str = Field(..., description="ユーザーからの質問")
    current_role: str = Field(
        default="", description="選定された回答ロール"
    )
    answers: Annotated[list[str], operator.add] = Field(
        default=[], description="回答履歴"
    )

from langchain_openai import ChatOpenAI
from langgraph.prebuilt import ToolNode
from langchain_core.runnables import ConfigurableField
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from typing import Any
from openai import OpenAI

llm = ChatOpenAI(model="gpt-4o", temperature=0.0)
llm = llm.configurable_fields(max_tokens=ConfigurableField(id='max_tokens'))


def selection_node(state) -> dict[str, Any]:
    query = state.query
    prompt = ChatPromptTemplate.from_template(
    """質問を分析し、最も適切なエージェントを選択してください。

    ### 選択肢
    1.web検索エージェント:AIがWeb検索を行い、最適な情報を提供します。
    2.社員情報検索エージェント：該当の技術スタックに精通する社員を検索し、その情報を提供する。

    ### 質問
     {query}

    ### 出力形式
    回答は選択肢の番号（1または2）のみを返してください。
    """.strip()
    )
    # 選択肢の番号のみを返すことを期待したいため、max_tokensの値を1に変更
    chain = prompt | llm.with_config(configurable=dict(max_tokens=1)) | StrOutputParser()
    role_number = chain.invoke({"query": query})

    return {"current_role": role_number}

def search_perplexity(state) -> dict[str, Any]:
    """2024年の時点の最新の情報をウェブ検索する関数"""
    query = state.query
    messages = [
        {
            "role": "system",
            "content": (
                "You are an artificial intelligence assistant and you need to "
                "engage in a helpful, detailed, polite conversation with a user."
            ),
        },
        {
            "role": "user",
            "content": (
                query
            ),
        },
    ]
    client = OpenAI(api_key="xxxx", base_url="https://api.perplexity.ai")

    response = client.chat.completions.create(
        model="llama-3.1-sonar-large-128k-online",
        messages=messages,
    )
    answer = response.choices[0].message.content

    return {"answers": [answer]}

def database_search_employee (state) -> dict[str, Any]:
    """社員情報検索を実行する関数"""
    query = state.query
    # 本来は検索機能を実装するべきだか、ここではダミーデータが返ってきたことを想定
    search_result = """Aさん
    - 生産管理システムの最適化
        - 生産ラインの稼働状況をリアルタイムで監視するダッシュボードの開発
        - AIを活用した生産計画の自動化と効率化
        - 現場スタッフ向けの使いやすいモバイルアプリの設計・導入

    Bさん
    - 製品検査プロセスのデジタル化
        - 画像認識技術を活用した自動検査システムの構築
        - データ解析による不良品の原因分析と予防策の提案
        - 検査結果の可視化ツールを用いた品質管理レポートの作成

    Cさん
    - IoTを活用した工場のスマート化推進
        - センサーデータを活用した設備稼働状況の予測分析
        - クラウドプラットフォームによる生産データの一元管理
        - 設備保全のためのAI予測保全モデルの導入"""
    
    prompt = ChatPromptTemplate.from_template(
    """あなたは社員情報検索エージェントです。以下の社員情報を元にユーザの質問に回答してください。

    ### 社員情報
    {search_result}

    ### 質問
    {query}
    """.strip()
    )
    chain = prompt | llm | StrOutputParser()
    answer = chain.invoke({"search_result": search_result,"query": query})
    return {"answers": [answer]}

def result_summary(state) -> dict[str, Any]:
    """前のノードで得られた結果を統合するノード"""
    summary = f"Question: {state.query}\nAnswer: {state.answers[-1]}"
    return {"answers": [summary]}