agent.py

import os
from dotenv import load_dotenv
import traceback
import time
import pprint

from typing import Annotated,Sequence, TypedDict

from langchain_core.messages import BaseMessage, HumanMessage
from langgraph.graph.message import add_messages # helper function to add messages to the state
from langchain_core.messages import ToolMessage
from langchain_core.runnables import RunnableConfig
from langgraph.graph import StateGraph, END
from langchain_google_genai import ChatGoogleGenerativeAI

# Local imports
from tools import get_search_tool, get_tavily_search_tool, get_wikipedia_tool, wikipedia_search, wikipedia_search_3,\
                  execute_python_code_from_file, download_taskid_file, analyze_excel_file, get_analyze_mp3_tool,\
                  get_analyze_image_tool, arxiv_search, get_youtube_transcript

# Nota: per i test in locale si usa il .env
#       su HuggingFace invece si usano le variabili definite in Settings/"Variables and secrets"
load_dotenv()
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
GEMINI_MODEL = os.environ.get("GEMINI_MODEL")
GEMINI_BASE_URL = os.environ.get("GEMINI_BASE_URL")
GEMINI_TEMPERATURE = float(os.environ.get("GEMINI_TEMPERATURE"))

TOOLS_CALL_DELAY = 1.5

# V1
# GENERAL_AGENT_INSTRUCTIONS = """You are a helpful assistant tasked with answering questions using a set of tools.
# Now, I will ask you a question. Analyze the question and provide your answer.
# Your 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.
# Provide only the answer, without notes, explanations or comments."""

# V2
# GENERAL_AGENT_INSTRUCTIONS = """
# You are a general AI assistant. Your purpose is to answer questions and complete tasks accurately and concisely.
# You have access to various tools to help you gather information and perform actions.
# Always prioritize using your tools to find factual information if a question requires it.
# If a question can be answered directly from your knowledge, do so.
# If you use a tool, provide only the direct result or answer based on the tool's output.
# Do not include any conversational filler, explanations of your thought process, or pleasantries unless specifically asked.
# """

# V3
# GENERAL_AGENT_INSTRUCTIONS = """
# You are a general AI assistant. Your purpose is to answer questions and complete tasks accurately and concisely.
# You have access to various tools to help you gather information and perform actions.
# Always prioritize using your tools to find factual information if a question requires it.
# Your 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.
# Provide only the answer, without notes, explanations or comments.
# To complete this task successfully, follow these steps carefully:
# 1. Comprehend the task and identify the intended goal.
# 2. Break the task into clear, logical steps.
# 3. Select and prepare the tools or resources you need.
# 4. Revise your plan if necessary based on feedback.
# 5. Maintain internal state and track progress.
# 6. Verify that the goal has been fully achieved.
# 7. Present the final result clearly and concisely."""

# V4
# GENERAL_AGENT_INSTRUCTIONS = """
# You are a general AI assistant. Your purpose is to answer questions and complete tasks accurately and concisely.
# You have access to various tools to help you gather information and perform actions.
# Always prioritize using your tools to find factual information if a question requires it.
# Fo instance, if the question mentions Wikpedia, use the wikpedia tool.
# Your 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.
# Provide only the answer, without notes, explanations or comments."""

# V5
GENERAL_AGENT_INSTRUCTIONS = """
You are a general AI assistant. Your purpose is to answer questions and complete tasks accurately and concisely.
You have access to various tools to help you gather information and perform actions.
Always prioritize using your tools to find factual information if a question requires it.
If the question mentions Wikpedia, use the wikpedia tool; if the question mentions a youtube url, use get_youtube_transcript tool.
Analyze the question and plan the necessary steps to get the answer.
Your 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.
Provide only the answer, without notes, explanations or comments."""

#
# Inizializza il modello e gli associa i tool
#

# ChatGoogleGenerativeAI è il package ufficiale di LangChain per interagire con i modelli Gemini
# https://python.langchain.com/docs/integrations/chat/google_generative_ai/
chat = ChatGoogleGenerativeAI(
    model=GEMINI_MODEL,
    google_api_key=GEMINI_API_KEY,
    temperature = GEMINI_TEMPERATURE)

