Final_Assignment_Agents

Sleeping

KarthikKB1998

Update agent.py

41062c1 verified over 1 year ago

8.34 kB

from typing import TypedDict, Annotated
from langgraph.graph.message import add_messages
from langchain_core.messages import AnyMessage, HumanMessage, AIMessage, SystemMessage
from langgraph.prebuilt import ToolNode
from langgraph.graph import START, StateGraph, END
from langchain_community.tools import DuckDuckGoSearchRun
from langchain_openai import ChatOpenAI
import requests
import base64
import json
import io
import contextlib
import os
import subprocess
import tempfile
from dotenv import load_dotenv

load_dotenv()


def replace(_, new_value):
    return new_value

search_tool = DuckDuckGoSearchRun()

tools = [search_tool]


llm = ChatOpenAI(
    model="gpt-4o",
)
llm_with_tools = llm.bind_tools(tools)


# Generate the AgentState and Agent graph
class AgentState(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]
    question: Annotated[str, replace]
    task_id: Annotated[str, replace]
    file_name: Annotated[str, replace]

def get_question(state: AgentState):
    print("get_question")
    question_dict = json.loads(state['messages'][-1].content)
    print(question_dict)
    state['question'] = question_dict["question"]
    state['task_id'] = question_dict["task_id"]
    state['file_name'] = question_dict["file_name"]
    return state


def reformulate_question(state: AgentState):
    print("reformulate_question")
    sys_msg = SystemMessage(content=f"You are a detective and you have to answer a hard question. The first step is to understand the question. Can you reformulate it to make it clear and concise ? The question is : {state['question']}")
    
    state['messages'] = state['messages'] + [llm_with_tools.invoke([sys_msg])]
    return state

def assistant(state: AgentState):
    print("assistant")
    # System message
    sys_msg="""
You are a helpful assistant that can answer questions about the world. You can use the internet through the search tool to find information.
"""
    sys_msg = SystemMessage(content=sys_msg)
    state['messages'] = state['messages'] + [llm_with_tools.invoke([sys_msg] + state["messages"])]
    return state

def get_final_answer(state: AgentState):
    print("get_final_answer")
    sys_msg = SystemMessage(content=f"""
Reply the answer and only the answer of the question. Do not make a sentence, just the answer. If the answer is a number, return it in numeric form.
If the question specifies a format, please return it in the specified format. Do not add unnecessary uppercasing or punctuation. Here is the question again: {state['question']} and here is the answer: {state['messages'][-1].content}
""")
    final_answer = state["messages"][-1].content  # The last assistant message
    print(f"final answer: {final_answer}")
    state['messages'] = state['messages'] + [llm.invoke([sys_msg])]
    return state

    
def python_interpreter(state: AgentState):
    print("python interpreter")

    question_id = state['task_id']
    url = f"https://agents-course-unit4-scoring.hf.space/files/{question_id}"
    
    # Step 1: Download the code
    response = requests.get(url)
    python_code = response.content.decode("utf-8")

    print("Running user code with subprocess...")

    # Step 2: Write code to temp file
    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as tmp_file:
        tmp_file.write(python_code)
        tmp_file_path = tmp_file.name

    try:
        # Step 3: Run using subprocess (set timeout if needed)
        result = subprocess.run(
            ["python3", tmp_file_path],
            capture_output=True,
            text=True,
            timeout=30  # seconds
        )
        if result.returncode != 0:
            output = f"⚠️ Error:\n{result.stderr.strip()}"
        else:
            output = result.stdout.strip() or "Code ran but produced no output."
    except subprocess.TimeoutExpired:
        output = "Execution timed out."
    except Exception as e:
        output = f"Unexpected error: {e}"

    # Step 4: Update agent state
    state['messages'] = state['messages'] + [
        AIMessage(content=f"The output of the Python code is:\n```\n{output}\n```")
    ]
    
    return state



# def python_interpreter(state: AgentState):
#     question_id = state['task_id']
#     url = f"https://agents-course-unit4-scoring.hf.space/files/{question_id}"

#     # Step 1: Fetch Python code as bytes and decode
#     response = requests.get(url)
#     python_code = response.content.decode("utf-8")
#     # print(python_code)

#     # Step 2: Capture stdout
#     stdout = io.StringIO()
#     with contextlib.redirect_stdout(stdout):
#         try:
#             exec(python_code)  # Use empty dict for isolated global context
#         except Exception as e:
#             output = f"Error during execution: {e}"
#         else:
#             output = stdout.getvalue().strip()
#     print("Code output : ", output)
#     # Step 3: Save output as message
#     state['messages'] = state['messages'] + [
#         AIMessage(content=f"The output of the python code is : {output}")
#     ]
#     return state

def image_interpreter(state: AgentState):
    print("image interpreter")
    """Answers a question about an image. Provide raw image bytes and a question."""
    question_id = state['task_id']
    url = f"https://agents-course-unit4-scoring.hf.space/files/{question_id}"
    image = requests.get(url).content
    image_base64 = base64.b64encode(image).decode("utf-8")

    messages=[
        {"role": "system", "content": "You are a helpful assistant that can answer questions based on images."},
        {
            "role": "human",
            "content": [
                {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_base64}"}},
                {"type": "text", "text": state['question']}
            ],
        },
    ]
    resp = llm.invoke(messages)
    print("Image response : ", resp)
    state['messages'] = state['messages'] + [resp]

    return state

