Final_Assignment_Template

Build error

Final_Assignment_Template / myagent_langraph.py

cbf4814 over 1 year ago

12.3 kB

import base64
from typing import List, TypedDict, Annotated, Optional
from langchain_openai import ChatOpenAI
from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage
from langgraph.graph.message import add_messages
from langgraph.graph import START, StateGraph
from langgraph.prebuilt import ToolNode, tools_condition
#from IPython.display import Image, display
from langchain_community.document_loaders import WikipediaLoader
from langchain_community.tools.tavily_search import TavilySearchResults
#from langchain_community.tools.youtube_search import YoutubeSearchResults
from langchain_core.tools import tool

from dotenv import load_dotenv
import os

# Load environment variables
load_dotenv()
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

# Validate API key
if not OPENAI_API_KEY:
    print("Warning: OPENAI_API_KEY environment variable not found or empty!")

class AgentState(TypedDict):
    # The document provided
    question_text: Optional[str]
    messages: Annotated[list[AnyMessage], add_messages]

### TOOLS ###
@tool
def search_wikipedia(query: str) -> str:
    """Searches Wikipedia for a given query and returns maximum 1 results."""
    try:
        # Ensure query is properly encoded
        if isinstance(query, bytes):
            query = query.decode('utf-8')
        
        # Use WikipediaLoader with the encoded query
        search_docs = WikipediaLoader(
            query=query, 
            load_max_docs=1, 
            doc_content_chars_max=10000
        ).load()
        
        # Format the results with proper encoding
        formatted_docs = []
        for doc in search_docs:
            source = doc.metadata.get("source", "").encode('utf-8', errors='replace').decode('utf-8')
            page = doc.metadata.get("page", "").encode('utf-8', errors='replace').decode('utf-8')
            content = doc.page_content.encode('utf-8', errors='replace').decode('utf-8')
            formatted_docs.append(
                f'<Document source="{source}" page="{page}"/>\n{content}\n</Document>'
            )
        
        return "\n\n---\n\n".join(formatted_docs)
    except Exception as e:
        print(f"Error in search_wikipedia: {e}")
        return f"Error searching Wikipedia: {str(e)}"

@tool
def multiply(a: int, b: int) -> int:
    """Multiply two numbers.
    Args:
        a: first int
        b: second int
    """
    return a * b

@tool
def add(a: int, b: int) -> int:
    """Add two numbers.
    
    Args:
        a: first int
        b: second int
    """
    return a + b

@tool
def subtract(a: int, b: int) -> int:
    """Subtract two numbers.
    
    Args:
        a: first int
        b: second int
    """
    return a - b

@tool
def divide(a: int, b: int) -> int:
    """Divide two numbers.
    
    Args:
        a: first int
        b: second int
    """
    if b == 0:
        raise ValueError("Cannot divide by zero.")
    return a / b

@tool
def modulus(a: int, b: int) -> int:
    """Get the modulus of two numbers.
    
    Args:
        a: first int
        b: second int
    """
    return a % b

@tool
def web_search(query: str) -> str:
    """Search the web for a query and return maximum 3 results.
    
    Args:
        query: The search query."""
    search_tool = TavilySearchResults(max_results=3)
    search_results = search_tool.invoke(query)
    formatted_search_docs = "\n\n---\n\n".join(
        [
            f'<Document source="{result["url"]}" title="{result.get("title", "")}"/>\n{result["content"]}\n</Document>'
            for result in search_results
        ])
    return formatted_search_docs

# Equip the butler with tools
tools = [
    multiply,
    subtract,
    add,
    divide,
    modulus,
    search_wikipedia,
    web_search
]

# Initialize the OpenAI model with explicit encoding parameters
try:
    # Use encoding='utf-8' to ensure proper handling of special characters
    llm = ChatOpenAI(
        model="gpt-4o-mini",
        openai_api_key=OPENAI_API_KEY,
        default_headers={"Content-Type": "application/json; charset=utf-8"}
    )
    print("Binding tools...")
    llm_with_tools = llm.bind_tools(tools, parallel_tool_calls=False)
except Exception as e:
    print(f"Error initializing OpenAI model: {e}")
    raise

