aparnavellala/LangGraphMultiAgent

1import sys 2print("Python version")3print (sys. version) 4 5 6from typing import Annotated, Sequence, TypedDict7import operator8import functools9 10from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder11from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage12from langchain_community.tools.tavily_search import TavilySearchResults13from langchain_experimental.tools import PythonREPLTool14from langchain.agents import create_openai_tools_agent15from langchain_huggingface import HuggingFacePipeline16from langgraph.graph import StateGraph, END17 18from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline19 20# SETUP: HuggingFace Model and Pipeline21#name = "meta-llama/Llama-3.2-1B"22#name="deepseek-ai/DeepSeek-R1-Distill-Qwen-32B"23#name="deepseek-ai/deepseek-llm-7b-chat"24name="openai-community/gpt2"25#name="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"26#name="microsoft/Phi-3.5-mini-instruct"27#name="Qwen/Qwen2.5-7B-Instruct-1M"28 29tokenizer = AutoTokenizer.from_pretrained(name,truncation=True)30tokenizer.pad_token = tokenizer.eos_token31model = AutoModelForCausalLM.from_pretrained(name)32 33pipe = pipeline(34    "text-generation",35    model=model,36    tokenizer=tokenizer,37    device_map="auto",38    max_new_tokens=500,  # text to generate for outputs39)40print ("pipeline is created")41 42# Wrap in LangChain's HuggingFacePipeline43llm = HuggingFacePipeline(pipeline=pipe)44 45# Members and Final Options46members = ["Researcher", "Coder"]47options = ["FINISH"] + members48 49# Supervisor prompt50system_prompt = (51    "You are a supervisor tasked with managing a conversation between the following workers: {members}."52    " Given the following user request, respond with the workers to act next. Each worker will perform a task"53    " and respond with their results and status. When all workers are finished, respond with FINISH."54)55 56# Prompt template required for the workflow57prompt = ChatPromptTemplate.from_messages(58    [59        ("system", system_prompt),60        MessagesPlaceholder(variable_name="messages"),61        ("system", "Given the conversation above, who should act next? Or Should we FINISH? Select one of: {options}"),62    ]63).partial(options=str(options), members=", ".join(members))64 65print ("Prompt Template created")66 67# Supervisor routing logic68def route_tool_response(llm_response):69    """70    Parse the LLM response to determine the next step based on routing logic.71    """72    if "FINISH" in llm_response:73        return "FINISH"74    for member in members:75        if member in llm_response:76            return member77    return "Unknown"78 79def supervisor_chain(state):80    """81    Supervisor logic to interact with HuggingFacePipeline and decide the next worker.82    """83    messages = state.get("messages", [])84    user_prompt = prompt.format(messages=messages)85 86    try:87        llm_response = pipe(user_prompt, max_new_tokens=500)[0]["generated_text"]88    except Exception as e:89        raise RuntimeError(f"LLM processing error: {e}")90 91    next_action = route_tool_response(llm_response)92    return {"next": next_action}93 94# AgentState definition95class AgentState(TypedDict):96    messages: Annotated[Sequence[BaseMessage], operator.add]97    next: str98 99# Create tools100tavily_tool = TavilySearchResults(max_results=5)101python_repl_tool = PythonREPLTool()102 103# Create agents with their respective prompts104research_agent = create_openai_tools_agent(105    llm=llm,106    tools=[tavily_tool],107    prompt=ChatPromptTemplate.from_messages(108        [109            SystemMessage(content="You are a web researcher."),110            MessagesPlaceholder(variable_name="messages"),111            MessagesPlaceholder(variable_name="agent_scratchpad"),  # Add required placeholder112        ]113    ),114)115 116print ("Created agents with their respective prompts")117 118code_agent = create_openai_tools_agent(119    llm=llm,120    tools=[python_repl_tool],121    prompt=ChatPromptTemplate.from_messages(122        [123            SystemMessage(content="You may generate safe Python code for analysis."),124            MessagesPlaceholder(variable_name="messages"),125            MessagesPlaceholder(variable_name="agent_scratchpad"),  # Add required placeholder126        ]127    ),128)129 130 131print ("create_openai_tools_agent")132 133 134# Create the workflow135workflow = StateGraph(AgentState)136 137# Nodes138workflow.add_node("Researcher", research_agent)  # Pass the agent directly (no .run required)139workflow.add_node("Coder", code_agent)          # Pass the agent directly140workflow.add_node("supervisor", supervisor_chain)141 142# Add edges for workflow transitions143for member in members:144    workflow.add_edge(member, "supervisor")145 146workflow.add_conditional_edges(147    "supervisor",148    lambda x: x["next"],149    {k: k for k in members} | {"FINISH": END}  # Dynamically map workers to their actions150)151 152# Define entry point153workflow.set_entry_point("supervisor")154 155print(workflow)156 157# Compile the workflow158graph = workflow.compile()159 160from IPython.display import display, Image161display(Image(graph.get_graph().draw_mermaid_png()))162 163# Properly formatted initial state164initial_state = {165    "messages": [166        #HumanMessage(content="Code hello world and print it to the terminal.")  # Correct format for user input167        HumanMessage(content="Write Code for printing \"hello world\" in Python. Keep it precise.")  # Correct format for user input168    ]169}170 171# Execute the workflow172result = graph.invoke(initial_state)173print("Workflow Result:", result)