Spaces:
Runtime error
Runtime error
| """Small USDA FoodData Central API client used by the nutrition pipeline.""" | |
| from __future__ import annotations | |
| import json | |
| import re | |
| from pathlib import Path | |
| from typing import Any | |
| from urllib.parse import urlencode | |
| from urllib.request import urlopen | |
| FDC_SEARCH_URL = "https://api.nal.usda.gov/fdc/v1/foods/search" | |
| NUTRIENT_NUMBERS = { | |
| "kcal": "208", | |
| "protein": "203", | |
| "fat": "204", | |
| "carbs": "205", | |
| } | |
| class FdcClient: | |
| def __init__( | |
| self, | |
| api_key: str, | |
| cache_path: Path, | |
| data_type: str = "Foundation,SR Legacy,FNDDS", | |
| page_size: int = 10, | |
| timeout: int = 30, | |
| ) -> None: | |
| self.api_key = api_key | |
| self.cache_path = cache_path | |
| self.data_type = data_type | |
| self.page_size = page_size | |
| self.timeout = timeout | |
| self.cache = self._load_cache() | |
| def search_first(self, query: str) -> dict[str, Any] | None: | |
| cache_key = f"ranked-v2|{self.data_type}|{self.page_size}|{query}".lower() | |
| if cache_key in self.cache: | |
| return self.cache[cache_key] | |
| params = { | |
| "api_key": self.api_key, | |
| "query": query, | |
| "pageSize": self.page_size, | |
| "dataType": self.data_type, | |
| } | |
| url = f"{FDC_SEARCH_URL}?{urlencode(params)}" | |
| with urlopen(url, timeout=self.timeout) as response: | |
| payload = json.loads(response.read().decode("utf-8")) | |
| foods = payload.get("foods", []) | |
| result = choose_best_food(query, foods) | |
| self.cache[cache_key] = result | |
| return result | |
| def save_cache(self) -> None: | |
| self.cache_path.parent.mkdir(parents=True, exist_ok=True) | |
| self.cache_path.write_text( | |
| json.dumps(self.cache, indent=2, sort_keys=True), | |
| encoding="utf-8", | |
| ) | |
| def _load_cache(self) -> dict[str, Any]: | |
| if not self.cache_path.exists(): | |
| return {} | |
| return json.loads(self.cache_path.read_text(encoding="utf-8")) | |
| def extract_macros(food: dict[str, Any] | None) -> dict[str, float | None]: | |
| values: dict[str, float | None] = {name: None for name in NUTRIENT_NUMBERS} | |
| if food is None: | |
| return values | |
| nutrients = food.get("foodNutrients", []) | |
| by_number = { | |
| str(nutrient.get("nutrientNumber")): nutrient.get("value") | |
| for nutrient in nutrients | |
| if nutrient.get("nutrientNumber") is not None | |
| } | |
| for macro_name, nutrient_number in NUTRIENT_NUMBERS.items(): | |
| value = by_number.get(nutrient_number) | |
| values[macro_name] = round(float(value), 2) if value is not None else None | |
| return values | |
| def choose_best_food(query: str, foods: list[dict[str, Any]]) -> dict[str, Any] | None: | |
| if not foods: | |
| return None | |
| ranked = sorted( | |
| enumerate(foods), | |
| key=lambda item: food_match_score(query, item[1], item[0]), | |
| reverse=True, | |
| ) | |
| return ranked[0][1] | |
| def food_match_score(query: str, food: dict[str, Any], index: int) -> float: | |
| description = str(food.get("description", "")).lower() | |
| category = str(food.get("foodCategory", "")).lower() | |
| haystack = f"{description} {category}" | |
| tokens = tokenize(query) | |
| score = 0.0 | |
| cooking_words = {"cooked", "grilled", "boiled", "raw", "roasted", "fried", "steamed"} | |
| for token in tokens: | |
| if token in haystack: | |
| score += 0.25 if token in cooking_words else 2.0 | |
| if "skin" not in tokens: | |
| if description.startswith("chicken, skin") or " skin (" in description: | |
| score -= 4.0 | |
| # Keep API ordering as a small tie-breaker after semantic token matches. | |
| score -= index * 0.01 | |
| return score | |
| def tokenize(text: str) -> list[str]: | |
| return [token for token in re.findall(r"[a-z0-9]+", text.lower()) if len(token) > 2] | |