agent.py

import os
from dotenv import load_dotenv
from langchain_core.messages import SystemMessage, HumanMessage
from langchain_openai import AzureChatOpenAI

from langgraph.graph import START, StateGraph, MessagesState
from langgraph.prebuilt import tools_condition, ToolNode
from langchain_core.runnables import RunnableConfig

# Wikepedia tool
from langchain_community.document_loaders import WikipediaLoader

# Tavily web search
from langchain_community.tools.tavily_search import TavilySearchResults

# Python loader
from langchain_community.document_loaders import PythonLoader

# Whisper
from langchain_community.document_loaders.parsers.audio import AzureOpenAIWhisperParser
from langchain_core.documents.base import Blob

# excel
from langchain_community.document_loaders import UnstructuredExcelLoader

load_dotenv()


# --- TOOLS ---
# Simple tool to test a tool call
def meaning_of_life(a: int, b: int) -> int:
    """Returns meaning of life

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


# https://www.restack.io/docs/langchain-knowledge-wikipedia-loader-cat-ai
# https://api.python.langchain.com/en/latest/community/document_loaders/langchain_community.document_loaders.wikipedia.WikipediaLoader.html#
# ¤ I ended up not using this tool since I could not get it to return the table data in the Markov question. The Taveli search tool also find wiki content
# Better approach could be to combine this tool (to get URL) +  a webreader to get content
def wikipedia_search(query: str) -> str:
    """Searches Wikipedia for a given query and fetches full document

    Args:
        query: the query to search for
    """
    loader = loader = WikipediaLoader(
        query=query,
        load_max_docs=1,
        doc_content_chars_max=4000,
        load_all_available_meta=False,
    )
    documents = loader.load()
    formatted_search_docs = "\n\n---\n\n"

    for next_doc in documents:
        formatted_doc = f'<Document source="{next_doc.metadata["source"]}" title="{next_doc.metadata.get("title", "")}"\n{next_doc.page_content}\n</Document>'
        formatted_search_docs = formatted_search_docs + formatted_doc

    result = f"{{wiki_results: {formatted_search_docs}}}"

    return result


def web_search(query: str) -> str:
    """Search Web with Tavily for a query and return results.

    Args:
        query: The search query."""
    tool = TavilySearchResults(
        max_results=3,
        include_answer=True,
        include_raw_content=True,
        include_images=True,
        # search_depth="advanced",
        # include_domains = []
        # exclude_domains = []
    )

    documents = tool.invoke(input=query)
    formatted_search_docs = "\n\n---\n\n"
    for next_doc in documents:
        url = next_doc["url"]
        title = next_doc["title"]
        content = next_doc["content"]
        formatted_doc = (
            f'<Document source="{url}" title="{title}"\n{content}\n</Document>'
        )
        formatted_search_docs = formatted_search_docs + formatted_doc

    result = f"{{web_results: {formatted_search_docs}}}"

    return result


# https://python.langchain.com/api_reference/community/document_loaders/langchain_community.document_loaders.python.PythonLoader.html
def python_file_reader(file_name: str) -> str:
    """Reads a python file and returns the content

    Args:
        file_name: the filename to read
    """
    file_path = os.path.join(os.path.dirname(__file__), "files", file_name)
    loader = PythonLoader(file_path=file_path)
    documents = loader.load()
    formatted_search_docs = "\n\n---\n\n"

    for next_doc in documents:
        formatted_doc = (
            f'<Document source="{file_name}"\n{next_doc.page_content}\n</Document>'
        )
        formatted_search_docs = formatted_search_docs + formatted_doc

    result = f"{{python_code: {formatted_search_docs}}}"

    return result


# https://python.langchain.com/docs/integrations/document_loaders/microsoft_excel/
# https://python.langchain.com/api_reference/community/document_loaders/langchain_community.document_loaders.excel.UnstructuredExcelLoader.html
def excel_file_reader(excel_file_name: str) -> str:
    """Reads an excel file and returns the content

    Args:
        excel_file_name: the filename to read
    """
    file_path = os.path.join(os.path.dirname(__file__), "files", excel_file_name)
    loader = UnstructuredExcelLoader(file_path, mode="elements")

    documents = loader.load()
    formatted_search_docs = "\n\n---\n\n"

    for next_doc in documents:
        formatted_doc = f'<Document source="{excel_file_name}"\n{next_doc.metadata["text_as_html"]}\n</Document>'
        formatted_search_docs = formatted_search_docs + formatted_doc

    result = f"{{python_code: {formatted_search_docs}}}"

    return result


# https://python.langchain.com/api_reference/community/document_loaders/langchain_community.document_loaders.parsers.audio.AzureOpenAIWhisperParser.html
def audio_to_text(audio_file_name: str) -> str:
    """Listen to audio and extract text from speech

    Args:
        audio_file_name: the audio filename to read
    """
    file_path = os.path.join(os.path.dirname(__file__), "files", audio_file_name)

    deployment_name = os.environ.get("AZURE_WHISPER_DEPLOYMENT")
    api_version = os.environ.get("AZURE_WHISPER_API_VERSION")
    api_key = os.environ.get("AZURE_WHISPER_API_KEY")
    azure_endpoint = os.environ.get("AZURE_WHISPER_ENDPOINT")

    whisper_parser = AzureOpenAIWhisperParser(
        deployment_name=deployment_name,
        api_version=api_version,
        api_key=api_key,
        azure_endpoint=azure_endpoint,
        # other params...
    )

    audio_blob = Blob(path=file_path)
    response = whisper_parser.parse(audio_blob)

    formatted_search_docs = "\n\n---\n\n"

    for next_doc in response:
        formatted_doc = f'<Document source="{audio_file_name}"\n{next_doc.page_content}\n</Document>'
        formatted_search_docs = formatted_search_docs + formatted_doc

    result = f"{{transscribed_audio: {formatted_search_docs}}}"

    return result


tools = [
    meaning_of_life,
    web_search,
    python_file_reader,
    audio_to_text,
    wikipedia_search,
    excel_file_reader,
]


# --- GRAPH ---
# This functions allow us to use the web interface to test the graph
def make_graph(config: RunnableConfig):
    graph = create_graph()
    return graph


# This function is used to create the graph
def create_graph():
    # Define LLM with bound tools
    azure_endpoint = os.environ.get("AZURE_ENDPOINT_LLM")
    api_key = os.environ.get("AZURE_API_KEY_LLM")
    api_version = os.environ.get("AZURE_API_VERSION_LLM")
    deployment = os.environ.get("AZURE_DEPLOYMENT_LLM")

    # Initialize LLM
    llm = AzureChatOpenAI(
        azure_deployment=deployment,
        api_version=api_version,
        temperature=0.01,
        max_tokens=None,
        timeout=None,
        max_retries=2,
        api_key=api_key,
        azure_endpoint=azure_endpoint,
    )
    llm_with_tools = llm.bind_tools(tools)

    # System message
    # original_system_prompt_txt = "You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER]. YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string."
    system_prompt_txt = "You are a general AI assistant that uses tools to answer questions. YOUR FINAL ANSWER should be a number represented as digits OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number or how many, only reply with a number represented as digits nothing else, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for an abbreviation or a code only reply with that. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string."

    sys_msg = SystemMessage(system_prompt_txt)

    # Node
    def assistant(state: MessagesState):
        return {"messages": [llm_with_tools.invoke([sys_msg] + state["messages"])]}

    # Build graph
    builder = StateGraph(MessagesState)
    builder.add_node("assistant", assistant)
    builder.add_node("tools", ToolNode(tools))
    builder.add_edge(START, "assistant")
    builder.add_conditional_edges(
        "assistant",
        # If the latest message (result) from assistant is a tool call -> tools_condition routes to tools
        # If the latest message (result) from assistant is a not a tool call -> tools_condition routes to END
        tools_condition,
    )
    builder.add_edge("tools", "assistant")

    # Compile graph
    graph = builder.compile()

    return graph


# You can run this script directly to test the graph
# Alternatively in a commandprompt run "langgraph dev" and it will allow you to interact with the graph in a web ui
if __name__ == "__main__":
    # Build the graph
    graph = create_graph()

    # Run the graph
    """ 
    print(f"******** TEST NORMAL LLM CALL ********")
    question = "What is an elephant? "
    messages = [HumanMessage(content=question)]
    messages = graph.invoke({"messages": messages})
    for m in messages["messages"]:
        m.pretty_print()

    print(f"******** TESTING MEANING OF LIFE TOOL ********")
    question = "What is meaning of life 10+10?"
    messages = [HumanMessage(content=question)]
    messages = graph.invoke({"messages": messages})
    for m in messages["messages"]:
        m.pretty_print()

 
    print("******** TESTING WIKEPEDIA TOOL ********")
    # expected answer is "Samuel"
    question = "Search Wikipedia and find out who is the recipient of the malko competition in 2024"
    messages = [HumanMessage(content=question)]
    messages = graph.invoke({"messages": messages})
    for m in messages["messages"]:
        m.pretty_print()



    print("******** TESTING WEB SEARCH TOOL ********")
    # expected answer is "Samuel"
    question = "Search web for information about mozart"
    messages = [HumanMessage(content=question)]
    messages = graph.invoke({"messages": messages})
    for m in messages["messages"]:
        m.pretty_print()


    print("******** PYTHON LOAD TOOL ********")
    question = "what does this python code do? filename is f918266a-b3e0-4914-865d-4faa564f1aef.py"
    messages = [HumanMessage(content=question)]
    messages = graph.invoke({"messages": messages})
    for m in messages["messages"]:
        m.pretty_print()


    print("******** TRANSSCRIBE AUDIO TOOL ********")
    question = "Hi, I was out sick from my classes on Friday, so I'm trying to figure out what I need to study for my Calculus mid-term next week. My friend from class sent me an audio recording of Professor Willowbrook giving out the recommended reading for the test, but my headphones are broken :( Could you please listen to the recording for me and tell me the page numbers I'm supposed to go over? I've attached a file called Homework.mp3 that has the recording. Please provide just the page numbers as a comma-delimited list. And please provide the list in ascending order. File to use is 1f975693-876d-457b-a649-393859e79bf3.mp3"
    messages = [HumanMessage(content=question)]
    messages = graph.invoke({"messages": messages})
    for m in messages["messages"]:
        m.pretty_print()


    print("******** EXCEL TOOL ********")
    question = "The attached Excel file contains the sales of menu items for a local fast-food chain. What were the total sales that the chain made from food (not including drinks)? Express your answer in USD with two decimal places. File to use is 7bd855d8-463d-4ed5-93ca-5fe35145f733.xlsx"
    messages = [HumanMessage(content=question)]
    messages = graph.invoke({"messages": messages})
    for m in messages["messages"]:
        m.pretty_print()
"""