File size: 3,364 Bytes
90a0674
 
 
 
ca1014c
90a0674
8fd8197
 
 
 
 
 
90a0674
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import os
import requests
import json
from calendar import monthrange
from smolagents.tools import Tool

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 = {'query': {'type': 'string', 'date range': 'The range of dates 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?"

    def get_nasa_neo_data(self, start_date: str, end_date: str) -> dict:
        """A function to get Near Earth Object data from NASA API. Try this function first before
        calling the fetch_neo_data_in_chunks function.
        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()
            return data
        else:
            print("Error: ", response.status_code)

    def split_month_into_chunks(self, year: int, month: int) -> list:
        days_in_month = monthrange(year, month)[1]
        first_day = 1
        last_day = 7
        week_dates = []
        while last_day <= days_in_month:
            start_date = f"{year}-{month}-{first_day}"
            end_date = f"{year}-{month}-{last_day}"
            week_dates.append((start_date, end_date))
            first_day += 7
            last_day += 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}-{days_in_month}"
            week_dates.append((start_date, end_date))
        return week_dates

    def fetch_neo_data_in_chunks(
        self, start_year: int, end_year: int
    ) -> list[tuple[str, list[dict]]]:
        """A function to fetch Near Earth Object data from NASA API in chunks. Call this function if the API
        returns an error when you try to pull data e.g. a 400 or 429.
        Args:
            start_year: An integer representing the start year of the data to be fetched.
            end_year: An integer representing the end year of the data to be fetched.
        Returns: A list of Near Earth Object data fetched from the API as a JSON-like dictionary.
        """
        neo_data = []
        for year in range(start_year, end_year + 1):
            for month in range(1, 13):
                for chunk in self.split_month_into_chunks(year, month):
                    start_date, end_date = chunk
                    print(f"Fetching data for {start_date} to {end_date}")
                    data = self.get_nasa_neo_data(
                        start_date=start_date, end_date=end_date
                    )
                    if data and "near_earth_objects" in data:
                        neo_data.extend(data["near_earth_objects"].items())
        neo_data_json = json.dumps(neo_data)
        return neo_data_json