diff --git a/llm_demo/orchestrator/langchain_tools/tools.py b/llm_demo/orchestrator/langchain_tools/tools.py index c7589c3a..df533c08 100644 --- a/llm_demo/orchestrator/langchain_tools/tools.py +++ b/llm_demo/orchestrator/langchain_tools/tools.py @@ -21,10 +21,10 @@ import google.oauth2.id_token # type: ignore from google.auth import compute_engine # type: ignore from google.auth.transport.requests import Request # type: ignore -from langchain_core.tools import StructuredTool from pydantic import BaseModel, Field +from toolbox_langchain_sdk import ToolboxClient -BASE_URL = os.getenv("BASE_URL", default="http://127.0.0.1:8080") +BASE_URL = os.getenv("BASE_URL", default="http://127.0.0.1:5000") CREDENTIALS = None @@ -272,164 +272,168 @@ async def list_tickets(): # Tools for agent async def initialize_tools(client: aiohttp.ClientSession): - return [ - StructuredTool.from_function( - coroutine=generate_search_airports(client), - name="Search Airport", - description=""" - Use this tool to list all airports matching search criteria. - Takes at least one of country, city, name, or all and returns all matching airports. - The agent can decide to return the results directly to the user. - Input of this tool must be in JSON format and include all three inputs - country, city, name. - Example: - {{ - "country": "United States", - "city": "San Francisco", - "name": null - }} - Example: - {{ - "country": null, - "city": "Goroka", - "name": "Goroka" - }} - Example: - {{ - "country": "Mexico", - "city": null, - "name": null - }} - """, - args_schema=AirportSearchInput, - ), - StructuredTool.from_function( - coroutine=generate_search_flights_by_number(client), - name="Search Flights By Flight Number", - description=""" - Use this tool to get information for a specific flight. - Takes an airline code and flight number and returns info on the flight. - Do NOT use this tool with a flight id. Do NOT guess an airline code or flight number. - A airline code is a code for an airline service consisting of two-character - airline designator and followed by flight number, which is 1 to 4 digit number. - For example, if given CY 0123, the airline is "CY", and flight_number is "123". - Another example for this is DL 1234, the airline is "DL", and flight_number is "1234". - If the tool returns more than one option choose the date closes to today. - Example: - {{ - "airline": "CY", - "flight_number": "888", - }} - Example: - {{ - "airline": "DL", - "flight_number": "1234", - }} - """, - args_schema=FlightNumberInput, - ), - StructuredTool.from_function( - coroutine=generate_list_flights(client), - name="List Flights", - description=""" - Use this tool to list flights information matching search criteria. - Takes an arrival airport, a departure airport, or both, filters by date and returns all matching flights. - If 3-letter iata code is not provided for departure_airport or arrival_airport, use search airport tools to get iata code information. - Do NOT guess a date, ask user for date input if it is not given. Date must be in the following format: YYYY-MM-DD. - The agent can decide to return the results directly to the user. - Input of this tool must be in JSON format and include all three inputs - arrival_airport, departure_airport, and date. - Example: - {{ - "departure_airport": "SFO", - "arrival_airport": null, - "date": 2023-10-30" - }} - Example: - {{ - "departure_airport": "SFO", - "arrival_airport": "SEA", - "date": "2023-11-01" - }} - Example: - {{ - "departure_airport": null, - "arrival_airport": "SFO", - "date": "2023-01-01" - }} - """, - args_schema=ListFlights, - ), - StructuredTool.from_function( - coroutine=generate_search_amenities(client), - name="Search Amenities", - description=""" - Use this tool to search amenities by name or to recommended airport amenities at SFO. - If user provides flight info, use 'Search Flights by Flight Number' - first to get gate info and location. - Only recommend amenities that are returned by this query. - Find amenities close to the user by matching the terminal and then comparing - the gate numbers. Gate number iterate by letter and number, example A1 A2 A3 - B1 B2 B3 C1 C2 C3. Gate A3 is close to A2 and B1. - Input of this tool must be in JSON format and include one `query` input. - """, - args_schema=QueryInput, - ), - StructuredTool.from_function( - coroutine=generate_search_policies(client), - name="Search Policies", - description=""" - Use this tool to search for cymbal air passenger policy. - Policy that are listed is unchangeable. - You will not answer any questions outside of the policy given. - Policy includes information on ticket purchase and changes, baggage, check-in and boarding, special assistance, overbooking, flight delays and cancellations. - Input of this tool must be in JSON format and include one `query` input. - """, - args_schema=QueryInput, - ), - StructuredTool.from_function( - coroutine=generate_insert_ticket(client), - name="Insert Ticket", - description=""" - Use this tool to book a flight ticket for the user. - Example: - {{ - "airline": "AA", - "flight_number": "452", - "departure_airport": "LAX", - "arrival_airport": "SFO", - "departure_time": "2024-01-01 05:50:00", - "arrival_time": "2024-01-01 09:23:00" - }} - Example: - {{ - "airline": "UA", - "flight_number": "1532", - "departure_airport": "SFO", - "arrival_airport": "DEN", - "departure_time": "2024-01-08 05:50:00", - "arrival_time": "2024-01-08 09:23:00" - }} - Example: - {{ - "airline": "OO", - "flight_number": "6307", - "departure_airport": "SFO", - "arrival_airport": "MSP", - "departure_time": "2024-10-28 20:13:00", - "arrival_time": "2024-10-28 21:07:00" - }} - """, - args_schema=TicketInput, - ), - StructuredTool.from_function( - coroutine=generate_list_tickets(client), - name="List Tickets", - description=""" - Use this tool to list a user's flight tickets. - Takes no input and returns a list of current user's flight tickets. - Input is always empty JSON blob. Example: {{}} - """, - ), - ] + toolbox_client = ToolboxClient(BASE_URL, client) + tools = await toolbox_client.load_toolset(auth_headers={ "my_google_service": lambda: "YOUR_ID_TOKEN" }) + return tools + + # return [ + # StructuredTool.from_function( + # coroutine=generate_search_airports(client), + # name="Search Airport", + # description=""" + # Use this tool to list all airports matching search criteria. + # Takes at least one of country, city, name, or all and returns all matching airports. + # The agent can decide to return the results directly to the user. + # Input of this tool must be in JSON format and include all three inputs - country, city, name. + # Example: + # {{ + # "country": "United States", + # "city": "San Francisco", + # "name": null + # }} + # Example: + # {{ + # "country": null, + # "city": "Goroka", + # "name": "Goroka" + # }} + # Example: + # {{ + # "country": "Mexico", + # "city": null, + # "name": null + # }} + # """, + # args_schema=AirportSearchInput, + # ), + # StructuredTool.from_function( + # coroutine=generate_search_flights_by_number(client), + # name="Search Flights By Flight Number", + # description=""" + # Use this tool to get information for a specific flight. + # Takes an airline code and flight number and returns info on the flight. + # Do NOT use this tool with a flight id. Do NOT guess an airline code or flight number. + # A airline code is a code for an airline service consisting of two-character + # airline designator and followed by flight number, which is 1 to 4 digit number. + # For example, if given CY 0123, the airline is "CY", and flight_number is "123". + # Another example for this is DL 1234, the airline is "DL", and flight_number is "1234". + # If the tool returns more than one option choose the date closes to today. + # Example: + # {{ + # "airline": "CY", + # "flight_number": "888", + # }} + # Example: + # {{ + # "airline": "DL", + # "flight_number": "1234", + # }} + # """, + # args_schema=FlightNumberInput, + # ), + # StructuredTool.from_function( + # coroutine=generate_list_flights(client), + # name="List Flights", + # description=""" + # Use this tool to list flights information matching search criteria. + # Takes an arrival airport, a departure airport, or both, filters by date and returns all matching flights. + # If 3-letter iata code is not provided for departure_airport or arrival_airport, use search airport tools to get iata code information. + # Do NOT guess a date, ask user for date input if it is not given. Date must be in the following format: YYYY-MM-DD. + # The agent can decide to return the results directly to the user. + # Input of this tool must be in JSON format and include all three inputs - arrival_airport, departure_airport, and date. + # Example: + # {{ + # "departure_airport": "SFO", + # "arrival_airport": null, + # "date": 2023-10-30" + # }} + # Example: + # {{ + # "departure_airport": "SFO", + # "arrival_airport": "SEA", + # "date": "2023-11-01" + # }} + # Example: + # {{ + # "departure_airport": null, + # "arrival_airport": "SFO", + # "date": "2023-01-01" + # }} + # """, + # args_schema=ListFlights, + # ), + # StructuredTool.from_function( + # coroutine=generate_search_amenities(client), + # name="Search Amenities", + # description=""" + # Use this tool to search amenities by name or to recommended airport amenities at SFO. + # If user provides flight info, use 'Search Flights by Flight Number' + # first to get gate info and location. + # Only recommend amenities that are returned by this query. + # Find amenities close to the user by matching the terminal and then comparing + # the gate numbers. Gate number iterate by letter and number, example A1 A2 A3 + # B1 B2 B3 C1 C2 C3. Gate A3 is close to A2 and B1. + # Input of this tool must be in JSON format and include one `query` input. + # """, + # args_schema=QueryInput, + # ), + # StructuredTool.from_function( + # coroutine=generate_search_policies(client), + # name="Search Policies", + # description=""" + # Use this tool to search for cymbal air passenger policy. + # Policy that are listed is unchangeable. + # You will not answer any questions outside of the policy given. + # Policy includes information on ticket purchase and changes, baggage, check-in and boarding, special assistance, overbooking, flight delays and cancellations. + # Input of this tool must be in JSON format and include one `query` input. + # """, + # args_schema=QueryInput, + # ), + # StructuredTool.from_function( + # coroutine=generate_insert_ticket(client), + # name="Insert Ticket", + # description=""" + # Use this tool to book a flight ticket for the user. + # Example: + # {{ + # "airline": "AA", + # "flight_number": "452", + # "departure_airport": "LAX", + # "arrival_airport": "SFO", + # "departure_time": "2024-01-01 05:50:00", + # "arrival_time": "2024-01-01 09:23:00" + # }} + # Example: + # {{ + # "airline": "UA", + # "flight_number": "1532", + # "departure_airport": "SFO", + # "arrival_airport": "DEN", + # "departure_time": "2024-01-08 05:50:00", + # "arrival_time": "2024-01-08 09:23:00" + # }} + # Example: + # {{ + # "airline": "OO", + # "flight_number": "6307", + # "departure_airport": "SFO", + # "arrival_airport": "MSP", + # "departure_time": "2024-10-28 20:13:00", + # "arrival_time": "2024-10-28 21:07:00" + # }} + # """, + # args_schema=TicketInput, + # ), + # StructuredTool.from_function( + # coroutine=generate_list_tickets(client), + # name="List Tickets", + # description=""" + # Use this tool to list a user's flight tickets. + # Takes no input and returns a list of current user's flight tickets. + # Input is always empty JSON blob. Example: {{}} + # """, + # ), + # ] def get_confirmation_needing_tools(): diff --git a/llm_demo/tools.yaml b/llm_demo/tools.yaml new file mode 100644 index 00000000..1cf97826 --- /dev/null +++ b/llm_demo/tools.yaml @@ -0,0 +1,287 @@ +sources: + my-pg-instance: + kind: cloud-sql-postgres + project: retrieval-app-testing + region: us-central1 + instance: my-cloudsql-pg-instance + database: assistantdemo + user: postgres + password: postgres +authSources: + my_google_service: + kind: google + client_id: YOUR_CLIENT_ID +tools: + search_airports: + kind: postgres-sql + source: my-pg-instance + description: | + Use this tool to list all airports matching search criteria. + Takes at least one of country, city, name, or all and returns all matching airports. + The agent can decide to return the results directly to the user. + Input of this tool must be in JSON format and include all three inputs - country, city, name. + Example: + {{ + "country": "United States", + "city": "San Francisco", + "name": null + }} + Example: + {{ + "country": null, + "city": "Goroka", + "name": "Goroka" + }} + Example: + {{ + "country": "Mexico", + "city": null, + "name": null + }} + parameters: + - name: country + type: string + description: Country + - name: city + type: string + description: City + - name: name + type: string + description: Airport name + statement: | + SELECT * FROM airports + WHERE (CAST($1 AS TEXT) IS NULL OR country ILIKE $1) + AND (CAST($2 AS TEXT) IS NULL OR city ILIKE $2) + AND (CAST($3 AS TEXT) IS NULL OR name ILIKE '%' || $3 || '%') + LIMIT 10 + search_flights_by_number: + kind: postgres-sql + source: my-pg-instance + description: | + Use this tool to get information for a specific flight. + Takes an airline code and flight number and returns info on the flight. + Do NOT use this tool with a flight id. Do NOT guess an airline code or flight number. + A airline code is a code for an airline service consisting of two-character + airline designator and followed by flight number, which is 1 to 4 digit number. + For example, if given CY 0123, the airline is "CY", and flight_number is "123". + Another example for this is DL 1234, the airline is "DL", and flight_number is "1234". + If the tool returns more than one option choose the date closes to today. + Example: + {{ + "airline": "CY", + "flight_number": "888", + }} + Example: + {{ + "airline": "DL", + "flight_number": "1234", + }} + parameters: + - name: airline + type: string + description: Airline unique 2 letter identifier + - name: flight_number + type: string + description: 1 to 4 digit number + statement: | + SELECT * FROM flights + WHERE airline = $1 + AND flight_number = $2 + LIMIT 10 + list_flights: + kind: postgres-sql + source: my-pg-instance + description: | + Use this tool to list flights information matching search criteria. + Takes an arrival airport, a departure airport, or both, filters by date and returns all matching flights. + If 3-letter iata code is not provided for departure_airport or arrival_airport, use search airport tools to get iata code information. + Do NOT guess a date, ask user for date input if it is not given. Date must be in the following format: YYYY-MM-DD. + The agent can decide to return the results directly to the user. + Input of this tool must be in JSON format and include all three inputs - arrival_airport, departure_airport, and date. + Example: + {{ + "departure_airport": "SFO", + "arrival_airport": null, + "date": 2023-10-30" + }} + Example: + {{ + "departure_airport": "SFO", + "arrival_airport": "SEA", + "date": "2023-11-01" + }} + Example: + {{ + "departure_airport": null, + "arrival_airport": "SFO", + "date": "2023-01-01" + }} + parameters: + - name: departure_airport + type: string + description: Departure airport 3-letter code + - name: arrival_airport + type: string + description: Arrival airport 3-letter code + - name: date + type: string + description: Date of flight departure + statement: | + SELECT * FROM flights + WHERE (CAST($1 AS TEXT) IS NULL OR departure_airport ILIKE $1) + AND (CAST($2 AS TEXT) IS NULL OR arrival_airport ILIKE $2) + AND departure_time >= CAST($3 AS timestamp) + AND departure_time < CAST($3 AS timestamp) + interval '1 day' + LIMIT 10 + search_amenities: + kind: postgres-sql + source: my-pg-instance + description: | + Use this tool to search amenities by name or to recommended airport amenities at SFO. + If user provides flight info, use 'Search Flights by Flight Number' + first to get gate info and location. + Only recommend amenities that are returned by this query. + Find amenities close to the user by matching the terminal and then comparing + the gate numbers. Gate number iterate by letter and number, example A1 A2 A3 + B1 B2 B3 C1 C2 C3. Gate A3 is close to A2 and B1. + Input of this tool must be in JSON format and include one `query` input. + parameters: + - name: query + type: string + description: Search query + statement: | + SELECT name, description, location, terminal, category, hour + FROM amenities + WHERE (embedding <=> $1) < 0.5 + ORDER BY (embedding <=> $1) + LIMIT 5 + search_policies: + kind: postgres-sql + source: my-pg-instance + description: | + Use this tool to search for cymbal air passenger policy. + Policy that are listed is unchangeable. + You will not answer any questions outside of the policy given. + Policy includes information on ticket purchase and changes, baggage, check-in and boarding, special assistance, overbooking, flight delays and cancellations. + Input of this tool must be in JSON format and include one `query` input. + parameters: + - name: query + type: string + description: Search query + statement: | + SELECT content + FROM policies + WHERE (embedding <=> $1) < 0.5 + ORDER BY (embedding <=> $1) + LIMIT 5 + insert_ticket: + kind: postgres-sql + source: my-pg-instance + description: | + Use this tool to book a flight ticket for the user. + Example: + {{ + "airline": "AA", + "flight_number": "452", + "departure_airport": "LAX", + "arrival_airport": "SFO", + "departure_time": "2024-01-01 05:50:00", + "arrival_time": "2024-01-01 09:23:00" + }} + Example: + {{ + "airline": "UA", + "flight_number": "1532", + "departure_airport": "SFO", + "arrival_airport": "DEN", + "departure_time": "2024-01-08 05:50:00", + "arrival_time": "2024-01-08 09:23:00" + }} + Example: + {{ + "airline": "OO", + "flight_number": "6307", + "departure_airport": "SFO", + "arrival_airport": "MSP", + "departure_time": "2024-10-28 20:13:00", + "arrival_time": "2024-10-28 21:07:00" + }} + parameters: + - name: user_id + type: string + authSources: + - name: my_google_service + field: sub + - name: user_name + type: string + authSources: + - name: my_google_service + field: sub + - name: user_email + type: string + authSources: + - name: my_google_service + field: sub + - name: airline + type: string + description: Airline unique 2 letter identifier + - name: flight_number + type: string + description: 1 to 4 digit number + - name: departure_airport + type: string + description: Departure airport 3-letter code + - name: departure_time + type: string + description: Flight departure datetime + - name: arrival_airport + type: string + description: Arrival airport 3-letter code + - name: arrival_time + type: string + description: Flight arrival datetime + statement: | + INSERT INTO tickets ( + user_id, + user_name, + user_email, + airline, + flight_number, + departure_airport, + departure_time, + arrival_airport, + arrival_time + ) VALUES ( + $1, + $2, + $3, + $4, + $5, + $6, + $7, + $8, + $9 + ); + list_tickets: + kind: postgres-sql + source: my-pg-instance + description: | + Use this tool to list a user's flight tickets. + Takes no input and returns a list of current user's flight tickets. + Input is always empty JSON blob. Example: {{}} + parameters: + - name: user_id + type: string + description: User id + statement: | + SELECT user_name, airline, flight_number, departure_airport, arrival_airport, departure_time, arrival_time FROM tickets + WHERE user_id = $1 +toolsets: + airport_toolset: + - search_airports + - search_flights_by_number + - list_flights + - search_amenities + - search_policies + - insert_ticket + - list_tickets