langgraph-test/devops.py - dev - zidanez/mcp-test

langgraph-test/devops.py
115 строк4 KB
zidanez
Develop prekol
17 мар 2025, 02:54
17 мар 2025, 02:54c726c1b
import io
import typing as t

from langgraph.constants import END
from typing_extensions import TypedDict

from langchain_core.messages import HumanMessage
from langgraph.prebuilt import create_react_agent
from langgraph.graph.message import add_messages, MessagesState
from langgraph.graph import StateGraph
from langchain_core.tools import tool
from langgraph.types import Command

from lllm import model

members = ["jenkinsAssistant", "dpmAssistant"]
options = members + ["FINISH"]


class GraphState(MessagesState):
    next: str


class Router(TypedDict):
    """Worker to route to next. If no workers needed, route to FINISH."""

    next: t.Literal[*options]


system_prompt = (
    f"Ты специалист по статистике. У тебя в подчинении есть: {members} Твоя задача:"
    f"используя подчиненных. Выполнить запрос пользователя. Для начала составь план, а затем выполни его, вызывая подчиненных."
    # " ответьте, какой работник будет действовать следующим, после Каждый работник будет выполнять задание"
    # " и сообщать о своих результатах"
    # " После "
    " Если контекст запроса пользователя не относится к твоей работе или ты не можешь выполнить его, используя подчиненных, ответь, что не можешь выполнить данный запроc"
)


@tool
def stub_tool(query: str) -> str:
    """Получение данных о сборках"""
    return f"Я выполнил твой запрос: {query}"


jenkins_tools = [stub_tool]
dpm_tools = [stub_tool]

jenkins_agent = create_react_agent(model=model,
                                   tools=jenkins_tools,
                                   state_modifier="Ты специалист по статистике Jenkins. Твоя задача: используя ТОЛЬКО свои функции выполнить требование пользователя. Используй только функции. Сначала построй план действий, а затем выполни его.")

dpm_agent = create_react_agent(model=model,
                               tools=dpm_tools,
                               state_modifier="Ты занимаешься статистикой по DPM. Твоя задача выполнить запрос пользователя. Используй только функции. Сначала построй план действий, а затем выполни его.")


def jenkins_node(state: GraphState):
    result = jenkins_agent.invoke(state)

    return Command(
        update={"messages": [
            HumanMessage(content=result["messages"][-1].content, name="jenkinsAssistant")
        ]},
        goto="planner"
    )


def dpm_node(state: GraphState):
    result = dpm_agent.invoke(state)

    return Command(
        update={"messages": [
            HumanMessage(content=result["messages"][-1].content, name="dpmAssistant")
        ]},
        goto="planner"
    )


def planner_node(state: GraphState) -> Command[t.Literal[*members, "__end__"]]:
    messages = [
                   {"role": "system", "content": system_prompt},
               ] + state["messages"]
    response = model.with_structured_output(Router).invoke(messages)
    print(response)
    goto = response["next"]
    if goto == "FINISH":
        goto = END

    return Command(goto=goto, update={"next": goto})


workflow = StateGraph(GraphState)

workflow.add_node("planner", planner_node)
workflow.add_node("jenkinsAssistant", jenkins_node)
workflow.add_node("dpmAssistant", dpm_node)

workflow.set_entry_point("planner")

graph = workflow.compile()

if __name__ == '__main__':
    # from PIL import Image
    #
    # img = Image.open(io.BytesIO(graph.get_graph().draw_mermaid_png()))
    #
    # img.show()

    for s in graph.stream(
            {"messages": [("user", "Какое количество сборок Jenkins было всего выполнено?")]}, subgraphs=True
            # {"messages": [("user", "Рецепт сырников")]}, subgraphs=True
    ):
        print(s)
        print("----")