utils/weather_forecast.py

from datetime import datetime
from typing import Optional

import httpx
from pydantic import BaseModel, Field
from simpleaichat import AIChat

from utils.env import get_from_env

WEATHERAPI_BASE_URL = "http://api.weatherapi.com/v1"


class WeatherAPI(BaseModel):
    """Wrapper for WeatherAPI

    Docs for using:

    1. Go to https://www.weatherapi.com/ and sign up for an API key
    2. Save your API WEATHERAPI_API_KEY env variable
    """

    weatherapi_key: Optional[str] = get_from_env("WEATHERAPI_API_KEY")
    forecast_response: dict = None
    warning: str = ""

    def _format_weather_info(self):
        current = self.forecast_response["current"]
        forecast = self.forecast_response["forecast"]["forecastday"][0]
        location_name = self.forecast_response["location"]["name"]
        current_status = current["condition"]["text"]
        wind_speed = current["wind_kph"]
        wind_degree = current["wind_degree"]
        humidity = current["humidity"]
        cloud = current["cloud"]
        feels_like = current["feelslike_c"]
        current_temp = current["temp_c"]
        min_temp = forecast["day"]["mintemp_c"]
        max_temp = forecast["day"]["maxtemp_c"]
        total_precip_mm = forecast["day"]["totalprecip_mm"]
        total_snow_cm = forecast["day"]["totalsnow_cm"]
        chance_of_rain = forecast["day"]["daily_chance_of_rain"]
        chance_of_snow = forecast["day"]["daily_chance_of_snow"]
        status = forecast["day"]["condition"]["text"]

        is_day = current["is_day"]
        sunrise = forecast["astro"]["sunrise"]
        sunset = forecast["astro"]["sunset"]
        # Convert the sunrise and sunset times to datetime objects
        sunrise_time = datetime.strptime(sunrise, "%H:%M %p").time()
        # sunset_time = datetime.strptime(sunset, "%I:%M %p").time()

        # Get the current time
        now = datetime.now().time()

        # Check if the current time is before or after sunrise
        if now < sunrise_time:
            next_event = f"The sun will rise at {sunrise}."
        else:
            next_event = f"The sun will set at {sunset}."
        rain_times = []
        snow_times = []
        for i, c in enumerate(forecast["hour"]):
            if c["will_it_rain"]:
                rain_times.append(c["time"])
            if c["will_it_snow"]:
                snow_times.append(c["time"])

        results = (
            f"In {location_name}, the current weather is as follows:\n"
            f"Detailed status: {current_status}\n"
            f"Wind speed: {wind_speed} kph, direction: {wind_degree}°\n"
            f"Humidity: {humidity}%\n"
            f"Temperature: \n"
            f"  - Current: {current_temp}°C\n"
            f"  - High: {max_temp}°C\n"
            f"  - Low: {min_temp}°C\n"
            f"  - Feels like: {feels_like}°C\n"
            f"Cloud cover: {cloud}%\n"
            f"Precipitation:\n"
            f"  - Total precipitation: {total_precip_mm} mm\n"
            f"  - Total snowfall: {total_snow_cm} cm\n"
            f"Chance of precipitation:\n"
            f"  - Chance of rain: {chance_of_rain}%\n"
            f"  - Chance of snow: {chance_of_snow}%\n"
            f"Weather status for the day: {status}\n"
            + (
                f"It is currently daytime.\n"
                if is_day
                else f"It is currently nighttime.\n"
            )
            + next_event
        )

        if rain_times:
            results += "There is a chance of rain at the following times:\n"
            for time in rain_times:
                results += f"- {time}\n"

        if snow_times:
            results += "There is a chance of snow at the following times:\n"
            for time in snow_times:
                results += f"- {time}\n"

        return f"{self.warning}\n{results}"

    def get_forecast(self, location: dict):
        """Get the current weather information for a specified location."""
        forecast_url = f"{WEATHERAPI_BASE_URL}/forecast.json"
        forecast_params = {
            "key": self.weatherapi_key,
            "q": location["city"],
            "format": "json",
            "days": 1,
        }

        r = httpx.get(forecast_url, params=forecast_params).json()
        self.forecast_response = r

    def run(self, query) -> str | None:
        location = get_location(query)
        if not location:
            self.warning = (
                "Could not identify any location in the query. Defaulted to Halifax."
            )
            location = {"city": "Halifax", "country": "CA"}
        self.get_forecast(location)
        weather_info = self._format_weather_info()
        return weather_info


class GetLocationMetadata(BaseModel):
    """Location information"""

    city: str = Field(description="The city of the location.")
    state: int = Field(
        description="The state or province of the location. Must be the full state name"
    )
    state_code: int = Field(
        description="The state or province of the location. Must be a 2-char string."
    )
    country: str = Field(
        description="The country of the location. Country must be a 2-char string"
    )


def get_location(query: str) -> dict | None:
    params = {"temperature": 0.0, "max_tokens": 100}
    system_prompt = (
        "You reply ONLY with the location information. If no location is detected, respond with None in "
        "each field"
    )
    ai = AIChat(system=system_prompt, params=params)
    # noinspection PyTypeChecker
    location: dict = ai(query, output_schema=GetLocationMetadata)
    if location["city"] == "None":
        return None
    return location