# Imposta i tool
#search_tool = get_search_tool()
search_tool = get_tavily_search_tool()
#wikipedia_tool = get_wikipedia_tool()
analyze_mp3_tool = get_analyze_mp3_tool(chat)
analyze_png_tool = get_analyze_image_tool(chat)

tools = [search_tool,
         wikipedia_search_3,
         execute_python_code_from_file,
         download_taskid_file,
         analyze_excel_file,
         analyze_mp3_tool,
         analyze_png_tool,
         arxiv_search,
         get_youtube_transcript]

# Bind tools to the model
chat_with_tools = chat.bind_tools(tools)

tools_by_name = {tool.name: tool for tool in tools}

# debug
print("Tools:")
for tool in tools:
    print("  {}".format(tool.name))

#
# Definisce il grafo
#

class AgentState(TypedDict):
    """The state of the agent."""
    messages: Annotated[Sequence[BaseMessage], add_messages]
    number_of_steps: int


# Define our tool node
def call_tool(state: AgentState):
    outputs = []
    # Iterate over the tool calls in the last message
    for i, tool_call in enumerate(state["messages"][-1].tool_calls):
        # Get the tool by name
        tool_result = tools_by_name[tool_call["name"]].invoke(tool_call["args"])

        # print(f"\n--- DEBUG: Tool Raw Output (Length: {len(tool_result)} chars) ---")
        # print(tool_result)
        # print("------------------------------------------------------------------\n")

        outputs.append(
            ToolMessage(
                content=tool_result,
                name=tool_call["name"],
                tool_call_id=tool_call["id"],
            )
        )

        # Add a delay after each tool call, but not after the very last one
        if i < len(state["messages"][-1].tool_calls) - 1:
            time.sleep(TOOLS_CALL_DELAY)

    return {"messages": outputs}


def call_model( state: AgentState, config: RunnableConfig):
    # Modo 1)
    # Invoke the model with the system prompt and the messages
    #response = chat_with_tools.invoke(state["messages"], config)

    # Modo 2) - aggiunge in fondo alcune istruzioni
    # Create a copy to avoid modifying the original state and append instruction to the end
    # messages = state["messages"][:]
    # messages.append(
    #     HumanMessage(content="Provide only the answer, without explanations or comments.")
    # )  # Append instruction to the end
    # response = chat_with_tools.invoke(messages, config)

    # Modo 3)
    # Create a new list for messages to send to the LLM
    # Start with the general instructions
    messages_to_send = [HumanMessage(content=GENERAL_AGENT_INSTRUCTIONS)]

    # Append all existing messages from the agent state
    messages_to_send.extend(state["messages"])

    response = chat_with_tools.invoke(messages_to_send)

    # We return a list, because this will get added to the existing messages state using the add_messages reducer
    return {"messages": [response]}


# Define the conditional edge that determines whether to continue or not
def should_continue(state: AgentState):
    messages = state["messages"]
    # If the last message is not a tool call, then we finish
    if not messages[-1].tool_calls:
        return "end"
    # default to continue
    return "continue"


def get_agent():
    # Creazione del grafo
    workflow = StateGraph(AgentState)

    # 1. Add our nodes 
    workflow.add_node("llm", call_model)
    workflow.add_node("tools",  call_tool)
    # 2. Set the entrypoint as `agent`, this is the first node called
    workflow.set_entry_point("llm")
    # 3. Add a conditional edge after the `llm` node is called.
    workflow.add_conditional_edges(
        # Edge is used after the `llm` node is called.
        "llm",
        # The function that will determine which node is called next.
        should_continue,
        # Mapping for where to go next, keys are strings from the function return, and the values are other nodes.
        # END is a special node marking that the graph is finish.
        {
            # If `tools`, then we call the tool node.
            "continue": "tools",
            # Otherwise we finish.
            "end": END,
        },
    )
    # 4. Add a normal edge after `tools` is called, `llm` node is called next.
    workflow.add_edge("tools", "llm")

    # 5. Now we can compile our graph
    react_graph = workflow.compile()

    return react_graph


# Riferimenti
#
# https://ai.google.dev/gemini-api/docs/langgraph-example