Fix pre-commit issues

Author: brightsparc

examples/pydantic_ai_examples/flight_booking_grok.py

@@ -4,9 +4,10 @@
 """
 
 import datetime
+import os
 from dataclasses import dataclass
 from typing import Literal
-import os
+
 import logfire
 from pydantic import BaseModel, Field
 from rich.prompt import Prompt
@@ -22,22 +23,22 @@
 logfire.instrument_httpx()
 
 # Configure for xAI API
-xai_api_key = os.getenv("XAI_API_KEY")
+xai_api_key = os.getenv('XAI_API_KEY')
 if not xai_api_key:
-    raise ValueError("XAI_API_KEY environment variable is required")
+    raise ValueError('XAI_API_KEY environment variable is required')
 
 
 # Create the model using the new GrokModelpwd
-model = GrokModel("grok-4-fast-non-reasoning", api_key=xai_api_key)
+model = GrokModel('grok-4-fast-non-reasoning', api_key=xai_api_key)
 
 
 class FlightDetails(BaseModel):
     """Details of the most suitable flight."""
 
     flight_number: str
     price: int
-    origin: str = Field(description="Three-letter airport code")
-    destination: str = Field(description="Three-letter airport code")
+    origin: str = Field(description='Three-letter airport code')
+    destination: str = Field(description='Three-letter airport code')
     date: datetime.date
 
 
@@ -58,15 +59,17 @@ class Deps:
     model=model,
     output_type=FlightDetails | NoFlightFound,  # type: ignore
     retries=4,
-    system_prompt=("Your job is to find the cheapest flight for the user on the given date. "),
+    system_prompt=(
+        'Your job is to find the cheapest flight for the user on the given date. '
+    ),
 )
 
 
 # This agent is responsible for extracting flight details from web page text.
 extraction_agent = Agent(
     model=model,
     output_type=list[FlightDetails],
-    system_prompt="Extract all the flight details from the given text.",
+    system_prompt='Extract all the flight details from the given text.',
 )
 
 
@@ -75,7 +78,7 @@ async def extract_flights(ctx: RunContext[Deps]) -> list[FlightDetails]:
     """Get details of all flights."""
     # we pass the usage to the search agent so requests within this agent are counted
     result = await extraction_agent.run(ctx.deps.web_page_text, usage=ctx.usage)
-    logfire.info("found {flight_count} flights", flight_count=len(result.output))
+    logfire.info('found {flight_count} flights', flight_count=len(result.output))
     return result.output
 
 
@@ -89,23 +92,25 @@ async def validate_output(
 
     errors: list[str] = []
     if output.origin != ctx.deps.req_origin:
-        errors.append(f"Flight should have origin {ctx.deps.req_origin}, not {output.origin}")
+        errors.append(
+            f'Flight should have origin {ctx.deps.req_origin}, not {output.origin}'
+        )
     if output.destination != ctx.deps.req_destination:
         errors.append(
-            f"Flight should have destination {ctx.deps.req_destination}, not {output.destination}"
+            f'Flight should have destination {ctx.deps.req_destination}, not {output.destination}'
         )
     if output.date != ctx.deps.req_date:
-        errors.append(f"Flight should be on {ctx.deps.req_date}, not {output.date}")
+        errors.append(f'Flight should be on {ctx.deps.req_date}, not {output.date}')
 
     if errors:
-        raise ModelRetry("\n".join(errors))
+        raise ModelRetry('\n'.join(errors))
     else:
         return output
 
 
 class SeatPreference(BaseModel):
     row: int = Field(ge=1, le=30)
-    seat: Literal["A", "B", "C", "D", "E", "F"]
+    seat: Literal['A', 'B', 'C', 'D', 'E', 'F']
 
 
 class Failed(BaseModel):
@@ -118,9 +123,9 @@ class Failed(BaseModel):
     output_type=SeatPreference | Failed,
     system_prompt=(
         "Extract the user's seat preference. "
-        "Seats A and F are window seats. "
-        "Row 1 is the front row and has extra leg room. "
-        "Rows 14, and 20 also have extra leg room. "
+        'Seats A and F are window seats. '
+        'Row 1 is the front row and has extra leg room. '
+        'Rows 14, and 20 also have extra leg room. '
     ),
 )
 
@@ -184,46 +189,46 @@ class Failed(BaseModel):
 async def main():
     deps = Deps(
         web_page_text=flights_web_page,
-        req_origin="SFO",
-        req_destination="ANC",
+        req_origin='SFO',
+        req_destination='ANC',
         req_date=datetime.date(2025, 1, 10),
     )
     message_history: list[ModelMessage] | None = None
     usage: RunUsage = RunUsage()
     # run the agent until a satisfactory flight is found
     while True:
         result = await search_agent.run(
-            f"Find me a flight from {deps.req_origin} to {deps.req_destination} on {deps.req_date}",
+            f'Find me a flight from {deps.req_origin} to {deps.req_destination} on {deps.req_date}',
             deps=deps,
             usage=usage,
             message_history=message_history,
             usage_limits=usage_limits,
         )
         if isinstance(result.output, NoFlightFound):
-            print("No flight found")
+            print('No flight found')
             break
         else:
             flight = result.output
-            print(f"Flight found: {flight}")
+            print(f'Flight found: {flight}')
             answer = Prompt.ask(
-                "Do you want to buy this flight, or keep searching? (buy/*search)",
-                choices=["buy", "search", ""],
+                'Do you want to buy this flight, or keep searching? (buy/*search)',
+                choices=['buy', 'search', ''],
                 show_choices=False,
             )
-            if answer == "buy":
+            if answer == 'buy':
                 seat = await find_seat(usage)
                 await buy_tickets(flight, seat)
                 break
             else:
                 message_history = result.all_messages(
-                    output_tool_return_content="Please suggest another flight"
+                    output_tool_return_content='Please suggest another flight'
                 )
 
 
 async def find_seat(usage: RunUsage) -> SeatPreference:
     message_history: list[ModelMessage] | None = None
     while True:
-        answer = Prompt.ask("What seat would you like?")
+        answer = Prompt.ask('What seat would you like?')
 
         result = await seat_preference_agent.run(
             answer,
@@ -234,15 +239,15 @@ async def find_seat(usage: RunUsage) -> SeatPreference:
         if isinstance(result.output, SeatPreference):
             return result.output
         else:
-            print("Could not understand seat preference. Please try again.")
+            print('Could not understand seat preference. Please try again.')
             message_history = result.all_messages()
 
 
 async def buy_tickets(flight_details: FlightDetails, seat: SeatPreference):
-    print(f"Purchasing flight {flight_details=!r} {seat=!r}...")
+    print(f'Purchasing flight {flight_details=!r} {seat=!r}...')
 
 
-if __name__ == "__main__":
+if __name__ == '__main__':
     import asyncio
 
     asyncio.run(main())