#### THE NODES #####
def assistant(state: AgentState):
    try:
        # Get the question text from state and ensure it's properly encoded
        question = state.get("question_text", "")
        
        # Build the system prompt piece by piece
        parts = [
            "You are a helpful assistant tasked with answering questions using a set of tools.",
            "Report your thoughts step by step.",
            "Once you have determined the answer, provide it in the following XML format:",
            "<FINAL_ANSWER>Your concise final answer here</FINAL_ANSWER>",
            "Your final answer should be as concise as possible - a number OR as few words as possible OR a comma separated list of numbers and/or strings.",
            "If asked for a number, don't use commas or units such as $ or percent sign unless specified otherwise.",
            "If asked for a string, don't use articles or abbreviations (e.g. for cities), and write digits in plain text unless specified otherwise.",
            "The answer inside <FINAL_ANSWER> tags must be complete, clear, and not truncated.",
            
            # Tool descriptions
            "You can use these provided tools:",
            "search_wikipedia(query: str) -> str:",
            "    Searches Wikipedia for a given query and returns maximum 1 results.",
            "web_search(query: str) -> str:",
            "    Searches the web for a query and returns maximum 3 results.",
            "multiply, add, subtract, divide, modulus: Basic math operations between two numbers."
        ]
        
        # Add the question
        parts.append(f"The question you have to answer is: {question}")
        
        # Join all parts with newlines and ensure proper encoding
        sys_prompt = "\n".join(parts).encode('utf-8', errors='replace').decode('utf-8')
        
        # Create system message with explicit encoding
        sys_msg = SystemMessage(content=sys_prompt)
        
        # Invoke the model with proper error handling
        messages = llm_with_tools.invoke([sys_msg] + state["messages"])
        return {
            "messages": messages,
            "question_text": question
        }
    except Exception as e:
        error_msg = f"Error in assistant node: {str(e)}"
        print(error_msg)
        return {
            "messages": [SystemMessage(content=f"An error occurred: {error_msg}")],
            "question_text": state.get("question_text", "")
        }

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

    # Define nodes: these do the work
    builder.add_node("assistant", assistant)
    builder.add_node("tools", ToolNode(tools))

    # Define edges: these determine how the control flow moves
    builder.add_edge(START, "assistant")
    builder.add_conditional_edges(
        "assistant",
        # If the latest message requires a tool, route to tools
        # Otherwise, provide a direct response
        tools_condition,
    )
    builder.add_edge("tools", "assistant")
    react_graph = builder.compile()
    return react_graph



class WalidzAgent:
    """A langgraph agent."""
    def __init__(self):
        print("BasicAgent initialized.")
        self.graph = build_graph()

    def __call__(self, question: str) -> str:
        if isinstance(question, dict) and "messages" in question and "question_text" in question:
            # We got a properly formatted input
            input_state = question
        else:
            # Create a proper input state from the question string
            human_msg = HumanMessage(content=question)
            input_state = {
                "messages": [human_msg],
                "question_text": question
            }

        # Stream the execution and capture final step
        final_state = None
        for step in self.graph.stream(input_state):
            print("--- STEP ---")
            print(step)
            final_state = step
        
        # Extract the answer from the final state
        if final_state:
            # Extract message content from the appropriate structure
            message_content = None
            
            # Check if the result has the nested 'assistant' structure
            if 'assistant' in final_state and 'messages' in final_state['assistant']:
                # Get message from assistant's messages field
                message = final_state['assistant']['messages']
                message_content = message.content
            # Fallback to direct messages list if available
            elif 'messages' in final_state:
                messages_list = final_state['messages']
                if messages_list and len(messages_list) > 0:
                    message_content = messages_list[-1].content
            
            if not message_content:
                print("Warning: Could not extract message content from response")
                return "Failed to extract message content from agent response"
            
            # First try XML tag extraction
            import re
            xml_pattern = re.compile(r'<FINAL_ANSWER>(.*?)</FINAL_ANSWER>', re.DOTALL)
            match = xml_pattern.search(message_content)
            
            if match:
                # Found XML tags
                answer = match.group(1).strip()
            else:
                # Fall back to "FINAL ANSWER:" format
                if "FINAL ANSWER:" in message_content:
                    # Split by "FINAL ANSWER:" and take the last part
                    final_answer_parts = message_content.split("FINAL ANSWER:")
                    answer = final_answer_parts[-1].strip()
                else:
                    # If no explicit final answer format is found, return the whole message
                    print("Warning: No final answer format found in response")
                    answer = message_content
            
            print(f"Extracted final answer: {answer}")
            return answer
            
        # Fallback if no valid answer found
        print("Warning: Could not extract answer from agent response")
        return "No valid answer could be extracted from agent response"
    