def not_handled(state: AgentState):
    print("not_handled")
    state['messages'] = state['messages'] + [AIMessage(content="I cannot handle this question.")]
    return state


def tools_condition(state: AgentState) -> str:
    print("tools_condition")
    last_msg = state["messages"][-1]

    if "tool_calls" in last_msg.additional_kwargs:
        # LLM is requesting a tool
        return "tools"

    # No tool requested → go to final summary
    return "get_final_answer"

def files_condition(state: AgentState) -> str:
    print("files_condition")
    if state["file_name"].endswith(".py"):
        return "python_interpreter"
    elif state["file_name"].endswith(".png"):
        return "image_interpreter"
    elif state["file_name"]=="":
        return "reformulate_question"
    else:
        return "not_handled"


def build_agent():
    ## The graph
    builder = StateGraph(AgentState)
    
    # Define nodes: these do the work
    builder.add_node("get_question", get_question)
    builder.add_node("reformulate_question", reformulate_question)
    builder.add_node("assistant", assistant)
    builder.add_node("tools", ToolNode(tools))
    builder.add_node("get_final_answer", get_final_answer)
    builder.add_node("python_interpreter", python_interpreter)
    builder.add_node("image_interpreter", image_interpreter)
    builder.add_node("not_handled", not_handled)
    
    # Define edges: these determine how the control flow moves
    builder.add_edge(START, "get_question")
    builder.add_conditional_edges(
        "get_question",
        files_condition,
        {
            "python_interpreter": "python_interpreter",
            "image_interpreter": "image_interpreter",
            "reformulate_question": "reformulate_question",
            "not_handled": "not_handled"
        }
    )
    builder.add_edge("python_interpreter", "get_final_answer")
    builder.add_edge("image_interpreter", "get_final_answer")
    builder.add_edge("reformulate_question", "assistant")
    builder.add_conditional_edges(
        "assistant",
        tools_condition,
        {
            "tools": "tools",
            "get_final_answer": "get_final_answer"
        }
    )
    builder.add_edge("tools", "assistant")
    builder.add_edge("get_final_answer", END)
    builder.add_edge("not_handled", END)
    mySuperAgent = builder.compile()
    return mySuperAgent

	from typing import TypedDict, Annotated
	from langgraph.graph.message import add_messages
	from langchain_core.messages import AnyMessage, HumanMessage, AIMessage, SystemMessage
	from langgraph.prebuilt import ToolNode
	from langgraph.graph import START, StateGraph, END
	from langchain_community.tools import DuckDuckGoSearchRun
	from langchain_openai import ChatOpenAI
	import requests
	import base64
	import json
	import io
	import contextlib
	import os
	import subprocess
	import tempfile
	from dotenv import load_dotenv

	load_dotenv()


	def replace(_, new_value):
	return new_value

	search_tool = DuckDuckGoSearchRun()

	tools = [search_tool]


	llm = ChatOpenAI(
	model="gpt-4o",
	)
	llm_with_tools = llm.bind_tools(tools)


	# Generate the AgentState and Agent graph
	class AgentState(TypedDict):
	messages: Annotated[list[AnyMessage], add_messages]
	question: Annotated[str, replace]
	task_id: Annotated[str, replace]
	file_name: Annotated[str, replace]

	def get_question(state: AgentState):
	print("get_question")
	question_dict = json.loads(state['messages'][-1].content)
	print(question_dict)
	state['question'] = question_dict["question"]
	state['task_id'] = question_dict["task_id"]
	state['file_name'] = question_dict["file_name"]
	return state


	def reformulate_question(state: AgentState):
	print("reformulate_question")
	sys_msg = SystemMessage(content=f"You are a detective and you have to answer a hard question. The first step is to understand the question. Can you reformulate it to make it clear and concise ? The question is : {state['question']}")

	state['messages'] = state['messages'] + [llm_with_tools.invoke([sys_msg])]
	return state

	def assistant(state: AgentState):
	print("assistant")
	# System message
	sys_msg="""
	You are a helpful assistant that can answer questions about the world. You can use the internet through the search tool to find information.
	"""
	sys_msg = SystemMessage(content=sys_msg)
	state['messages'] = state['messages'] + [llm_with_tools.invoke([sys_msg] + state["messages"])]
	return state

	def get_final_answer(state: AgentState):
	print("get_final_answer")
	sys_msg = SystemMessage(content=f"""
	Reply the answer and only the answer of the question. Do not make a sentence, just the answer. If the answer is a number, return it in numeric form.
	If the question specifies a format, please return it in the specified format. Do not add unnecessary uppercasing or punctuation. Here is the question again: {state['question']} and here is the answer: {state['messages'][-1].content}
	""")
	final_answer = state["messages"][-1].content # The last assistant message
	print(f"final answer: {final_answer}")
	state['messages'] = state['messages'] + [llm.invoke([sys_msg])]
	return state


	def python_interpreter(state: AgentState):
	print("python interpreter")

	question_id = state['task_id']
	url = f"https://agents-course-unit4-scoring.hf.space/files/{question_id}"

	# Step 1: Download the code
	response = requests.get(url)
	python_code = response.content.decode("utf-8")

	print("Running user code with subprocess...")

	# Step 2: Write code to temp file
	with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as tmp_file:
	tmp_file.write(python_code)
	tmp_file_path = tmp_file.name

	try:
	# Step 3: Run using subprocess (set timeout if needed)
	result = subprocess.run(
	["python3", tmp_file_path],
	capture_output=True,
	text=True,
	timeout=30 # seconds
	)
	if result.returncode != 0:
	output = f"⚠️ Error:\n{result.stderr.strip()}"
	else:
	output = result.stdout.strip() or "Code ran but produced no output."
	except subprocess.TimeoutExpired:
	output = "Execution timed out."
	except Exception as e:
	output = f"Unexpected error: {e}"

	# Step 4: Update agent state
	state['messages'] = state['messages'] + [
	AIMessage(content=f"The output of the Python code is:\n```\n{output}\n```")
	]

	return state



	# def python_interpreter(state: AgentState):
	# question_id = state['task_id']
	# url = f"https://agents-course-unit4-scoring.hf.space/files/{question_id}"

	# # Step 1: Fetch Python code as bytes and decode
	# response = requests.get(url)
	# python_code = response.content.decode("utf-8")
	# # print(python_code)

	# # Step 2: Capture stdout
	# stdout = io.StringIO()
	# with contextlib.redirect_stdout(stdout):
	# try:
	# exec(python_code) # Use empty dict for isolated global context
	# except Exception as e:
	# output = f"Error during execution: {e}"
	# else:
	# output = stdout.getvalue().strip()
	# print("Code output : ", output)
	# # Step 3: Save output as message
	# state['messages'] = state['messages'] + [
	# AIMessage(content=f"The output of the python code is : {output}")
	# ]
	# return state

	def image_interpreter(state: AgentState):
	print("image interpreter")
	"""Answers a question about an image. Provide raw image bytes and a question."""
	question_id = state['task_id']
	url = f"https://agents-course-unit4-scoring.hf.space/files/{question_id}"
	image = requests.get(url).content
	image_base64 = base64.b64encode(image).decode("utf-8")

	messages=[
	{"role": "system", "content": "You are a helpful assistant that can answer questions based on images."},
	{
	"role": "human",
	"content": [
	{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_base64}"}},
	{"type": "text", "text": state['question']}
	],
	},
	]
	resp = llm.invoke(messages)
	print("Image response : ", resp)
	state['messages'] = state['messages'] + [resp]

	return state

	def not_handled(state: AgentState):
	print("not_handled")
	state['messages'] = state['messages'] + [AIMessage(content="I cannot handle this question.")]
	return state


	def tools_condition(state: AgentState) -> str:
	print("tools_condition")
	last_msg = state["messages"][-1]

	if "tool_calls" in last_msg.additional_kwargs:
	# LLM is requesting a tool
	return "tools"

	# No tool requested → go to final summary
	return "get_final_answer"

	def files_condition(state: AgentState) -> str:
	print("files_condition")
	if state["file_name"].endswith(".py"):
	return "python_interpreter"
	elif state["file_name"].endswith(".png"):
	return "image_interpreter"
	elif state["file_name"]=="":
	return "reformulate_question"
	else:
	return "not_handled"


	def build_agent():
	## The graph
	builder = StateGraph(AgentState)

	# Define nodes: these do the work
	builder.add_node("get_question", get_question)
	builder.add_node("reformulate_question", reformulate_question)
	builder.add_node("assistant", assistant)
	builder.add_node("tools", ToolNode(tools))
	builder.add_node("get_final_answer", get_final_answer)
	builder.add_node("python_interpreter", python_interpreter)
	builder.add_node("image_interpreter", image_interpreter)
	builder.add_node("not_handled", not_handled)

	# Define edges: these determine how the control flow moves
	builder.add_edge(START, "get_question")
	builder.add_conditional_edges(
	"get_question",
	files_condition,
	{
	"python_interpreter": "python_interpreter",
	"image_interpreter": "image_interpreter",
	"reformulate_question": "reformulate_question",
	"not_handled": "not_handled"
	}
	)
	builder.add_edge("python_interpreter", "get_final_answer")
	builder.add_edge("image_interpreter", "get_final_answer")
	builder.add_edge("reformulate_question", "assistant")
	builder.add_conditional_edges(
	"assistant",
	tools_condition,
	{
	"tools": "tools",
	"get_final_answer": "get_final_answer"
	}
	)
	builder.add_edge("tools", "assistant")
	builder.add_edge("get_final_answer", END)
	builder.add_edge("not_handled", END)
	mySuperAgent = builder.compile()
	return mySuperAgent