Tutorial: Restaurant Reservation Agent¶

"""before_model_callback for the Bella Notte slot filling agent."""

import random
import string


# ── Executor functions ────────────────────────────────────────────────────────

def _find_available_times(filled: dict) -> dict:
    """Mock availability lookup. Replace with real API in production."""
    # Simulate different availability by date
    date = filled.get('preferred_date', '')
    times_by_day = {
        'Mon': '6:00 PM, 7:30 PM, 9:00 PM',
        'Tue': '6:00 PM, 7:00 PM, 8:30 PM',
        'Wed': '5:30 PM, 7:00 PM, 8:00 PM, 9:00 PM',
    }
    import datetime
    try:
        day_abbr = datetime.date.fromisoformat(date).strftime('%a')
        times = times_by_day.get(day_abbr, '6:00 PM, 7:00 PM, 8:00 PM')
    except ValueError:
        times = '6:00 PM, 7:00 PM, 8:00 PM'
    return {'success': True, 'times': times}


def _book_reservation(filled: dict) -> dict:
    """Mock booking API. Replace with real API in production."""
    confirmation = 'BN-' + ''.join(random.choices(string.digits, k=6))
    return {'success': True, 'confirmation': confirmation}


# ── DAG helpers ───────────────────────────────────────────────────────────────

def _task_inputs_ready(sm: dict, task_name: str) -> bool:
    """Check whether all inputs for a task are in sm['filled']."""
    filled = sm.get('filled', {})
    task_results = sm.get('task_results', {})
    inputs = {
        'FindAvailableTimes': ['party_size', 'preferred_date'],
        'BookReservation': [
            'party_size', 'preferred_date', 'selected_time',
            'guest_name', 'special_requests',
        ],
    }
    if task_name in task_results:
        return False  # Already ran successfully
    return all(s in filled for s in inputs.get(task_name, []))


def _next_question(sm: dict) -> tuple:
    """Return (question_text, slot_name) for the first unfilled slot."""
    filled = sm.get('filled', {})

    order = [
        ('party_size',    'How many guests will be joining you?'),
        ('preferred_date', 'What date were you thinking?'),
    ]

    if 'available_times' in filled:
        times = filled['available_times']
        order.append(
            ('selected_time', f'We have openings at {times}. Which time works for you?')
        )

    order += [
        ('guest_name',       'What name should I put the reservation under?'),
        ('special_requests', 'Do you have any special requests? '
                             "(If not, just say 'none'.)"),
    ]

    for slot_name, question in order:
        if slot_name not in filled:
            return question, slot_name

    return 'All information has been collected.', None


# ── Main callback ─────────────────────────────────────────────────────────────

def before_model_callback(callback_context, llm_request):  # (1)
    sm = callback_context.state.get('sm', {})
    if not sm:
        sm = {
            'filled': {}, 'pending': {}, 'task_results': {},
            '_retries': {}, '_slot_errors': [],
            '_system_message': '', 'status': 'in_progress',
        }
        callback_context.state['sm'] = sm

    # 1. Terminal state guard — let the LLM respond freely
    if sm.get('status') in ('complete', 'escalated'):
        return None

    # 2. Handle slot validation errors from setter tools
    errors = sm.pop('_slot_errors', [])
    if errors:
        error_messages = {
            ('party_size', 'out_of_range'):    'We accept parties of 1 to 8 guests. How many guests will be dining?',
            ('party_size', 'parse_error'):     "I didn't catch the number of guests. How many will be in your party?",
            ('preferred_date', 'past_date'):   'That date is in the past. Could you provide a future date?',
            ('preferred_date', 'invalid_format'): "I couldn't parse that date. Could you say it in another way?",
            ('selected_time', 'not_available'): 'That time is not available. Please choose from the options shown.',
            ('selected_time', 'prereq_not_met'): 'Please wait for available times to be presented first.',
        }
        msgs = []
        for err in errors:
            key = (err.get('slot'), err.get('code'))
            msg = error_messages.get(key, 'There was an issue with that value. Please try again.')
            msgs.append(msg)
        combined = ' '.join(msgs)
        sm['_system_message'] = combined
        if llm_request.contents and len(llm_request.contents) > 1:
            from google.cloud.aiplatform_v1beta1.types.content import Part
            from google.cloud.aiplatform_v1beta1.types.prediction_service import GenerateContentResponse as LlmResponse
            return LlmResponse.from_parts(parts=[Part.from_text(text=combined)])
        return None

    # 3. Check if FindAvailableTimes should fire
    task_fired = False
    if _task_inputs_ready(sm, 'FindAvailableTimes'):
        result = _find_available_times(sm['filled'])
        sm['task_results']['FindAvailableTimes'] = result
        if result.get('success'):
            sm['filled']['available_times'] = result['times']
            task_fired = True

    # 4. Check if BookReservation should fire
    if _task_inputs_ready(sm, 'BookReservation'):
        result = _book_reservation(sm['filled'])
        sm['task_results']['BookReservation'] = result
        if result.get('success'):
            sm['filled']['confirmation_number'] = result['confirmation']
            sm['status'] = 'complete'
            msg = (
                f"Wonderful! Your reservation at Bella Notte is confirmed. "
                f"Your confirmation number is {result['confirmation']}. "
                f"We look forward to welcoming you!"
            )
            sm['_system_message'] = msg
            task_fired = True

    # 5. Set _system_message with the next question
    if not sm.get('status') == 'complete':
        next_q, _ = _next_question(sm)
        sm['_system_message'] = next_q

    # 6. Preempt the LLM when a task just fired (skip model generation)
    if task_fired and llm_request.contents and len(llm_request.contents) > 1:
        from google.cloud.aiplatform_v1beta1.types.content import Part
        from google.cloud.aiplatform_v1beta1.types.prediction_service import GenerateContentResponse as LlmResponse
        return LlmResponse.from_parts(
            parts=[Part.from_text(text=sm['_system_message'])]
        )

    return None