examples: schedule-extract.py update

Author: pchalasani

examples/basic/schedule-extract.py

@@ -4,21 +4,25 @@
 Enter vague, unstructured info like:
 
 M-F 8-3pm at home or Tue/Wed 9-1030am at daycare
+
+Run like this --
+
+```bash
+uv run examples/basic/schedule-extract.py -m
 """
 
 import langroid as lr
 import langroid.language_models as lm
-from enum import Enum
 from langroid.agent.tools.orchestration import FinalResultTool
-from typing import List, Dict, Tuple
+from typing import List, Dict, Literal, Tuple
 from langroid.pydantic_v1 import BaseModel, Field
 from rich.prompt import Prompt
 from fire import Fire
 
 
 class Slot(BaseModel):
-    start_time: float = Field(..., description="start time of the slot, e.g. 11:30AM")
-    duration: float = Field(..., description="duration of the slot in MINUTES")
+    start_time: str = Field(..., description="start time of the slot, e.g. 11:30AM")
+    end_time: str = Field(..., description="end time of the slot, e.g. 12:30PM")
     location: str = Field(..., description="location of the slot or UNKNOWN")
 
 
@@ -30,28 +34,19 @@ class DaySchedule(BaseModel):
     slots: List[Slot] = Field(..., description="List of time slots for the day")
 
 
-class Weekday(int, Enum):
-    """
-    A class to represent a weekday.
-    """
-
-    MON = 0
-    TUE = 1
-    WED = 2
-    THU = 3
-    FRI = 4
+Weekday = Literal["Mon", "Tue", "Wed", "Thu", "Fri"]
 
 
 class Availability(BaseModel):
     """
     A class to represent schedule information.
     """
 
-    week_availability: Dict[int, DaySchedule] = Field(
+    week_availability: Dict[Weekday, DaySchedule] = Field(
         ...,
         description="""
         Dictionary mapping weekday to DaySchedule,
-        where 0 = Monday, 1 = Tuesday, ... 4 = Friday
+        where weekday is one of "Mon", "Tue", "Wed", "Thu", "Fri"
         """,
     )
 
@@ -77,17 +72,27 @@ def examples(cls) -> List["lr.ToolMessage" | Tuple[str, "lr.ToolMessage"]]:
                 cls(
                     availabilities=Availability(
                         week_availability={
-                            Weekday.MON: DaySchedule(
+                            "Mon": DaySchedule(
                                 slots=[
-                                    Slot(start_time=10, duration=360, location="home"),
                                     Slot(
-                                        start_time=15, duration=60, location="daycare"
+                                        start_time="10:00",
+                                        end_time="16:00",
+                                        location="home",
+                                    ),
+                                    Slot(
+                                        start_time="15:00",
+                                        end_time="16:00",
+                                        location="daycare",
                                     ),
                                 ]
                             ),
-                            Weekday.WED: DaySchedule(
+                            "Wed": DaySchedule(
                                 slots=[
-                                    Slot(start_time=10, duration=360, location="home")
+                                    Slot(
+                                        start_time="10:00",
+                                        end_time="16:00",
+                                        location="home",
+                                    )
                                 ]
                             ),
                         }
@@ -110,7 +115,7 @@ def handle(self) -> str:
 
 def make_schedule_task(model: str = ""):
     llm_config = lm.OpenAIGPTConfig(
-        chat_model=model or lm.GeminiModel.GEMINI_2_FLASH_LITE,
+        chat_model=model or lm.OpenAIChatModel.GPT4o_MINI,
     )
     agent = lr.ChatAgent(
         lr.ChatAgentConfig(