# LangGraph è un framework che consente di creare applicazioni pronte per la produzione offrendo strumenti di controllo sul flusso del proprio agente.
# LangGraph è un framework sviluppato da LangChain per gestire il flusso di controllo delle applicazioni che integrano un LLM.
# https://www.langchain.com/
# LangChain fornisce un'interfaccia standard per interagire con modelli e altri componenti, utile per il recupero, le chiamate LLM e le chiamate agli strumenti. 
# Le classi di LangChain possono essere utilizzate in LangGraph, ma non è obbligatorio.

# Quando usare LangGraph? Quando si progettano applicazioni di intelligenza artificiale, ci si trova di fronte a un compromesso fondamentale tra controllo e libertà:
# Gli agenti, come quelli che si possono trovare in smolagents, sono molto liberi. Possono richiamare più strumenti in un'unica fase di azione, creare i propri strumenti, ecc.
# Tuttavia, questo comportamento può renderli meno prevedibili e meno controllabili di un normale agente che lavora con JSON!
# LangGraph è all'altro estremo dello spettro, ed è utile quando hai bisogno di "controllo" sull'esecuzione del tuo agente.
# In parole povere, se la tua applicazione prevede una serie di passaggi che devono essere orchestrati in un modo specifico, con decisioni prese a ogni punto di congiunzione, LangGraph fornisce la struttura di cui hai bisogno.

# Come funziona LangGraph?
# Nodes : I nodi rappresentano singole fasi di elaborazione (come la chiamata di un LLM, l'utilizzo di uno strumento o la presa di una decisione).
# Edges : definiscono le possibili transizioni tra i passaggi.
# State : È definito e gestito dall'utente e trasmesso tra i nodi durante l'esecuzione. Quando decidiamo quale nodo indirizzare successivamente, questo è lo stato attuale che prendiamo in considerazione.

# Building Blocks of LangGraph
# Per creare applicazioni con LangGraph, è necessario comprenderne i componenti principali. Esploriamo gli elementi fondamentali che compongono un'applicazione LangGraph.
#Un'applicazione in LangGraph inizia da un punto di ingresso e, a seconda dell'esecuzione, il flusso può passare a una funzione o all'altra fino a raggiungere la FINE.
# 1. State Lo stato è il concetto centrale di LangGraph. Rappresenta tutte le informazioni che fluiscono attraverso l'applicazione.
# 2. Nodes I nodi sono funzioni Python. Ogni nodo: Accetta lo stato come input Esegue un'operazione Restituisce aggiornamenti allo stato
# Ad esempio, i nodi possono contenere: Chiamate LLM: generare testo o prendere decisioni Chiamate strumenti: interagire con sistemi esterni Logica condizionale: determinare i passaggi successivi Intervento umano: ottenere input dagli utenti
# 3. Edges  collegano i nodi e definiscono i possibili percorsi attraverso il grafico
# 4. StateGraph è il contenitore che contiene l'intero flusso di lavoro dell'agente:

# Esempio: https://huggingface.co/agents-course/notebooks/blob/main/unit2/langgraph/mail_sorting.ipynb

# pip install langgraph langchain_openai

import datasets
from langchain.docstore.document import Document
from langchain_community.retrievers import BM25Retriever
from langchain.tools import Tool
from typing import TypedDict, Annotated
from langgraph.graph.message import add_messages
from langchain_core.messages import AnyMessage, HumanMessage, AIMessage
from langgraph.prebuilt import ToolNode
from langgraph.graph import START, StateGraph
from langgraph.prebuilt import tools_condition
#from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
from langchain_huggingface import HuggingFaceEndpoint
from langchain_community.tools import DuckDuckGoSearchRun
import gradio as gr
import os

#######################################
## Give Your Agent Access to the Web ##
#######################################

search_tool = DuckDuckGoSearchRun()
# results = search_tool.invoke("Who's the current President of France?")
# print(results)

##########################################
## Step 1: Load and Prepare the Dataset ##
##########################################
# Per prima cosa, dobbiamo trasformare i dati grezzi dei nostri ospiti in un formato ottimizzato per il recupero.
# Utilizzeremo la libreria di set di dati Hugging Face per caricare il set di dati e convertirlo in un elenco di oggetti Document dal modulo langchain.docstore.document.