if __name__ == "__main__":

    question_text = "How many studio albums were published by Mercedes Sosa between 2000 and 2009 (included)? You can use the latest 2022 version of english wikipedia."
    question_text = "What country had the least number of athletes at the 1928 Summer Olympics? If there's a tie for a number of athletes, return the first in alphabetical order. Give the IOC country code as your answer."

    agent = WalidzAgent()
    # Prepare input state
    human_msg = HumanMessage(content=question_text)
    input_state = {
        "messages": [human_msg],
        "question_text": question_text
    }

    agent(input_state)


    # Save the graph visualization to disk
    #graph_image = react_graph.get_graph(xray=True).draw_mermaid_png()
    #with open("graph_visualization.png", "wb") as f:
    #    f.write(graph_image)
    #print("Graph visualization saved to 'graph_visualization.png'")

    # Also display it if in an environment that supports display
    #display(Image(graph_image))


    #question_sample = question_sample.encode('utf-8', errors='replace').decode('utf-8')


    """print("Running with sample question:", question_sample)
    try:
        # Create message with proper encoding
        human_msg = HumanMessage(content=question_sample)
        
        # Prepare input state
        input_state = {
            "messages": [human_msg],
            "question_text": question_sample
        }
        
        # Invoke graph with proper encoding handling
        print("Invoking graph...")
        messages = react_graph.invoke(input_state)
        
        # Show the messages
        for m in messages['messages']:
            m.pretty_print()
    except Exception as e:
        print(f"Error in Agent execution: {str(e)}")
        import traceback
        traceback.print_exc()"""

	import base64
	from typing import List, TypedDict, Annotated, Optional
	from langchain_openai import ChatOpenAI
	from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage
	from langgraph.graph.message import add_messages
	from langgraph.graph import START, StateGraph
	from langgraph.prebuilt import ToolNode, tools_condition
	#from IPython.display import Image, display
	from langchain_community.document_loaders import WikipediaLoader
	from langchain_community.tools.tavily_search import TavilySearchResults
	#from langchain_community.tools.youtube_search import YoutubeSearchResults
	from langchain_core.tools import tool

	from dotenv import load_dotenv
	import os

	# Load environment variables
	load_dotenv()
	OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

	# Validate API key
	if not OPENAI_API_KEY:
	print("Warning: OPENAI_API_KEY environment variable not found or empty!")

	class AgentState(TypedDict):
	# The document provided
	question_text: Optional[str]
	messages: Annotated[list[AnyMessage], add_messages]

	### TOOLS ###
	@tool
	def search_wikipedia(query: str) -> str:
	"""Searches Wikipedia for a given query and returns maximum 1 results."""
	try:
	# Ensure query is properly encoded
	if isinstance(query, bytes):
	query = query.decode('utf-8')

	# Use WikipediaLoader with the encoded query
	search_docs = WikipediaLoader(
	query=query,
	load_max_docs=1,
	doc_content_chars_max=10000
	).load()

	# Format the results with proper encoding
	formatted_docs = []
	for doc in search_docs:
	source = doc.metadata.get("source", "").encode('utf-8', errors='replace').decode('utf-8')
	page = doc.metadata.get("page", "").encode('utf-8', errors='replace').decode('utf-8')
	content = doc.page_content.encode('utf-8', errors='replace').decode('utf-8')
	formatted_docs.append(
	f'<Document source="{source}" page="{page}"/>\n{content}\n</Document>'
	)

	return "\n\n---\n\n".join(formatted_docs)
	except Exception as e:
	print(f"Error in search_wikipedia: {e}")
	return f"Error searching Wikipedia: {str(e)}"

	@tool
	def multiply(a: int, b: int) -> int:
	"""Multiply two numbers.
	Args:
	a: first int
	b: second int
	"""
	return a * b

	@tool
	def add(a: int, b: int) -> int:
	"""Add two numbers.

	Args:
	a: first int
	b: second int
	"""
	return a + b

	@tool
	def subtract(a: int, b: int) -> int:
	"""Subtract two numbers.

	Args:
	a: first int
	b: second int
	"""
	return a - b

	@tool
	def divide(a: int, b: int) -> int:
	"""Divide two numbers.

	Args:
	a: first int
	b: second int
	"""
	if b == 0:
	raise ValueError("Cannot divide by zero.")
	return a / b

	@tool
	def modulus(a: int, b: int) -> int:
	"""Get the modulus of two numbers.

	Args:
	a: first int
	b: second int
	"""
	return a % b

	@tool
	def web_search(query: str) -> str:
	"""Search the web for a query and return maximum 3 results.

	Args:
	query: The search query."""
	search_tool = TavilySearchResults(max_results=3)
	search_results = search_tool.invoke(query)
	formatted_search_docs = "\n\n---\n\n".join(
	[
	f'<Document source="{result["url"]}" title="{result.get("title", "")}"/>\n{result["content"]}\n</Document>'
	for result in search_results
	])
	return formatted_search_docs

	# Equip the butler with tools
	tools = [
	multiply,
	subtract,
	add,
	divide,
	modulus,
	search_wikipedia,
	web_search
	]

	# Initialize the OpenAI model with explicit encoding parameters
	try:
	# Use encoding='utf-8' to ensure proper handling of special characters
	llm = ChatOpenAI(
	model="gpt-4o-mini",
	openai_api_key=OPENAI_API_KEY,
	default_headers={"Content-Type": "application/json; charset=utf-8"}
	)
	print("Binding tools...")
	llm_with_tools = llm.bind_tools(tools, parallel_tool_calls=False)
	except Exception as e:
	print(f"Error initializing OpenAI model: {e}")
	raise

	#### THE NODES #####
	def assistant(state: AgentState):
	try:
	# Get the question text from state and ensure it's properly encoded
	question = state.get("question_text", "")

	# Build the system prompt piece by piece
	parts = [
	"You are a helpful assistant tasked with answering questions using a set of tools.",
	"Report your thoughts step by step.",
	"Once you have determined the answer, provide it in the following XML format:",
	"<FINAL_ANSWER>Your concise final answer here</FINAL_ANSWER>",
	"Your final answer should be as concise as possible - a number OR as few words as possible OR a comma separated list of numbers and/or strings.",
	"If asked for a number, don't use commas or units such as $ or percent sign unless specified otherwise.",
	"If asked for a string, don't use articles or abbreviations (e.g. for cities), and write digits in plain text unless specified otherwise.",
	"The answer inside <FINAL_ANSWER> tags must be complete, clear, and not truncated.",

	# Tool descriptions
	"You can use these provided tools:",
	"search_wikipedia(query: str) -> str:",
	" Searches Wikipedia for a given query and returns maximum 1 results.",
	"web_search(query: str) -> str:",
	" Searches the web for a query and returns maximum 3 results.",
	"multiply, add, subtract, divide, modulus: Basic math operations between two numbers."
	]

	# Add the question
	parts.append(f"The question you have to answer is: {question}")

	# Join all parts with newlines and ensure proper encoding
	sys_prompt = "\n".join(parts).encode('utf-8', errors='replace').decode('utf-8')

	# Create system message with explicit encoding
	sys_msg = SystemMessage(content=sys_prompt)

	# Invoke the model with proper error handling
	messages = llm_with_tools.invoke([sys_msg] + state["messages"])
	return {
	"messages": messages,
	"question_text": question
	}
	except Exception as e:
	error_msg = f"Error in assistant node: {str(e)}"
	print(error_msg)
	return {
	"messages": [SystemMessage(content=f"An error occurred: {error_msg}")],
	"question_text": state.get("question_text", "")
	}

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

	# Define nodes: these do the work
	builder.add_node("assistant", assistant)
	builder.add_node("tools", ToolNode(tools))

	# Define edges: these determine how the control flow moves
	builder.add_edge(START, "assistant")
	builder.add_conditional_edges(
	"assistant",
	# If the latest message requires a tool, route to tools
	# Otherwise, provide a direct response
	tools_condition,
	)
	builder.add_edge("tools", "assistant")
	react_graph = builder.compile()
	return react_graph



	class WalidzAgent:
	"""A langgraph agent."""
	def __init__(self):
	print("BasicAgent initialized.")
	self.graph = build_graph()

	def __call__(self, question: str) -> str:
	if isinstance(question, dict) and "messages" in question and "question_text" in question:
	# We got a properly formatted input
	input_state = question
	else:
	# Create a proper input state from the question string
	human_msg = HumanMessage(content=question)
	input_state = {
	"messages": [human_msg],
	"question_text": question
	}

	# Stream the execution and capture final step
	final_state = None
	for step in self.graph.stream(input_state):
	print("--- STEP ---")
	print(step)
	final_state = step

	# Extract the answer from the final state
	if final_state:
	# Extract message content from the appropriate structure
	message_content = None

	# Check if the result has the nested 'assistant' structure
	if 'assistant' in final_state and 'messages' in final_state['assistant']:
	# Get message from assistant's messages field
	message = final_state['assistant']['messages']
	message_content = message.content
	# Fallback to direct messages list if available
	elif 'messages' in final_state:
	messages_list = final_state['messages']
	if messages_list and len(messages_list) > 0:
	message_content = messages_list[-1].content

	if not message_content:
	print("Warning: Could not extract message content from response")
	return "Failed to extract message content from agent response"

	# First try XML tag extraction
	import re
	xml_pattern = re.compile(r'<FINAL_ANSWER>(.*?)</FINAL_ANSWER>', re.DOTALL)
	match = xml_pattern.search(message_content)

	if match:
	# Found XML tags
	answer = match.group(1).strip()
	else:
	# Fall back to "FINAL ANSWER:" format
	if "FINAL ANSWER:" in message_content:
	# Split by "FINAL ANSWER:" and take the last part
	final_answer_parts = message_content.split("FINAL ANSWER:")
	answer = final_answer_parts[-1].strip()
	else:
	# If no explicit final answer format is found, return the whole message
	print("Warning: No final answer format found in response")
	answer = message_content

	print(f"Extracted final answer: {answer}")
	return answer

	# Fallback if no valid answer found
	print("Warning: Could not extract answer from agent response")
	return "No valid answer could be extracted from agent response"

	if __name__ == "__main__":

	question_text = "How many studio albums were published by Mercedes Sosa between 2000 and 2009 (included)? You can use the latest 2022 version of english wikipedia."
	question_text = "What country had the least number of athletes at the 1928 Summer Olympics? If there's a tie for a number of athletes, return the first in alphabetical order. Give the IOC country code as your answer."

	agent = WalidzAgent()
	# Prepare input state
	human_msg = HumanMessage(content=question_text)
	input_state = {
	"messages": [human_msg],
	"question_text": question_text
	}

	agent(input_state)


	# Save the graph visualization to disk
	#graph_image = react_graph.get_graph(xray=True).draw_mermaid_png()
	#with open("graph_visualization.png", "wb") as f:
	# f.write(graph_image)
	#print("Graph visualization saved to 'graph_visualization.png'")

	# Also display it if in an environment that supports display
	#display(Image(graph_image))


	#question_sample = question_sample.encode('utf-8', errors='replace').decode('utf-8')


	"""print("Running with sample question:", question_sample)

	try:
	# Create message with proper encoding
	human_msg = HumanMessage(content=question_sample)

	# Prepare input state
	input_state = {
	"messages": [human_msg],
	"question_text": question_sample
	}

	# Invoke graph with proper encoding handling
	print("Invoking graph...")
	messages = react_graph.invoke(input_state)

	# Show the messages
	for m in messages['messages']:
	m.pretty_print()
	except Exception as e:
	print(f"Error in Agent execution: {str(e)}")
	import traceback
	traceback.print_exc()"""