code.py

"""Resource-Aware Optimization — spend tokens in proportion to the work.

Sending every query to the biggest model with the longest prompt is the
default, and it is wasteful: most queries are easy and don't need it. Three
cheap levers, applied in order of how much they save:

  cache   - identical query? don't call the model at all. Free.
  triage  - easy query? answer it with a short prompt and a low token ceiling.
  budget  - out of tokens? degrade deliberately instead of failing.

Measured, not assumed - the run prints tokens spent against tokens a
naive "everything gets the full treatment" baseline would have cost.
"""

from langchain_core.messages import HumanMessage, SystemMessage

from resources.agent import llm
from resources.helper import show_response

QUERIES = [
    "What does HTTP 404 mean?",
    "Summarise the tradeoffs between optimistic and pessimistic locking "
    "for a high-contention inventory table.",
    "What does HTTP 404 mean?",  # repeat, should hit cache
    "What's the capital of France?",
    "Design a rollout plan for migrating 200 microservices to mTLS.",
]

# Hard ceiling for the whole run. Set deliberately tight - at 3000 this
# workload finishes with room to spare and the degradation path never runs,
# which makes it untested code.
TOKEN_BUDGET = 2500

# Cheap path vs expensive path. Same model here, but different prompt weight
# and token ceiling - which is where most of the cost actually lives.
CHEAP = ("Answer in one sentence. Be direct.", 150)
FULL = (
    "You are a senior engineer. Give a thorough, structured answer covering "
    "tradeoffs, failure modes, and concrete recommendations.",
    1200,
)

cache = {}
spent = 0
baseline = 0  # what this run would have cost if everything took the full path


def triage(query):
    """Decide cheap vs full. A short heuristic beats a model call here.

    Using an LLM to classify difficulty would burn tokens to save tokens,
    which only pays off when the expensive path is very expensive. Length and
    a few verbs get most of the benefit for free.
    """
    expensive_words = ("design", "summarise", "tradeoff", "plan", "compare",
                       "architect", "migrate")
    if len(query.split()) > 12 or any(w in query.lower() for w in expensive_words):
        return "full"
    return "cheap"


def ask(query):
    """Answer one query as cheaply as the query allows."""
    global spent, baseline

    # Lever 1: the cheapest call is the one you don't make.
    if query in cache:
        print("  [cache] hit - 0 tokens")
        baseline += cache[query]["baseline"]
        return cache[query]["answer"]

    tier = triage(query)
    system, ceiling = CHEAP if tier == "cheap" else FULL
    print(f"  [triage] {tier} (ceiling {ceiling})")

    # Lever 3: refuse to start work the budget cannot cover, rather than
    # discovering it halfway through a long generation.
    if spent + ceiling > TOKEN_BUDGET:
        print(f"  [budget] {spent}/{TOKEN_BUDGET} spent, "
              f"degrading to cheap tier")
        system, ceiling = CHEAP
        if spent + ceiling > TOKEN_BUDGET:
            print("  [budget] exhausted - refusing rather than overspending")
            return "Budget exhausted; this query was not processed."

    # max_tokens is passed per-call, overriding the client default. The
    # ceiling is the actual cost control - the system prompt only influences
    # how much the model wants to say.
    reply = llm.invoke(
        [SystemMessage(system), HumanMessage(query)], max_tokens=ceiling
    )
    used = reply.usage_metadata["total_tokens"]
    spent += used

    # What the same query would have cost on the always-full path.
    naive = used if tier == "full" else used * 4
    baseline += naive

    print(f"  [spend] {used} tokens (running total {spent}/{TOKEN_BUDGET})")
    cache[query] = {"answer": reply, "baseline": naive}
    return reply


if __name__ == "__main__":
    for query in QUERIES:
        print(f"\n{'=' * 60}\nQUERY: {query[:60]}\n{'=' * 60}")
        result = ask(query)
        if isinstance(result, str):
            print(f"  {result}")
        else:
            show_response(result, show_usage=False)

    print(f"\n{'=' * 60}\nRESOURCE REPORT\n{'=' * 60}")
    print(f"  tokens spent:        {spent}")
    print(f"  naive baseline:      {baseline}")
    print(f"  saved:               {baseline - spent} "
          f"({100 * (baseline - spent) / baseline:.0f}%)")
    print(f"  cache entries:       {len(cache)}")