# Load the dataset
guest_dataset = datasets.load_dataset("agents-course/unit3-invitees", split="train")

# Convert dataset entries into Document objects
docs = [
    Document(
        page_content="\n".join([
            f"Name: {guest['name']}",
            f"Relation: {guest['relation']}",
            f"Description: {guest['description']}",
            f"Email: {guest['email']}"
        ]),
        metadata={"name": guest["name"]}
    )
    for guest in guest_dataset
]

# Nel codice sopra: Carichiamo il dataset Convertiamo ogni voce ospite in un oggetto Documento con contenuto formattato Memorizziamo gli oggetti Documento in un elenco
# Ciò significa che tutti i nostri dati sono prontamente disponibili e possiamo iniziare a configurare il recupero.

##########################################
## Step 2: Create the Retriever Tool    ##
##########################################
# Ora creiamo uno strumento personalizzato che Alfred potrà utilizzare per cercare le informazioni sui nostri ospiti.
# Utilizzeremo BM25Retriever dal modulo langchain_community.retrievers per creare uno strumento di recupero.
# BM25Retriever è un ottimo punto di partenza per il recupero, ma per una ricerca semantica più avanzata, potresti prendere in considerazione l'utilizzo di retriever basati sull'incorporamento come quelli di sentence-transformers.

bm25_retriever = BM25Retriever.from_documents(docs)

def extract_text(query: str) -> str:
    """Retrieves detailed information about gala guests based on their name or relation."""
    results = bm25_retriever.invoke(query)
    if results:
        return "\n\n".join([doc.page_content for doc in results[:3]])
    else:
        return "No matching guest information found."

guest_info_tool = Tool(
    name="guest_info_retriever",
    func=extract_text,
    description="Retrieves detailed information about gala guests based on their name or relation."
)

# Analizziamo questo strumento passo dopo passo. 
# Il nome e la descrizione aiutano l'agente a capire quando e come utilizzare questo strumento. 
# I decoratori di tipo definiscono i parametri che lo strumento si aspetta (in questo caso, una query di ricerca).
# Utilizziamo BM25Retriever, un potente algoritmo di recupero del testo che non richiede incorporamenti. 
# Il metodo elabora la query e restituisce le informazioni più rilevanti sull'ospite.

############################################
## Step 2: Integrate the Tool with Alfred ##
############################################

# Infine, mettiamo insieme il tutto creando il nostro agente e dotandolo del nostro strumento personalizzato:

# Generate the chat interface, including the tools
llm = HuggingFaceEndpoint(
    repo_id="Qwen/Qwen2.5-Coder-32B-Instruct",
    huggingfacehub_api_token=os.getenv("HF_TOKEN"),
)

#chat = ChatHuggingFace(llm=llm, verbose=True)
#tools = [guest_info_tool,search_tool]
tools = [search_tool]
#chat_with_tools = chat.bind_tools(tools)

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

# def assistant(state: AgentState):
#     return {
#         "messages": [chat_with_tools.invoke(state["messages"])],
#     }

def assistant(state: AgentState):
    # Prendi l'ultimo messaggio umano
    human_msg = [msg for msg in state["messages"] if isinstance(msg, HumanMessage)][-1]
    # Chiama il modello LLM direttamente con il testo
    response_text = llm.invoke(human_msg.content)
    # Crea una risposta AIMessage
    ai_message = HumanMessage(content=response_text)  # o AIMessage se preferisci
    return {
        "messages": state["messages"] + [ai_message]
    }    

## 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")
alfred = builder.compile()

# Funzione Gradio
def run_agent(input_text):
    try:
        response = alfred.invoke({"messages": input_text})
        return response['messages'][-1].content
    except Exception as e:
        return f"Errore: {str(e)}"
    
iface = gr.Interface(
    fn=run_agent,
    inputs="text",
    outputs="text",
    title="Wikipedia AI Agent",
    description="Scrivi un argomento, l'agente recupera un riassunto da Wikipedia e lo spiega usando un LLM."
)

iface.launch()

# messages = [HumanMessage(content="Tell me about our guest named 'Lady Ada Lovelace'.")]
# response = alfred.invoke({"messages": messages})

# print("🎩 Alfred's Response:")
# print(response['messages'][-1].content)