File size: 5,733 Bytes
90a0674
 
 
ca1014c
c3a9d5d
 
90a0674
320acdc
 
c3a9d5d
 
8fd8197
 
c3a9d5d
 
 
 
 
 
 
 
 
 
7384219
c3a9d5d
 
90a0674
 
c3a9d5d
5bd2d87
320acdc
c3a9d5d
 
 
 
320acdc
c3a9d5d
 
90a0674
c3a9d5d
 
90a0674
 
 
 
 
 
 
 
 
 
c3a9d5d
90a0674
 
 
 
c3a9d5d
 
 
90a0674
 
c3a9d5d
 
90a0674
c3a9d5d
90a0674
 
c3a9d5d
90a0674
 
 
c3a9d5d
90a0674
 
 
c3a9d5d
 
 
 
 
 
 
 
 
 
 
90a0674
c3a9d5d
 
 
 
 
 
 
 
320acdc
 
 
c3a9d5d
 
 
 
 
320acdc
90a0674
 
c3a9d5d
 
90a0674
c3a9d5d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
320acdc
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
import os
import requests
from calendar import monthrange
from smolagents.tools import Tool
from typing import Generator
from dotenv import load_dotenv

load_dotenv()


class NasaNeoDataFetcher(Tool):
    name = "nasa_neo_data_fetcher"
    description = "Retrieves Neo data from the NASA API for the dates in your query then returns the available Neo data."
    inputs = {
        "start_date": {
            "type": "string",
            "description": "The start date of the range to query.",
        },
        "end_date": {
            "type": "string",
            "description": "The end date of the range to query.",
        },
    }
    output_type = "string"

    def __init__(self):
        self.api_key = os.getenv("NASA_API_KEY")
        self.root_url = "https://api.nasa.gov/neo/rest/v1/feed?"
        self.is_initialized = False

    def forward(self, start_date: str, end_date: str) -> list[tuple[str, dict]]:
        """A function to fetch Near Earth Object data from NASA API for a given date range.
        Args:
            start_date: A string representing the start date of the data to be fetched.
            end_date: A string representing the end date of the data to be fetched.
        Returns: The data fetched from the API as a list of tuples containing the date and the JSON-like dictionaries.
        """
        return self._fetch_neo_data_in_chunks(start_date, end_date)

    def _get_nasa_neo_data(self, start_date: str, end_date: str) -> dict:
        """A function to get Near Earth Object data from NASA API for a given date range.
        Args:
            start_date: A string representing the start date of the data to be fetched.
            end_date: A string representing the end date of the data to be fetched.
        Returns: The data fetched from the API as a JSON-like dictionary.
        """
        params = f"start_date={start_date}&end_date={end_date}&api_key={self.api_key}"
        url = self.root_url + params
        response = requests.get(url)
        if repr(response.status_code) == "200":
            data = response.json()
            print(f"Successfully fetched data for {start_date} to {end_date}")
            return data
        else:
            print("Error: ", response.status_code)

    def _split_date_range_into_chunks(
        self, year: int, month: int, first_day: int, last_day: int
    ) -> list:
        days_in_month = monthrange(year, month)[1]
        week_dates = []
        last_day_in_chunk = first_day + 6
        while last_day_in_chunk <= last_day:
            start_date = f"{year}-{month}-{first_day}"
            end_date = f"{year}-{month}-{last_day_in_chunk}"
            week_dates.append((start_date, end_date))
            first_day += 7
            last_day_in_chunk += 7
        remaining_days = days_in_month - first_day + 1
        if remaining_days > 0:
            start_date = f"{year}-{month}-{first_day}"
            end_date = f"{year}-{month}-{last_day}"
            week_dates.append((start_date, end_date))
        return week_dates

    def _split_date_into_ymd(self, date: str) -> tuple[int, int, int]:
        split_date = date.split("-")
        year = int(split_date[0])
        month = int(split_date[1])
        day = int(split_date[2])
        return year, month, day

    def _get_neo_data_chunks(
        self, *args: tuple[str, str]
    ) -> Generator[dict, None, None]:
        """A function that splits a date range into chunks of 7 days.
        Args:
            *args: A variable number of date ranges to split.
        Returns: A generator yielding NASA NEO data dictionaries.
        """
        for chunk in self._split_date_range_into_chunks(*args):
            start_date, end_date = chunk
            data = self._get_nasa_neo_data(start_date=start_date, end_date=end_date)
            yield data

    def _fetch_neo_data_in_chunks(
        self, start_date: str, end_date: str
    ) -> list[tuple[str, dict]]:
        """A function that fetches Near Earth Object data from NASA API in chunks of 7 days. NB: The API
        returns an error when you try to pull data covering longer than 7 days.
        Args:
            start_date: A string representing the start date of the data to be fetched.
            end_date: A string representing the end date of the data to be fetched.
        Returns: A list of tuples containing the date and the JSON-like dictionaries of Near Earth Object data fetched from the API.
        """
        neo_data = []
        start_year, start_month, start_day = self._split_date_into_ymd(start_date)
        end_year, end_month, end_day = self._split_date_into_ymd(end_date)
        for year in range(start_year, end_year + 1):
            for month in range(start_month, end_month + 1):
                days_in_month = monthrange(year, month)[1]
                if month == start_month and start_month != end_month:
                    args = (year, month, start_day, days_in_month)
                    data = list(self._get_neo_data_chunks(*args))

                elif month == end_month and start_day != 1:
                    args = (year, month, start_day, end_day)
                    data = list(self._get_neo_data_chunks(*args))

                elif month == end_month:
                    args = (year, month, 1, end_day)
                    data = list(self._get_neo_data_chunks(*args))

                else:
                    args = (year, month, 1, days_in_month)
                    data = list(self._get_neo_data_chunks(*args))

                if data:
                    for item in data:
                        if "near_earth_objects" in item:
                            neo_data.extend(item["near_earth_objects"].items())
        return neo_data