agents_langgraph_multiple_parallel_tools.py

## Tools/functions are copied from https://github.com/thoraxe/autogen-ols/blob/main/main.py
## Experimental code for multiple tools in parallel with single controller.

import json

with open("../config/cred_keys.json", "r") as f:
    cred = json.load(f)

import operator
import subprocess
from typing import TypedDict, Annotated

from langchain_openai import ChatOpenAI
from langchain_openai import AzureChatOpenAI

from langchain.tools import tool
from langchain_core.messages import BaseMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.agents import AgentAction

from langgraph.graph import StateGraph, END

# from IPython.display import Image, display
# from langchain_core.runnables.graph import MermaidDrawMethod

# def visualize_graph(app_graph):
#     display(
#         Image(
#             app_graph.get_graph().draw_mermaid_png(
#                 draw_method=MermaidDrawMethod.API,
#             )
#         )
#     )


llm_openai = ChatOpenAI(
    # model="gpt-3.5-turbo-0125",
    model="gpt-4o-mini",
    base_url=cred["openai"]["url"],
    openai_api_key=cred["openai"]["token"],
)
llm_azure = AzureChatOpenAI(
    model="gpt-4o-mini",
    azure_endpoint=cred["azure_openai"]["url"],
    api_key=cred["azure_openai"]["token"],
    api_version=cred["azure_openai"]["api_version"],
    deployment_name=cred["azure_openai"]["deployment_name"]["gpt-4o-mini"],
)
llm = llm_azure


@tool("get_pod_status")
def get_pod_status(namespace: str, pod: str) -> str:
    """
    Fetch the status information for a specific pod in the cluster by namespace.
    Only returns pod status and no other details.

    Args:
        namespace: the namespace where the pod is
        pod: the name of the pod to get the status for
    """
    output = subprocess.run(
        [
            "oc",
            "get",
            "pod",
            "-n",
            namespace,
            pod,
            "-o",
            "jsonpath='\{.status\}'",
        ],
        capture_output=True,
        timeout=2,
    )
    return output.stdout


@tool("get_object_details")
def get_object_details(namespace: str, kind: str, name: str) -> str:
    """
    Fetch the details for a specific object in the cluster.

    Args:
        namespace: the namespace where the object is
        kind: the kind of the object
        name: the name of the object
    """
    output = subprocess.run(
        ["oc", "get", kind, "-n", namespace, name, "-o", "yaml"],
        capture_output=True,
        timeout=2,
    )
    return output.stdout


@tool("get_nonrunning_pods")
def get_nonrunning_pods() -> str:
    """Fetch a list of pods that are not currently running"""
    output = subprocess.run(
        [
            "oc",
            "get",
            "pods",
            "-A",
            "--field-selector",
            "status.phase!=Running",
            "-o",
            "custom-columns=NAMESPACE:.metadata.namespace,NAME:.metadata.name",
        ],
        capture_output=True,
        timeout=2,
    )
    return output.stdout


@tool("get_object_namespace_list")
def get_object_namespace_list(kind: str, namespace: str) -> str:
    """
    Fetch a list of all instance of a specific type of kubernetes/openshift
    object in a specific namespace.

    Args:
        kind: the kubernetes/openshift objects to get
        namespace: the namespace containing the objects
    """
    output = subprocess.run(
        ["oc", "get", kind, "-n", namespace, "-o", "name"],
        capture_output=True,
        timeout=2,
    )
    return output.stdout


@tool("get_object_cluster_wide_list")
def get_object_cluster_wide_list(kind: str) -> str:
    """
    Fetch a list of all instances of a specific type of kubernetes/openshift
    object in the cluster.

    Args:
        kind: the kubernetes/openshift objects to get
    """
    print(f"provided kind: {kind}")
    output = subprocess.run(
        ["oc", "get", kind, "-A", "-o", "name"],
        capture_output=True,
        timeout=2,
    )
    return output.stdout


@tool("get_namespaces")
def get_namespaces() -> str:
    """Fetch the list of all namespaces in the cluster"""
    output = subprocess.run(["oc", "get", "namespaces"], capture_output=True, timeout=2)
    return output.stdout


@tool("final_answer")
def final_answer(
    introduction: str,
    research_steps: str,
    main_body: str,
    conclusion: str,
    sources: str,
):
    """Give final response as plain text"""
    return ""  # placeholder


tools = [
    get_pod_status,
    get_object_details,
    get_nonrunning_pods,
    get_object_namespace_list,
    get_object_cluster_wide_list,
    get_namespaces,
    final_answer,
]

