from strands import Agent
from strands.models import LlamaAPIModel
from strands import Agent, tool

from rich.pretty import pprint

def main():
    # Llama API
    llama_model = LlamaAPIModel(
        model_id="Llama-4-Maverick-17B-128E-Instruct-FP8",
    )
    agent = Agent(model=llama_model)
    response = agent("Tell me about Agentic AI")

    print()
    pprint(response)


def main_with_get_weather_tool():
    @tool
    def get_weather(location: str) -> str:
        """Get the weather for a location.
        
        :param location: The city, state, or country for which to fetch the temperature. 
        """
        return f"The weather in {location} is sunny."

    llama_model = LlamaAPIModel(
        model_id="Llama-4-Maverick-17B-128E-Instruct-FP8",
    )
    agent = Agent(model=llama_model, tools=[get_weather])
    response = agent("What is the weather in San Francisco?")

    print()
    pprint(response)

def main_with_word_count_tool():
    @tool
    def word_count(text: str) -> int:
        """Count the number of words in a text.
        
        :param text: The text to count the words of. 
        """
        return len(text.split())

    llama_model = LlamaAPIModel(
        model_id="Llama-4-Maverick-17B-128E-Instruct-FP8",
    )
    agent = Agent(model=llama_model, tools=[word_count])
    response = agent("How many words are in this sentence?")

    print()
    pprint(response)



if __name__ == "__main__":
    # main()
    main_with_word_count_tool()
    # main_with_get_weather_tool()