from langchain_core.tools import tool from pydantic.v1 import BaseModel, Field from datetime import datetime from database.db_handler import ( add_appointment, search_appointments, delete_appointments, update_appointment, get_all_appointments ) # --- Pydantic Models for Tool Inputs --- # (These are unchanged) class ScheduleAppointmentInput(BaseModel): customer_name: str = Field(description="The full name of the customer.") customer_phone: str = Field(description="The contact phone number of the customer.") appointment_datetime: str = Field(description="The date and time in 'YYYY-MM-DD HH:MM:SS' format.") service_type: str = Field(description="The type of service, e.g., 'Dental Checkup', 'Haircut'.") class SearchAppointmentInput(BaseModel): criteria: str = Field(description="The field to search by: 'customer_name', 'service_type', or 'date'.") value: str = Field(description="The value to search for. For date, use 'YYYY-MM-DD' format.") class DeleteAppointmentInput(BaseModel): criteria: str = Field( description="The field to identify appointments to delete: 'customer_name', 'service_type', or 'date'.") value: str = Field(description="The value to identify the appointments by.") class UpdateAppointmentInput(BaseModel): identifier_value: str = Field(description="The name of the customer whose appointment needs updating.") field_to_update: str = Field( description="The specific field to change, e.g., 'customer_phone', 'appointment_datetime'.") new_value: str = Field(description="The new value for the specified field.") # --- Tool Definitions --- @tool(args_schema=ScheduleAppointmentInput) def schedule_appointment(customer_name: str, customer_phone: str, appointment_datetime: str, service_type: str) -> str: """Schedules a new appointment.""" # <-- ADDED THIS DOCSTRING BACK try: datetime.strptime(appointment_datetime, '%Y-%m-%d %H:%M:%S') except ValueError: return "Error: Invalid datetime format. Please use 'YYYY-MM-DD HH:MM:SS'." result = add_appointment(customer_name, customer_phone, appointment_datetime, service_type) if result["status"] == "success": return f"Success! Appointment for {customer_name} is scheduled. The appointment ID is {result['id']}." else: return f"Error: {result['message']}" @tool(args_schema=SearchAppointmentInput) def search_for_appointments(criteria: str, value: str) -> str: """Searches for existing appointments based on name, service type, or date.""" # <-- ADDED THIS DOCSTRING BACK results = search_appointments(criteria, value) if not results: return f"No confirmed appointments found for {criteria} = {value}." details = "\n".join([ f"- ID: {r['id']}, Name: {r['customer_name']}, DateTime: {r['appointment_datetime']}, Service: {r['service_type']}" for r in results]) return f"Found the following appointments:\n{details}" @tool def list_all_appointments() -> str: """Lists all confirmed appointments currently in the system. Use this when the user asks to see all appointments without specifying any criteria.""" results = get_all_appointments() if not results: return "There are currently no confirmed appointments in the system." details = "\n".join([ f"- ID: {r['id']}, Name: {r['customer_name']}, DateTime: {r['appointment_datetime']}, Service: {r['service_type']}" for r in results ]) return f"Here are all the confirmed appointments:\n{details}" @tool(args_schema=DeleteAppointmentInput) def delete_appointment_records(criteria: str, value: str) -> str: """Deletes appointment records by name, service type, or date.""" # <-- ADDED THIS DOCSTRING BACK count = delete_appointments(criteria, value) if count == 0: return f"No appointments found to delete for {criteria} = {value}." else: return f"Successfully deleted {count} appointment(s)." @tool(args_schema=UpdateAppointmentInput) def update_appointment_record(identifier_value: str, field_to_update: str, new_value: str) -> str: """Updates an appointment's information, like phone number or time.""" # <-- ADDED THIS DOCSTRING BACK count = update_appointment(identifier_value, field_to_update, new_value) if count == 0: return f"Error: Could not find or update an appointment for customer '{identifier_value}'." else: return f"Successfully updated {count} appointment(s) for {identifier_value}."