system_prompt = """You are the controller, the great AI decision maker.
Given the user's query you must decide what to do with it based on the
list of tools provided to you.

If you see that a tool has been used (in the scratchpad) with a particular
query, do NOT use that same tool with the same query again. Also, do NOT use
any tool more than twice (ie, if the tool appears in the scratchpad twice, do
not use it again).

You should aim to collect information from a diverse range of sources before
providing the answer to the user. Once you have collected plenty of information
to answer the user's question (stored in the scratchpad) use the final_answer
tool.
"""

prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system_prompt),
        MessagesPlaceholder(variable_name="chat_history"),
        ("user", "{input}"),
        #    ("user", "Context:\n{context}"),
        ("assistant", "scratchpad: {scratchpad}"),
    ]
)


class AgentState(TypedDict):
    input: str
    chat_history: list[BaseMessage]
    context: str
    intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add]


def create_scratchpad(intermediate_steps: list[AgentAction]):
    research_steps = []
    for i, action in enumerate(intermediate_steps):
        if action.log != "TBD":
            # this was the ToolExecution
            research_steps.append(
                f"Tool: {action.tool}, input: {action.tool_input}\n"
                f"Output: {action.log}"
            )
    return "\n---\n".join(research_steps)


controller = (
    {
        "input": lambda x: x["input"],
        "chat_history": lambda x: x["chat_history"],
        #    "context": lambda x: x["context"],
        "scratchpad": lambda x: create_scratchpad(
            intermediate_steps=x["intermediate_steps"]
        ),
    }
    | prompt
    | llm.bind_tools(tools)
    #    | llm.bind_tools(tools, tool_choice="any")
    # | llm.bind_tools(tools, tool_choice="auto")
)


def run_controller(state: TypedDict):
    print("run controller")
    print(f"intermediate_steps: {state['intermediate_steps']}")
    out = controller.invoke(state)
    tool_name = out.tool_calls[0]["name"]
    tool_args = out.tool_calls[0]["args"]
    action_out = AgentAction(tool=tool_name, tool_input=tool_args, log="TBD")
    return {"intermediate_steps": [action_out]}


def router(state: TypedDict):
    if isinstance(state["intermediate_steps"], list):
        return state["intermediate_steps"][-1].tool
    else:
        print("Router invalid format")
        return "final_answer"


tool_str_to_func = {
    "get_pod_status": get_pod_status,
    "get_object_details": get_object_details,
    "get_nonrunning_pods": get_nonrunning_pods,
    "get_object_namespace_list": get_object_namespace_list,
    "get_object_cluster_wide_list": get_object_cluster_wide_list,
    "get_namespaces": get_namespaces,
    "final_answer": final_answer,
}


def run_tool(state: TypedDict):
    # use this as helper function so we repeat less code
    tool_name = state["intermediate_steps"][-1].tool
    tool_args = state["intermediate_steps"][-1].tool_input
    print(f"{tool_name}.invoke(input={tool_args})")
    # run tool
    out = tool_str_to_func[tool_name].invoke(input=tool_args)
    action_out = AgentAction(tool=tool_name, tool_input=tool_args, log=str(out))
    intermediate = {"intermediate_steps": [action_out]}
    if tool_name == "get_rag_content":
        intermediate["context"] = state["intermediate_steps"][-1].log
    else:
        intermediate["context"] = ""

    return intermediate


# initialize the graph with our AgentState
graph = StateGraph(AgentState)

# add nodes
graph.add_node("controller", run_controller)
graph.add_node("get_pod_status", run_tool)
graph.add_node("get_object_details", run_tool)
graph.add_node("get_nonrunning_pods", run_tool)
graph.add_node("get_object_namespace_list", run_tool)
graph.add_node("get_object_cluster_wide_list", run_tool)
graph.add_node("get_namespaces", run_tool)
graph.add_node("final_answer", run_tool)

# specify the entry node
graph.set_entry_point("controller")

# add the conditional edges which use the router
graph.add_conditional_edges(
    source="controller",  # where in graph to start
    path=router,  # function to determine which node is called
)

# create edges from each tool back to the controller
for tool_obj in tools:
    if tool_obj.name != "final_answer":
        graph.add_edge(tool_obj.name, "controller")

# if anything goes to final answer, it must then move to END
graph.add_edge("final_answer", END)

# finally, we compile our graph
runnable = graph.compile()
# visualize_graph(runnable)

runnable.invoke(
    {
        # "input": "get me failed pods",
        # "input": "get me status of the node",
        "input": "why pod is failing ?",
        "chat_history": [],
    }
)

# def _args_parser(args):
#     """Arguments parser."""
#     parser = argparse.ArgumentParser()
#     parser.add_argument("--query", type=str)
#     return parser.parse_args(args)

# def main():
#     args = _args_parser(sys.argv[1:])

#     runnable.invoke({
#         "input": args.query,
#         "chat_history": [],
#     })

# if __name__ == "__main__":
#     main()