"""Search module — now with smart auto-detection and graceful UX."""
from typing import Dict, Any
from config import CATEGORIES
from utils.budget_core import (
budget_exists_for_month,
get_current_month,
get_budget_for_month,
calculate_remaining_budget
)
from utils.date_utils import parse_flexible_month
from utils.category_detection import detect_category # NEW ✔
def register_search_tools(mcp):
# =====================================================================
# 1. SEARCH ITEMS WITHIN BUDGET (SMART)
# =====================================================================
@mcp.tool()
def search_items_within_budget(
item_query: str,
category: str = '',
use_total_budget: bool = False,
max_results: int = 10,
month: str = ''
) -> Dict[str, Any]:
# --------------------------------------------------
# Smart category detection (if user didn't specify)
# --------------------------------------------------
detected_meta = None
if not category:
det = detect_category(item_query)
if det["auto_assign"]:
category = det["detected_category"]
detected_meta = {
"auto_detected": True,
"category": category,
"confidence": det["confidence"],
}
else:
return {
"status": "needs_clarification",
"message": f"I found multiple possible categories for '{item_query}'.",
"top_guess": det["detected_category"],
"alternatives": det["alternatives"],
"next_steps": [
"Pick one of the suggestions.",
"Search using total budget instead.",
"Say 'ignore category' to bypass category mode."
]
}
# If category provided but invalid → soften
if category not in CATEGORIES:
guess = detect_category(category)
return {
"status": "invalid_category",
"message": f"'{category}' is not a valid budget category.",
"closest_match": guess["detected_category"],
"alternatives": guess["alternatives"],
"suggestion": f"Try '{guess['detected_category']}'?",
"valid_categories": list(CATEGORIES)
}
# --------------------------------------------------
# Month parsing
# --------------------------------------------------
try:
target_month = parse_flexible_month(month) if month else get_current_month()
except Exception as e:
return {
"status": "needs_clarification",
"message": f"Couldn't understand '{month}'.",
"error": str(e),
"suggestion": "Say something like 'next month', 'Jan 2025', or leave it empty."
}
# --------------------------------------------------
# If no budget exists → offer options
# --------------------------------------------------
if not budget_exists_for_month(target_month):
return {
"status": "no_budget",
"message": f"No budget exists for {target_month}.",
"options": [
"Create a budget now",
"Search without budget constraints",
"Set a temporary spending limit"
],
"detected_category": detected_meta
}
# --------------------------------------------------
# Compute remaining budget
# --------------------------------------------------
success, b = calculate_remaining_budget(target_month, category)
if not success:
return {
"status": "soft_error",
"message": f"Couldn't compute remaining budget for {target_month}.",
"suggestion": "Would you like to search without a budget?"
}
is_total_only = (b.get("budget_mode") == "total_only")
# Apply mode
if use_total_budget or is_total_only:
max_price = b.get("remaining_total", 0)
budget_type = "total budget"
else:
max_price = b.get("remaining", b.get("remaining_total", 0))
budget_type = f"{category} budget"
# --------------------------------------------------
# Handle zero or negative budget gently
# --------------------------------------------------
if max_price <= 0:
return {
"status": "budget_exhausted",
"message": f"Your {budget_type} is exhausted.",
"overspend": abs(max_price),
"options": [
"Search anyway (ignore budget)",
"Adjust your budget",
"Switch to total budget",
"Ask for cheaper alternatives"
],
"suggestion": f"Shall I still search for '{item_query}'?"
}
# --------------------------------------------------
# Normal response: Request search
# --------------------------------------------------
return {
"status": "requires_web_search",
"message": f"Searching '{item_query}' within your {budget_type} of ₹{max_price:,.2f}.",
"web_search_query": f"{item_query} price India under ₹{int(max_price)}",
"detected_category": detected_meta,
"budget_context": {
"month": target_month,
"max_price": max_price,
"budget_type": budget_type,
"currency": "INR",
"category": category
},
"filter_instructions": {
"max_price": max_price,
"currency": "INR",
"max_results": max_results,
"preferred_sites": ["amazon.in", "flipkart.com", "croma.com"]
}
}
# =====================================================================
# 2. FIND AFFORDABLE ALTERNATIVES
# =====================================================================
@mcp.tool()
def find_affordable_alternatives(
item_query: str,
current_price: float,
category: str = '',
max_results: int = 5
) -> Dict[str, Any]:
# Smart detect category when not provided
if not category:
det = detect_category(item_query, current_price)
if det["auto_assign"]:
category = det["detected_category"]
else:
category = None # fallback to total-budget search
target_month = get_current_month()
if not budget_exists_for_month(target_month):
return {
"status": "no_budget",
"message": "No budget found. I can still find cheaper alternatives.",
"search_query": f"affordable {item_query} India",
}
# Get remaining budgets
success, b = calculate_remaining_budget(target_month, category)
if not success:
return {
"status": "soft_error",
"message": "Couldn't determine remaining budget.",
"suggestion": "Search for alternatives without using budget?"
}
max_price = b.get("remaining_total", 0) if not category else b.get("remaining", b.get("remaining_total", 0))
# If item already affordable
if current_price <= max_price:
return {
"status": "within_budget",
"message": f"Great! ₹{current_price} fits within your budget.",
"remaining_after_purchase": max_price - current_price,
"next_steps": [
"Want me to compare with similar items?",
"Want reviews or alternatives?"
]
}
# Otherwise suggest cheaper alternatives
difference = current_price - max_price
return {
"status": "needs_cheaper_options",
"message": (
f"₹{current_price:.2f} exceeds your budget by ₹{difference:.2f}. "
"Let me find more affordable alternatives."
),
"suggested_search_query": f"affordable {item_query} India under ₹{int(max_price)}",
"filter_instructions": {
"max_price": max_price,
"keywords": ["budget", "cheap", "value"],
"currency": "INR",
"max_results": max_results
}
}
# =====================================================================
# 3. COMPARE ITEMS WITHIN BUDGET
# =====================================================================
@mcp.tool()
def compare_items_within_budget(
item1: str,
item2: str,
category: str = '',
use_total_budget: bool = False,
month: str = ''
) -> Dict[str, Any]:
# Auto-detect category if missing
if not category:
det1 = detect_category(item1)
det2 = detect_category(item2)
# If both agree → auto-assign
if det1["detected_category"] == det2["detected_category"] and det1["auto_assign"]:
category = det1["detected_category"]
# Month parsing
try:
target_month = parse_flexible_month(month) if month else get_current_month()
except Exception:
target_month = get_current_month()
# No budget? → Still compare!
if not budget_exists_for_month(target_month):
return {
"status": "no_budget",
"message": "No budget exists — but I can still compare the two items.",
"comparison_items": [item1, item2],
"queries": [
f"{item1} price India",
f"{item2} price India"
]
}
# Compute remaining budget
success, b = calculate_remaining_budget(target_month, category)
if not success:
return {
"status": "soft_error",
"message": "Couldn't compute remaining budget.",
"suggestion": "Compare without budget instead?"
}
# Select correct remaining
if use_total_budget or not category:
max_price = b.get("remaining_total", 0)
budget_type = "total budget"
else:
max_price = b.get("remaining", b.get("remaining_total", 0))
budget_type = f"{category} budget"
# If exhausted → ask user
if max_price <= 0:
return {
"status": "budget_exhausted",
"message": f"You have no remaining {budget_type}. Compare anyway?",
"options": ["Compare anyway", "Adjust budget", "Switch to total budget"],
"items": [item1, item2]
}
# Normal flow
return {
"status": "requires_comparison_search",
"message": f"Comparing '{item1}' vs '{item2}' within your {budget_type} of ₹{max_price:.2f}.",
"comparison_items": [item1, item2],
"search_queries": [
f"{item1} price India under ₹{int(max_price)}",
f"{item2} price India under ₹{int(max_price)}"
]
}
# =====================================================================
# 4. ESTIMATE PRICE RANGE
# =====================================================================
@mcp.tool()
def estimate_item_price_range(item_query: str, category: str = '') -> Dict[str, Any]:
# Smart detection
if not category:
det = detect_category(item_query)
if det["auto_assign"]:
category = det["detected_category"]
target_month = get_current_month()
# If no budget → fallback
if not budget_exists_for_month(target_month):
return {
"status": "no_budget",
"message": "No budget found — estimating typical price ranges.",
"search_query": f"{item_query} price range India"
}
success, b = calculate_remaining_budget(target_month, category)
if not success:
return {
"status": "soft_error",
"message": "Couldn't compute budget.",
"search_query": f"{item_query} price range India"
}
# determine limit
if category and category in CATEGORIES:
max_price = b.get("remaining", None)
else:
max_price = b.get("remaining_total", None)
q = f"{item_query} price range India"
if max_price and max_price > 0:
q += f" under ₹{int(max_price)}"
return {
"status": "requires_price_research",
"message": f"Estimating price range for '{item_query}'.",
"search_query": q,
"budget_context": {
"category": category or "total",
"remaining_budget": max_price,
"month": target_month
}
}
