Spaces:

prats010
/

kronector

Runtime error

from agents.data_agent import data_agent

# Query with natural language
result = data_agent("Will Verstappen win Monaco 2023?")

# Access results
query = result["query"]                    # Original query
intent = result["intent"]                  # Parsed: season, round, driver
rows = result["rows"]                      # List of dicts
df = result["dataframe"]                   # pandas DataFrame

print(f"Found {len(rows)} row(s)")
print(df.to_string())

With ML Pipeline

from agents.data_agent import data_agent
from ml.predict import load_model_and_encoders, predict_dataframe

# 1. Query data
result = data_agent("Hamilton Silverstone 2023")
df = result["dataframe"]

# 2. Load model
model, encoders = load_model_and_encoders(run_id="abc123")

# 3. Predict
predictions = predict_dataframe(df, model, encoders)

# 4. Results
print(predictions[["driver_id", "driver_name", "win_probability"]])

Command Line

# Basic query
python -m agents.data_agent "Verstappen Bahrain 2023"

# JSON output
python -m agents.data_agent "Hamilton Monaco" --json

# Custom dataset
python -m agents.data_agent "Norris Austin" --data-path /path/to/races.parquet

Testing

Without API Key (Mock Parser)

from agents.data_agent import data_agent, QueryIntent

def mock_parser(query: str) -> QueryIntent:
    return {
        "season": 2023,
        "round": 1,
        "driver_id": "VER",
    }

result = data_agent("test", parser=mock_parser)  # No API key needed!

Run Test Suite

# All DataAgent tests
python -m pytest tests/test_data_agent.py -v

# Integration tests (end-to-end)
python -m pytest tests/test_integration_agent_predict.py -v -s

# Both
python -m pytest tests/test_data_agent.py tests/test_integration_agent_predict.py -v

Output Format

DataAgent returns DataAgentOutput TypedDict:

{
    "query": "Will Verstappen win Monaco 2023?",
    "intent": {
        "season": 2023,
        "round": 6,
        "driver_id": "VER",
        "driver_name": "Max Verstappen"
    },
    "rows": [
        {
            "season": 2023,
            "round": 6,
            "driver_id": "VER",
            "driver_name": "Max Verstappen",
            "team": "Red Bull Racing",
            "grid_position": 1.0,
            "finish_position": 1.0,
            "circuit_id": "Monaco Grand Prix",
            # ... 12 more columns
        }
    ],
    "dataframe": <pandas.DataFrame>  # 1 row, 20 columns
}

Supported Queries

The agent works with natural language like:

✅ "What was Max's win probability at Monaco 2023?"
✅ "Predict Hamilton Silverstone"
✅ "Verstappen Bahrain 2023 win chance"
✅ "Formula 1 prediction: Norris Austin 2024"
✅ "All drivers at Monza 2023" (no driver filter)

Error Handling

try:
    result = data_agent("Some race")
except ValueError as e:
    print(f"Data not found: {e}")
    # No rows found for season=X round=Y
    # No rows matched driver intent

except RuntimeError as e:
    print(f"Configuration error: {e}")
    # GROQ_API_KEY is required

Main Functions

Function	Purpose	Input	Output
`parse_query_with_groq()`	Parse NL → intent	str	QueryIntent
`build_prediction_dataframe()`	Filter data by intent	QueryIntent	DataFrame
`data_agent()`	Full pipeline	str	DataAgentOutput
`main()`	CLI entry point	sys.argv	stdout

Key Types

from agents.data_agent import (
    QueryIntent,           # season, round, driver_id, driver_name
    PredictionInputRow,    # 20 columns for prediction
    DataAgentOutput,       # query, intent, rows, dataframe
    IntentParser,          # Callable[[str], QueryIntent]
)

Configuration

File: agents/data_agent.py

DEFAULT_DATA_PATH = Path("data_output/fastf1_races.parquet")
GROQ_MODEL = "llama3-70b-8192"

Override at runtime:

result = data_agent(
    query="Verstappen Monaco",
    data_path="/custom/races.parquet"
)

Groq Settings

Model: llama3-70b-8192
Temperature: 0 (deterministic JSON output)
Max Tokens: Default (2048)
Context Window: 8K tokens

Perfect for structured output extraction.

Pipeline Compatibility

DataAgent output is 100% compatible with:

✅ prepare_model_data() — Feature engineering
✅ predict_dataframe() — ML inference
✅ load_model_and_encoders() — Model loading
✅ SHAP explanations

No intermediate transformations needed.

Files

File	Purpose
agents/data_agent.py	Main implementation
tests/test_data_agent.py	Unit tests
tests/test_integration_agent_predict.py	Integration tests
AGENT_ARCHITECTURE.md	Full technical docs

Examples

Single Driver Query

result = data_agent("Verstappen Bahrain 2023")
# Returns 1 row for VER at 2023 R1

Multiple Drivers (All in Race)

def no_driver_parser(q: str):
    return {"season": 2023, "round": 1, "driver_id": None}

result = data_agent("Bahrain 2023", parser=no_driver_parser)
# Returns 20 rows (all drivers at 2023 Bahrain)

Offline Testing

def mock_parser(q: str):
    return {"season": 2023, "round": 1, "driver_id": "HAM"}

result = data_agent("mock query", parser=mock_parser)
# Works without GROQ_API_KEY

Status

✅ Production-ready
✅ Fully tested (8 tests passing)
✅ ML pipeline integrated
✅ Type-safe (TypedDict throughout)
✅ Error handling complete

Ready to deploy!