import gradio as gr
import random
from smolagents import GradioUI, CodeAgent, InferenceClientModel
from smolagents import LiteLLMModel
from smolagents import ToolCallingAgent, PythonInterpreterTool, PromptTemplates
import os
from PIL import Image
#from smolagents.src.smolagents.prompts 

# Import our custom tools from their modules
from tools import evaluate_consumption, evaluate_consumption_example
from retriever import FrugalAI_methods

#HF_TOKEN = os.environ.get("HF_TOKEN") 
try:
    API_KEY = os.environ.get("API_KEY") 
except:
    print("Please provide the Anthropic API_KEY in the secret variables")
# Initialize the Hugging Face model
model = LiteLLMModel(
    model_id="anthropic/claude-3-5-sonnet-latest",
    temperature=0.2,
    api_key=API_KEY
)

evaluate_consumption = evaluate_consumption()
evaluate_consumption_example = evaluate_consumption_example()
FrugalAI_methods=FrugalAI_methods()

# Create agent with all the tools
fruggy = CodeAgent(
    tools=[evaluate_consumption, evaluate_consumption_example, FrugalAI_methods], 
    model=model,
    additional_authorized_imports=['os'],
    add_base_tools=True,  # Add any additional base tools
    planning_interval=10   
)

with gr.Blocks(theme=gr.themes.Ocean()) as demo:
    with gr.Row():
        with gr.Column(scale=1):
            image = gr.Image(value=Image.open("frugal.jpg"), label="Frugal Image")
            hello = gr.Interface(
                fn=fruggy,
                inputs="text",
                outputs="text",
                title="Frugalize it!",
                examples = [
                "You are an AI agent that receives Python code from a manager. "
                "Your task is to test its resource consumption using a custom tool. If no code is provided, use evaluate_consumption_example. If a code is provided by the manager, use evaluate_consumption."
                "Then, propose frugal alternatives, such as pruning, using your dedicated tool — always use it. "
                "Additionally, feel free to search the internet for frugal methods and present them to the manager. "
                "Examples include Knowledge Distillation, Transfer Learning, and others. "
                "Show me what you can do using your custom tool evaluate_consumption_example, and provide examples based on the code in that tool to optimize it for frugality. "
                "Use both your tools and web searches if needed. If implementation guidance is given by your custom tools, always give it in your final answer to the manager.",
            
                "You are an AI agent that receives Python code from a manager. "
                "Your task is to test its resource consumption using a custom tool. If no code is provided, use evaluate_consumption_example. If a code is provided by the manager, use evaluate_consumption."
                "Then, propose frugal alternatives, such as pruning, using your dedicated tool — always use it. "
                "Additionally, feel free to search the internet for frugal methods and present them to the manager. "
                "Examples include Knowledge Distillation, Transfer Learning, and others. If implementation guidance is given by your custom tools, always give it in your final answer to the manager. "
                "Here is my code: "
                "import anthropic "
                "API_KEY = os.environ.get('API_KEY') "
                "client = anthropic.Anthropic(api_key=API_KEY) "
                "message = client.messages.create(model='claude-sonnet-4-20250514',max_tokens=1024,messages=[{'role': 'user', 'content': 'Hello, Claude'}]) "
                "print(message.content)     "
                "Please give me frugal alternatives."
            ],
                description="Share your Python code with this AI agent! It will track its CO2 emissions using CodeCarbon and recommend greener, frugal AI alternatives."
            )
    
if __name__ == "__main__":
    demo.launch(mcp_server=True, share=True)