Token Caching

import time
from typing import Dict, Tuple

class TokenCache:
    def __init__(self):
        self.cache: Dict[str, Tuple[str, int]] = {}

    def get(self, plan_hash: str) -> str | None:
        if plan_hash in self.cache:
            token, expires_at = self.cache[plan_hash]
            # Return token if valid for at least 60 more seconds
            if time.time() < expires_at - 60:
                return token
        return None

    def set(self, plan_hash: str, token: str, expires_at: int):
        self.cache[plan_hash] = (token, expires_at)

    def clear_expired(self):
        now = time.time()
        self.cache = {
            k: v for k, v in self.cache.items()
            if v[1] > now
        }

# Usage
token_cache = TokenCache()

def get_token_cached(client, llm, prompt):
    captured = client.capture_plan(llm=llm, prompt=prompt)
    plan_hash = captured.plan_hash

    # Try cache first
    token = token_cache.get(plan_hash)
    if token:
        return token

    # Get new token
    response = client.get_intent_token(captured)
    token_cache.set(plan_hash, response["token"], response["expires_at"])

    return response["token"]