app.py

from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool
from datetime import datetime, time, timezone
import requests
import pytz
import yaml
import os
from tools.final_answer import FinalAnswerTool

from Gradio_UI import GradioUI


@tool
def get_cloud_costs(start_date: str, end_date: str, team: str) -> float | None:
    """
    Requests cloud costs data from the Finout API for a specific team.

    Args:
        start_date: The start date for the cost data in the format 'YYYY-MM-DD'.
        end_date: The end date for the cost data in the format 'YYYY-MM-DD'.
        team: The name of the team to get costs for. Use get_teams() to get a list of valid team names.

    Returns:
        float: The cost amount in USD for the specified team.
        None: If the team is not found in the response. This might indicate that either the team name is incorrect
              or there are no costs recorded for this team in the specified date range.

    Raises:
        requests.RequestException: If the API request fails.
    """
    # Convert strings to datetime objects
    start_dt = datetime.strptime(start_date, '%Y-%m-%d').replace(tzinfo=timezone.utc)
    end_dt = datetime.strptime(end_date, '%Y-%m-%d').replace(tzinfo=timezone.utc)
    
    # If same day, set start to beginning of day and end to end of day
    if start_date == end_date:
        start_dt = datetime.combine(start_dt.date(), time.min, tzinfo=timezone.utc)  # Sets time to 00:00:00 UTC
        end_dt = datetime.combine(end_dt.date(), time.max, tzinfo=timezone.utc)    # Sets time to 23:59:59.999999 UTC
    
    # Convert to Unix timestamps in milliseconds
    start_time = int(start_dt.timestamp() * 1000)
    end_time = int(end_dt.timestamp() * 1000)

    url = "https://app.finout.io/v1/cost/query-by-view"
    headers = {
        "x-finout-client-id": os.getenv('FINOUT_ACCESS'),
        "x-finout-secret-key": os.getenv('FINOUT_SECRET')
    }
    body = {
        "viewId": os.getenv('VIEW_PER_TEAM'),
        "date": {
            "unixTimeMillSecondsStart": start_time,
            "unixTimeMillSecondsEnd": end_time
        },
        "costType": "AmortizedCost"
    }

    response = requests.post(url, headers=headers, json=body)
    response.raise_for_status()  # Raise an exception for bad status codes
    
    # Extract the cost for the team
    for item in response.json()['data']:
        if item['name'] == team:
            return item['data'][0]['cost']
    return None


@tool
def get_teams(date: str) -> list[str]:   
    """
    Retrieves the list of team names from the Finout API for a specific date to prepare the request for the get_cloud_costs function.

    Args:
        date: The date for the data in the format 'YYYY-MM-DD'.

    Returns:
        list[str]: A list of team names that can be used as input for get_cloud_costs().
                  The list will be empty if no teams have cost data for the specified date.

    Raises:
        requests.RequestException: If the API request fails.

    
    Example:
        >>> teams = get_teams('2024-03-20')
        >>> for team in teams:
        ...     cost = get_cloud_costs('2024-03-20', '2024-03-20', team)
        ...     print(f"{team}: ${cost}")
    """
    # Convert string to datetime object
    date_dt = datetime.strptime(date, '%Y-%m-%d').replace(tzinfo=timezone.utc)
    start_dt = datetime.combine(date_dt.date(), time.min, tzinfo=timezone.utc)
    end_dt = datetime.combine(date_dt.date(), time.max, tzinfo=timezone.utc)

    url = "https://app.finout.io/v1/cost/query-by-view"
    headers = {
        "x-finout-client-id": os.getenv('FINOUT_ACCESS'),
        "x-finout-secret-key": os.getenv('FINOUT_SECRET')
    }
    body = {
        "viewId": os.getenv('VIEW_PER_TEAM'),
        "date": {
            "unixTimeMillSecondsStart": int(start_dt.timestamp() * 1000),
            "unixTimeMillSecondsEnd": int(end_dt.timestamp() * 1000)
        },
        "costType": "AmortizedCost"
    }

    response = requests.post(url, headers=headers, json=body)
    response.raise_for_status()  # Raise an exception for bad status codes
    
    return [item['name'] for item in response.json()['data']]

@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """A tool that fetches the current local time in a specified timezone.
    Args:
        timezone: A string representing a valid timezone (e.g., 'America/New_York').
    """
    try:
        # Create timezone object
        tz = pytz.timezone(timezone)
        # Get current time in that timezone
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        return f"The current local time in {timezone} is: {local_time}"
    except Exception as e:
        return f"Error fetching time for timezone '{timezone}': {str(e)}"


final_answer = FinalAnswerTool()

# If the agent does not answer, the model is overloaded, please use another model or the following Hugging Face Endpoint that also contains qwen2.5 coder:
# model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud' 

model = HfApiModel(
max_tokens=2096,
temperature=0.5,
model_id='Qwen/Qwen2.5-Coder-32B-Instruct',# it is possible that this model may be overloaded
custom_role_conversions=None,
)

with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)
    
agent = CodeAgent(
    model=model,
    tools=[final_answer, get_current_time_in_timezone, get_cloud_costs, get_teams], ## add your tools here (don't remove final answer)
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)


GradioUI(agent).launch()