import networkx as nx
from smolagents import CodeAgent, HfApiModel, tool, GradioUI

# Define a tool for analyzing the Florentine Families graph
@tool
def analyze_florentine_graph(metric: str) -> str:
    """
    Analyzes the Florentine Families graph based on the chosen centrality metric.

    Args:
        metric: The centrality metric to calculate. Valid options: 'degree', 'betweenness', or 'closeness'.

    Returns:
        A textual analysis of the chosen metric for the graph with formatted numerical results.
    """
    graph = nx.florentine_families_graph()

    if metric == "degree":
        centrality = nx.degree_centrality(graph)
    elif metric == "betweenness":
        centrality = nx.betweenness_centrality(graph)
    elif metric == "closeness":
        centrality = nx.closeness_centrality(graph)
    else:
        return "Invalid metric. Please choose 'degree', 'betweenness', or 'closeness'."

    analysis = f"Analysis of {metric} centrality:\n"
    for node, value in centrality.items():
        analysis += f"- {node}: {value:.3f}\n"
    return analysis

# Initialize the agent with proper configuration
model = HfApiModel("meta-llama/Llama-3.3-70B-Instruct")
agent = CodeAgent(
    tools=[analyze_florentine_graph],
    model=model,
    additional_authorized_imports=["networkx"],
    add_base_tools=True
)

# Launch the Gradio interface using smolagents' built-in UI component
GradioUI(agent).launch()