Switched to new neo tool code

tools/nasa_neo_data_fetcher.py

@@ -3,31 +3,42 @@ import requests
 import json
 from calendar import monthrange
 from smolagents.tools import Tool
+from typing import Generator
+from dotenv import load_dotenv
 
-class NASANeoDataFetcher(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', 'description': 'The range of dates to query.'}}
+    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, query):
+
+    def __init__(self):
+        load_dotenv()
         self.api_key = os.getenv("NASA_API_KEY")
         self.root_url = "https://api.nasa.gov/neo/rest/v1/feed?"
-        try:
-            get_data = self.get_nasa_neo_data(query)
-        except Error:
-            get_data = self.fetch_neo_data_in_chunks(query)
-        data = get_data(query)
-            
+        self.is_initialized = False
 
-    def forward(self, query):
-        results = self.data(query)
-        return results
-        
+    def forward(self, start_date: str, end_date: str) -> str:
+        """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 JSON-like dictionary.
+        """
+        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. Try this function first before
-        calling the fetch_neo_data_in_chunks function.
+    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.
@@ -38,49 +49,83 @@ class NASANeoDataFetcher(Tool):
         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:
-            fetch_neo_data_in_chunks(start_date, end_date)
             print("Error: ", response.status_code)
 
-    def split_month_into_chunks(self, year: int, month: int) -> list:
+    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]
-        first_day = 1
-        last_day = 7
         week_dates = []
-        while last_day <= days_in_month:
+        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}"
+            end_date = f"{year}-{month}-{last_day_in_chunk}"
             week_dates.append((start_date, end_date))
             first_day += 7
-            last_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}-{days_in_month}"
+            end_date = f"{year}-{month}-{last_day}"
             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.
+    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:
-            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.
+            *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) -> str:
+        """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 JSON-like dictionary 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(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())
+            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())
         neo_data_json = json.dumps(neo_data)
         return neo_data_json