Simple LangGraph Bike Training Agent

main.py

import json
import os
from datetime import date, datetime
from pathlib import Path
from typing import Annotated

from langchain_litellm import ChatLiteLLM
from langchain_core.messages import HumanMessage
from langchain_core.tools import tool
from langgraph.prebuilt import create_react_agent

WORKOUTS_FILE = Path("workouts.json")


def _load_workouts() -> list[dict]:
    if not WORKOUTS_FILE.exists():
        return []
    return json.loads(WORKOUTS_FILE.read_text())


def _save_workouts(workouts: list[dict]) -> None:
    WORKOUTS_FILE.write_text(json.dumps(workouts, indent=2))


@tool
def log_workout(
    workout_type: Annotated[str, "Type of workout: ride, interval, recovery, race, etc."],
    duration_min: Annotated[int, "Duration in minutes"],
    distance_km: Annotated[float | None, "Distance in kilometers, if applicable"] = None,
    avg_power_w: Annotated[int | None, "Average power in watts, if applicable"] = None,
    avg_hr_bpm: Annotated[int | None, "Average heart rate in BPM, if applicable"] = None,
    notes: Annotated[str | None, "Any notes about the workout"] = None,
    workout_date: Annotated[str | None, "Date in YYYY-MM-DD format, defaults to today"] = None,
) -> str:
    """Log a completed bike workout."""
    entry = {
        "date": workout_date or date.today().isoformat(),
        "type": workout_type,
        "duration_min": duration_min,
        "distance_km": distance_km,
        "avg_power_w": avg_power_w,
        "avg_hr_bpm": avg_hr_bpm,
        "notes": notes,
    }
    workouts = _load_workouts()
    workouts.append(entry)
    _save_workouts(workouts)
    return f"Logged {workout_type} workout on {entry['date']} ({duration_min} min)."


@tool
def get_recent_workouts(
    days: Annotated[int, "Number of days to look back"] = 14,
) -> str:
    """Get recent workouts from the training log."""
    workouts = _load_workouts()
    if not workouts:
        return "No workouts logged yet."
    cutoff = date.today().toordinal() - days
    recent = [w for w in workouts if date.fromisoformat(w["date"]).toordinal() >= cutoff]
    if not recent:
        return f"No workouts in the last {days} days."
    lines = []
    for w in sorted(recent, key=lambda x: x["date"], reverse=True):
        parts = [f"{w['date']} | {w['type']} | {w['duration_min']} min"]
        if w.get("distance_km"):
            parts.append(f"{w['distance_km']} km")
        if w.get("avg_power_w"):
            parts.append(f"{w['avg_power_w']}W avg")
        if w.get("avg_hr_bpm"):
            parts.append(f"{w['avg_hr_bpm']} bpm")
        if w.get("notes"):
            parts.append(f"— {w['notes']}")
        lines.append(" | ".join(parts))
    return "\n".join(lines)


@tool
def get_training_stats() -> str:
    """Get summary training statistics across all logged workouts."""
    workouts = _load_workouts()
    if not workouts:
        return "No workouts logged yet."

    total_rides = len(workouts)
    total_min = sum(w["duration_min"] for w in workouts)
    total_km = sum(w.get("distance_km") or 0 for w in workouts)
    power_entries = [w["avg_power_w"] for w in workouts if w.get("avg_power_w")]
    avg_power = sum(power_entries) / len(power_entries) if power_entries else None

    by_type: dict[str, int] = {}
    for w in workouts:
        by_type[w["type"]] = by_type.get(w["type"], 0) + 1

    lines = [
        f"Total workouts: {total_rides}",
        f"Total time: {total_min // 60}h {total_min % 60}m",
        f"Total distance: {total_km:.1f} km",
    ]
    if avg_power:
        lines.append(f"Average power: {avg_power:.0f}W")
    lines.append("By type: " + ", ".join(f"{k} ({v})" for k, v in by_type.items()))
    return "\n".join(lines)


SYSTEM_PROMPT = """You are a knowledgeable and encouraging cycling coach assistant.
Help the user log workouts, review their training history, and provide training advice.
When giving recommendations, consider periodization, recovery, and progressive overload.
Keep responses concise and practical. Today's date is {today}."""


def main() -> None:
    api_base = os.environ.get("DATAROBOT_ENDPOINT", "https://app.datarobot.com/api/v2")
    api_key = os.environ.get("DATAROBOT_API_TOKEN", "")

    llm = ChatLiteLLM(
        model="datarobot/bedrock/anthropic.claude-opus-4-8",
        api_base=api_base,
        api_key=api_key,
    )

    tools = [log_workout, get_recent_workouts, get_training_stats]
    system = SYSTEM_PROMPT.format(today=date.today().isoformat())
    agent = create_react_agent(llm, tools, prompt=system)

    print("Bike Training Assistant (type 'quit' to exit)\n")
    messages = []
    while True:
        try:
            user_input = input("You: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nRide on!")
            break
        if not user_input:
            continue
        if user_input.lower() in ("quit", "exit", "q"):
            print("Ride on!")
            break

        messages.append(HumanMessage(content=user_input))
        result = agent.invoke({"messages": messages})
        messages = result["messages"]
        reply = messages[-1].content
        print(f"\nCoach: {reply}\n")


if __name__ == "__main__":
    main()