feat: add new 5 tools to the agent

app.py

@@ -10,15 +10,10 @@ from rich.console import Console
 from langchain_google_genai import ChatGoogleGenerativeAI
 from gradio import ChatMessage
 import textwrap
-from tools import GetTelemetry
+from tools import *
 load_dotenv()
 os.environ['LANGCHAIN_PROJECT'] = 'gradio-test'
 
-console = Console(style="chartreuse1 on grey7")
-
-# * Initialize database
-db = SQLDatabase.from_uri("sqlite:///db/Bahrain_2023_Q.db")
-
 # * Initialize LLM
 llm = ChatGoogleGenerativeAI(
     model="gemini-1.5-flash",
@@ -32,8 +27,17 @@ llm = ChatGoogleGenerativeAI(
 toolkit = SQLDatabaseToolkit(db=db, llm=llm)
 tools = toolkit.get_tools()
 
+get_driver_performance_tool = GetDriverPerformance()
+get_event_performance_tool = GetEventPerformance()
 get_telemetry_tool = GetTelemetry()
+get_tyre_performance_tool = GetTyrePerformance()
+get_weather_impact_tool = GetWeatherImpact()
+
+tools.append(get_driver_performance_tool)
+tools.append(get_event_performance_tool)
 tools.append(get_telemetry_tool)
+tools.append(get_tyre_performance_tool)
+tools.append(get_weather_impact_tool)
 
 # * Initialize agent
 agent_prompt = open("agent_prompt.txt", "r")

tools/__init__.py

@@ -1,82 +1,21 @@
-from pydantic import BaseModel, Field
-from typing import Type
-from langchain_core.tools import BaseTool
 from langchain_community.utilities import SQLDatabase
 from rich.console import Console
+from .driver_performance import GetDriverPerformance
+from .event_performance import GetEventPerformance
+from .telemetry_analysis import GetTelemetry
+from .tyre_performance import GetTyrePerformance
+from .weather_impact import GetWeatherImpact
 
 console = Console(style="chartreuse1 on grey7")
 
 db = SQLDatabase.from_uri("sqlite:///db/Bahrain_2023_Q.db")
 
-
-class GetTelemetryAndWeatherInput(BaseModel):
-    """Input for the get_telemetry_and_weather tool"""
-    driver_name: str = Field(
-        description="Name of the driver to analyze (e.g., 'VER', 'HAM', 'LEC', etc.)")
-    lap_number: int = Field(description="Lap number to analyze")
-
-
-class GetTelemetryAndWeatherOutput(BaseModel):
-    """Output for the get_telemetry_and_weather tool"""
-    lap_id: int = Field(description="Lap ID")
-    lap_number: int = Field(description="Lap number")
-    lap_time_in_seconds: float | None = Field(
-        description="Lap time in seconds")
-    avg_speed: float = Field(description="Average speed in km/h")
-    max_speed: float = Field(description="Maximum speed in km/h")
-    avg_RPM: float = Field(description="Average RPM")
-    max_RPM: float = Field(description="Maximum RPM")
-    avg_throttle: float = Field(description="Average throttle")
-    brake_percentage: float = Field(description="Brake percentage")
-    drs_usage_percentage: float = Field(description="Drs usage percentage")
-    off_track_percentage: float = Field(description="Off track percentage")
-    avg_air_temp: float | None = Field(
-        description="Average air temperature in celsius")
-    avg_track_temp: float | None = Field(
-        description="Average track temperature in celsius")
-    avg_wind_speed: float | None = Field(
-        description="Average wind speed in meters per second")
-
-
-class GetTelemetry(BaseTool):
-    name: str = "get_telemetry"
-    description: str = "useful for when you need to answer questions about telemetry for a given driver and lap"
-    args_schema: Type[BaseModel] = GetTelemetryAndWeatherInput
-
-    def _run(
-        self, driver_name: str, lap_number: int
-    ) -> GetTelemetryAndWeatherOutput:
-        # """Use the tool."""
-        sql_file = open("tools/telemetry_and_weather_query.sql", "r")
-        sql_query = sql_file.read()
-        sql_file.close()
-        console.print("getting telemetry")
-        response = db.run(sql_query, parameters={
-            "driver_name": driver_name,
-            "lap_number": lap_number})
-
-        if not isinstance(response, str):
-            response = str(response)
-
-        clean_response = response.strip('[]()').split(',')
-        # Convert to appropriate types and create dictionary
-        return GetTelemetryAndWeatherOutput(
-            lap_id=int(float(clean_response[0])),
-            lap_number=int(float(clean_response[1])),
-            lap_time_in_seconds=float(
-                clean_response[2]) if clean_response[2].strip() != 'None' else None,
-            avg_speed=float(clean_response[3]),
-            max_speed=float(clean_response[4]),
-            avg_RPM=float(clean_response[5]),
-            max_RPM=float(clean_response[6]),
-            avg_throttle=float(clean_response[7]),
-            brake_percentage=float(clean_response[8]),
-            drs_usage_percentage=float(clean_response[9]),
-            off_track_percentage=float(clean_response[10]),
-            avg_air_temp=float(
-                clean_response[11]) if clean_response[11].strip() != 'None' else None,
-            avg_track_temp=float(
-                clean_response[12]) if clean_response[12].strip() != 'None' else None,
-            avg_wind_speed=float(
-                clean_response[13]) if clean_response[13].strip() != 'None' else None
-        )
+__all__ = [
+    "GetDriverPerformance",
+    "GetEventPerformance",
+    "GetTelemetry",
+    "GetTyrePerformance",
+    "GetWeatherImpact",
+    "console",
+    "db"
+]

tools/driver_performance.py

@@ -0,0 +1,88 @@
+from pydantic import BaseModel, Field
+from langchain_core.tools import BaseTool
+from . import db, console
+
+
+class GetDriverPerformanceOutput(BaseModel):
+    """Output for the get_driver_performance tool"""
+    driver_name: str = Field(description="Name of the driver")
+    event_name: str = Field(description="Name of the event")
+    session_type: str = Field(
+        description="Type of session (Practice, Qualifying, Race)")
+    track_name: str = Field(description="Name of the track")
+    total_laps: int = Field(description="Total number of laps completed")
+    avg_lap_time: float | None = Field(
+        description="Average lap time in seconds")
+    best_lap_time: float | None = Field(description="Best lap time in seconds")
+    avg_sector1_time: float | None = Field(
+        description="Average sector 1 time in seconds")
+    avg_sector2_time: float | None = Field(
+        description="Average sector 2 time in seconds")
+    avg_sector3_time: float | None = Field(
+        description="Average sector 3 time in seconds")
+    avg_finish_line_speed: float | None = Field(
+        description="Average finish line speed in km/h")
+    personal_best_laps: int = Field(description="Number of personal best laps")
+    avg_air_temp: float | None = Field(
+        description="Average air temperature in celsius")
+    avg_track_temp: float | None = Field(
+        description="Average track temperature in celsius")
+    rain_percentage: float = Field(
+        description="Percentage of time it rained during the session")
+
+
+class GetDriverPerformance(BaseTool):
+    name: str = "get_driver_performance"
+    description: str = "useful for when you need to analyze driver performance statistics across different sessions and events"
+
+    def _run(self) -> list[GetDriverPerformanceOutput]:
+        """Use the tool."""
+        sql_file = open("tools/sql/driver_performance.query.sql", "r")
+        sql_query = sql_file.read()
+        sql_file.close()
+
+        console.print("getting driver performance data")
+        response = db.run(sql_query)
+
+        if not isinstance(response, str):
+            response = str(response)
+
+        # Remove the outer brackets and split by rows
+        rows = response.strip('[]').split('), (')
+
+        results = []
+        for row in rows:
+            # Clean up the row string and split by columns
+            clean_row = row.strip('()').split(',')
+
+            # Convert to appropriate types and create output object
+            driver_data = GetDriverPerformanceOutput(
+                driver_name=clean_row[0].strip("'"),
+                event_name=clean_row[1].strip("'"),
+                session_type=clean_row[2].strip("'"),
+                track_name=clean_row[3].strip("'"),
+                total_laps=int(float(clean_row[4])),
+                avg_lap_time=float(
+                    clean_row[5]) if clean_row[5].strip() != 'None' else None,
+                best_lap_time=float(
+                    clean_row[6]) if clean_row[6].strip() != 'None' else None,
+                avg_sector1_time=float(
+                    clean_row[7]) if clean_row[7].strip() != 'None' else None,
+                avg_sector2_time=float(
+                    clean_row[8]) if clean_row[8].strip() != 'None' else None,
+                avg_sector3_time=float(
+                    clean_row[9]) if clean_row[9].strip() != 'None' else None,
+                avg_finish_line_speed=float(
+                    clean_row[10]) if clean_row[10].strip() != 'None' else None,
+                personal_best_laps=int(
+                    float(clean_row[11])) if clean_row[11].strip() != 'None' else 0,
+                avg_air_temp=float(
+                    clean_row[12]) if clean_row[12].strip() != 'None' else None,
+                avg_track_temp=float(
+                    clean_row[13]) if clean_row[13].strip() != 'None' else None,
+                rain_percentage=float(
+                    clean_row[14]) if clean_row[14].strip() != 'None' else 0.0
+            )
+            results.append(driver_data)
+
+        return results

tools/event_performance.py

@@ -0,0 +1,74 @@
+from pydantic import BaseModel, Field
+from typing import Type
+from langchain_core.tools import BaseTool
+from . import db, console
+
+
+class GetEventPerformanceOutput(BaseModel):
+    """Output for the get_event_performance tool"""
+    event_name: str = Field(description="Name of the event")
+    country: str = Field(description="Country where the event took place")
+    location: str = Field(description="Specific location of the event")
+    session_type: str = Field(
+        description="Type of session (Practice, Qualifying, Race)")
+    driver_count: int = Field(
+        description="Number of drivers that participated")
+    avg_lap_time: float = Field(description="Average lap time in seconds")
+    best_lap_time: float = Field(description="Best lap time in seconds")
+    max_finish_line_speed: float = Field(
+        description="Maximum speed at finish line in km/h")
+    avg_air_temp: float | None = Field(
+        description="Average air temperature in celsius")
+    avg_track_temp: float | None = Field(
+        description="Average track temperature in celsius")
+    rain_percentage: float = Field(
+        description="Percentage of time it rained during the session")
+
+
+class GetEventPerformance(BaseTool):
+    name: str = "get_event_performance"
+    description: str = "useful for when you need to get performance statistics for Formula 1 events"
+
+    def _run(self) -> list[GetEventPerformanceOutput]:
+        """Use the tool."""
+        sql_file = open("tools/sql/event_performance.query.sql", "r")
+        sql_query = sql_file.read()
+        sql_file.close()
+
+        console.print("getting event performance data")
+        response = db.run(sql_query)
+
+        if not isinstance(response, str):
+            response = str(response)
+
+        # Remove the outer brackets and split by rows
+        rows = response.strip('[]').split('), (')
+
+        results = []
+        for row in rows:
+            # Clean up the row string and split by columns
+            clean_row = row.strip('()').split(',')
+
+            # Convert to appropriate types and create output object
+            event_data = GetEventPerformanceOutput(
+                event_name=clean_row[0].strip("'"),
+                country=clean_row[1].strip("'"),
+                location=clean_row[2].strip("'"),
+                session_type=clean_row[3].strip("'"),
+                driver_count=int(float(clean_row[4])),
+                avg_lap_time=float(
+                    clean_row[5]) if clean_row[5].strip() != 'None' else 0.0,
+                best_lap_time=float(
+                    clean_row[6]) if clean_row[6].strip() != 'None' else 0.0,
+                max_finish_line_speed=float(
+                    clean_row[7]) if clean_row[7].strip() != 'None' else 0.0,
+                avg_air_temp=float(
+                    clean_row[8]) if clean_row[8].strip() != 'None' else None,
+                avg_track_temp=float(
+                    clean_row[9]) if clean_row[9].strip() != 'None' else None,
+                rain_percentage=float(
+                    clean_row[10]) if clean_row[10].strip() != 'None' else 0.0
+            )
+            results.append(event_data)
+
+        return results

tools/telemetry_analysis.py

@@ -0,0 +1,77 @@
+from pydantic import BaseModel, Field
+from typing import Type
+from langchain_core.tools import BaseTool
+from . import db, console
+
+
+class GetTelemetryAndWeatherInput(BaseModel):
+    """Input for the get_telemetry_and_weather tool"""
+    driver_name: str = Field(
+        description="Name of the driver to analyze (e.g., 'VER', 'HAM', 'LEC', etc.)")
+    lap_number: int = Field(description="Lap number to analyze")
+
+
+class GetTelemetryAndWeatherOutput(BaseModel):
+    """Output for the get_telemetry_and_weather tool"""
+    lap_id: int = Field(description="Lap ID")
+    lap_number: int = Field(description="Lap number")
+    lap_time_in_seconds: float | None = Field(
+        description="Lap time in seconds")
+    avg_speed: float = Field(description="Average speed in km/h")
+    max_speed: float = Field(description="Maximum speed in km/h")
+    avg_RPM: float = Field(description="Average RPM")
+    max_RPM: float = Field(description="Maximum RPM")
+    avg_throttle: float = Field(description="Average throttle")
+    brake_percentage: float = Field(description="Brake percentage")
+    drs_usage_percentage: float = Field(description="Drs usage percentage")
+    off_track_percentage: float = Field(description="Off track percentage")
+    avg_air_temp: float | None = Field(
+        description="Average air temperature in celsius")
+    avg_track_temp: float | None = Field(
+        description="Average track temperature in celsius")
+    avg_wind_speed: float | None = Field(
+        description="Average wind speed in meters per second")
+
+
+class GetTelemetry(BaseTool):
+    name: str = "get_telemetry"
+    description: str = "useful for when you need to answer questions about telemetry for a given driver and lap"
+    args_schema: Type[BaseModel] = GetTelemetryAndWeatherInput
+
+    def _run(
+        self, driver_name: str, lap_number: int
+    ) -> GetTelemetryAndWeatherOutput:
+        # """Use the tool."""
+        sql_file = open("tools/sql/telemetry_analysis.query.sql", "r")
+        sql_query = sql_file.read()
+        sql_file.close()
+        console.print("getting telemetry")
+        response = db.run(sql_query, parameters={
+            "driver_name": driver_name,
+            "lap_number": lap_number})
+
+        if not isinstance(response, str):
+            response = str(response)
+
+        clean_response = response.strip('[]()').split(',')
+        # Convert to appropriate types and create dictionary
+        return GetTelemetryAndWeatherOutput(
+            lap_id=int(float(clean_response[0])),
+            lap_number=int(float(clean_response[1])),
+            lap_time_in_seconds=float(
+                clean_response[2]) if clean_response[2].strip() != 'None' else None,
+            avg_speed=float(clean_response[3]),
+            max_speed=float(clean_response[4]),
+            avg_RPM=float(clean_response[5]),
+            max_RPM=float(clean_response[6]),
+            avg_throttle=float(clean_response[7]),
+            brake_percentage=float(clean_response[8]),
+            drs_usage_percentage=float(clean_response[9]),
+            off_track_percentage=float(clean_response[10]),
+            avg_air_temp=float(
+                clean_response[11]) if clean_response[11].strip() != 'None' else None,
+            avg_track_temp=float(
+                clean_response[12]) if clean_response[12].strip() != 'None' else None,
+            avg_wind_speed=float(
+                clean_response[13]) if clean_response[13].strip() != 'None' else None
+        )

tools/tyre_performance.py

@@ -0,0 +1,82 @@
+from pydantic import BaseModel, Field
+from typing import Type
+from langchain_core.tools import BaseTool
+from . import db, console
+
+
+class GetTyrePerformanceInput(BaseModel):
+    """Input for the get_tyre_performance tool"""
+    driver_name: str = Field(description="Name of the driver to analyze")
+
+
+class GetTyrePerformanceOutput(BaseModel):
+    """Output for the get_tyre_performance tool"""
+    driver_name: str = Field(description="Name of the driver")
+    lap_number: int = Field(description="Lap number")
+    tyre_compound: str = Field(description="Type of tyre compound used")
+    avg_tyre_life: float | None = Field(
+        description="Average tyre life in laps")
+    avg_lap_time: float | None = Field(
+        description="Average lap time in seconds")
+    avg_top_speed: float | None = Field(
+        description="Average top speed in longest straight in km/h")
+    fresh_tyre_laps: int = Field(
+        description="Number of laps done with fresh tyres")
+    used_tyre_laps: int = Field(
+        description="Number of laps done with used tyres")
+    avg_track_temp: float | None = Field(
+        description="Average track temperature in celsius")
+    avg_air_temp: float | None = Field(
+        description="Average air temperature in celsius")
+
+
+class GetTyrePerformance(BaseTool):
+    name: str = "get_tyre_performance"
+    description: str = "useful for when you need to analyze tyre performance and degradation for a specific driver across all their laps"
+    args_schema: Type[BaseModel] = GetTyrePerformanceInput
+
+    def _run(self, driver_name: str) -> list[GetTyrePerformanceOutput]:
+        """Use the tool."""
+        sql_file = open("tools/sql/tyre_performance.query.sql", "r")
+        sql_query = sql_file.read()
+        sql_file.close()
+
+        console.print("getting tyre performance data")
+        response = db.run(sql_query, parameters={
+            "driver_name": driver_name
+        })
+
+        if not isinstance(response, str):
+            response = str(response)
+
+        # Remove the outer brackets and split by rows
+        rows = response.strip('[]').split('), (')
+
+        results = []
+        for row in rows:
+            # Clean up the row string and split by columns
+            clean_row = row.strip('()').split(',')
+
+            # Convert to appropriate types and create output object
+            tyre_data = GetTyrePerformanceOutput(
+                driver_name=clean_row[0].strip("'"),
+                lap_number=int(float(clean_row[1])),
+                tyre_compound=clean_row[2].strip("'"),
+                avg_tyre_life=float(
+                    clean_row[3]) if clean_row[3].strip() != 'None' else None,
+                avg_lap_time=float(
+                    clean_row[4]) if clean_row[4].strip() != 'None' else None,
+                avg_top_speed=float(
+                    clean_row[5]) if clean_row[5].strip() != 'None' else None,
+                fresh_tyre_laps=int(
+                    float(clean_row[6])) if clean_row[6].strip() != 'None' else 0,
+                used_tyre_laps=int(
+                    float(clean_row[7])) if clean_row[7].strip() != 'None' else 0,
+                avg_track_temp=float(
+                    clean_row[8]) if clean_row[8].strip() != 'None' else None,
+                avg_air_temp=float(
+                    clean_row[9]) if clean_row[9].strip() != 'None' else None
+            )
+            results.append(tyre_data)
+
+        return results

tools/weather_impact.py

@@ -0,0 +1,76 @@
+from pydantic import BaseModel, Field
+from typing import Type
+from langchain_core.tools import BaseTool
+from rich.console import Console
+from . import db
+
+console = Console(style="chartreuse1 on grey7")
+
+
+class GetWeatherImpactInput(BaseModel):
+    """Input for the get_weather_impact tool"""
+    # Note: This tool doesn't require input parameters as it returns weather impact data
+    pass
+
+
+class GetWeatherImpactOutput(BaseModel):
+    """Output for the get_weather_impact tool"""
+    event_name: str = Field(description="Name of the event")
+    session_type: str = Field(
+        description="Type of session (Practice, Qualifying, Race)")
+    track_name: str = Field(description="Name of the track")
+    avg_air_temp: float | None = Field(
+        description="Average air temperature in celsius")
+    avg_track_temp: float | None = Field(
+        description="Average track temperature in celsius")
+    avg_humidity: float | None = Field(
+        description="Average relative humidity percentage")
+    avg_wind_speed: float | None = Field(
+        description="Average wind speed in meters per second")
+    rain_percentage: float = Field(
+        description="Percentage of time it rained during the session")
+    avg_lap_time: float | None = Field(
+        description="Average lap time in seconds")
+    best_lap_time: float | None = Field(description="Best lap time in seconds")
+
+
+class GetWeatherImpact(BaseTool):
+    name: str = "get_weather_impact"
+    description: str = "useful for when you need to analyze how weather conditions impact Formula 1 session performance"
+    args_schema: Type[BaseModel] = GetWeatherImpactInput
+
+    def _run(self) -> GetWeatherImpactOutput:
+        """Use the tool."""
+        sql_file = open("tools/sql/weather_impact.query.sql", "r")
+        sql_query = sql_file.read()
+        sql_file.close()
+
+        console.print("getting weather impact data")
+        response = db.run(sql_query)
+
+        if not isinstance(response, str):
+            response = str(response)
+
+        # Clean up the single row response
+        clean_row = response.strip('[]()').split(',')
+
+        # Convert to appropriate types and create output object
+        return GetWeatherImpactOutput(
+            event_name=clean_row[0].strip("'"),
+            session_type=clean_row[1].strip("'"),
+            track_name=clean_row[2].strip("'"),
+            avg_air_temp=float(
+                clean_row[3]) if clean_row[3].strip() != 'None' else None,
+            avg_track_temp=float(
+                clean_row[4]) if clean_row[4].strip() != 'None' else None,
+            avg_humidity=float(
+                clean_row[5]) if clean_row[5].strip() != 'None' else None,
+            avg_wind_speed=float(
+                clean_row[6]) if clean_row[6].strip() != 'None' else None,
+            rain_percentage=float(
+                clean_row[7]) if clean_row[7].strip() != 'None' else 0.0,
+            avg_lap_time=float(
+                clean_row[8]) if clean_row[8].strip() != 'None' else None,
+            best_lap_time=float(
+                clean_row[9]) if clean_row[9].strip() != 'None' else None
+        )