evenkeel 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,12 @@
1
+ .venv/
2
+ __pycache__/
3
+ *.pyc
4
+ .pytest_cache/
5
+ .ruff_cache/
6
+ dist/
7
+ evals/results/
8
+ .DS_Store
9
+
10
+ # Working documents: specs, plans and notes stay on this machine
11
+ docs/
12
+ .superpowers/
@@ -0,0 +1,20 @@
1
+ # Changelog
2
+
3
+ All notable changes to this project are documented here.
4
+
5
+ The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
6
+ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
7
+
8
+ ## [0.1.0] - 2026-10-13
9
+
10
+ First release.
11
+
12
+ ### Added
13
+ - `EvenKeelMiddleware` for LangChain agents (`create_agent`): classifies failed model and tool calls and
14
+ backs off, retries with a hint, fails, or escalates to a person with `interrupt()`.
15
+ - Dollar budgets from token usage and a user-supplied `PriceTable`, with a pause for approval when over budget.
16
+ - Attempt records in agent state (`evenkeel_attempts`).
17
+ - `RuleClassifier` and `JevClassifier` (optional `jev` extra), ported from keel.
18
+ - Contract tests for the LangChain, LangGraph and typesafe-sdk behaviour the middleware relies on.
19
+
20
+ [0.1.0]: https://github.com/yodablocks/evenkeel/releases/tag/v0.1.0
evenkeel-0.1.0/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 yodablocks
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,144 @@
1
+ Metadata-Version: 2.5
2
+ Name: evenkeel
3
+ Version: 0.1.0
4
+ Summary: Failure classification, retries with hints, dollar budgets and escalation for LangChain agents. Ported from keel.
5
+ Project-URL: Repository, https://github.com/yodablocks/evenkeel
6
+ Project-URL: keel, https://github.com/yodablocks/keel
7
+ Author: yodablocks
8
+ License: MIT
9
+ License-File: LICENSE
10
+ Keywords: agents,budgets,human-in-the-loop,jev,langchain,langgraph,middleware,retries
11
+ Classifier: Development Status :: 3 - Alpha
12
+ Classifier: Intended Audience :: Developers
13
+ Classifier: License :: OSI Approved :: MIT License
14
+ Classifier: Programming Language :: Python :: 3
15
+ Classifier: Programming Language :: Python :: 3.10
16
+ Classifier: Programming Language :: Python :: 3.11
17
+ Classifier: Programming Language :: Python :: 3.12
18
+ Classifier: Programming Language :: Python :: 3.13
19
+ Requires-Python: >=3.10
20
+ Requires-Dist: langchain<2,>=1.3
21
+ Requires-Dist: langgraph<2,>=1
22
+ Requires-Dist: typing-extensions>=4.7
23
+ Provides-Extra: jev
24
+ Requires-Dist: typesafe-sdk<0.8,>=0.7; extra == 'jev'
25
+ Description-Content-Type: text/markdown
26
+
27
+ # evenkeel
28
+
29
+ Keeps a LangChain agent on an even keel when things fail: failures are classified and handled by kind,
30
+ bad output is retried with a hint, spend is tracked in dollars with a budget, and failures that need a person
31
+ pause the run for approval.
32
+
33
+ Ported from [keel](https://github.com/yodablocks/keel), a durable execution engine for AI agents.
34
+ An independent project, not affiliated with LangChain. Experimental: 0.x, expect changes.
35
+
36
+ ## What it adds to LangChain's built-in middleware
37
+
38
+ | Need | Built in | EvenKeelMiddleware |
39
+ |---|---|---|
40
+ | Retries | `ModelRetryMiddleware`, `ToolRetryMiddleware` retry with the same input | Classifies the failure first: back off, retry **with a hint**, fail at once, or escalate |
41
+ | Limits | `ModelCallLimitMiddleware` counts calls | Tracks **dollars** from token usage and pauses at a budget |
42
+ | Approvals | `HumanInTheLoopMiddleware` asks before chosen tools | Escalates **when a failure needs a person**, from any model or tool call |
43
+
44
+ Use it alongside the built-ins: `ModelFallbackMiddleware` for switching models, `HumanInTheLoopMiddleware`
45
+ for approvals fixed in advance.
46
+
47
+ ## Install
48
+
49
+ pip install evenkeel # rule-based classification
50
+ pip install "evenkeel[jev]" # adds TypeSafe's Jev classifier
51
+
52
+ ## Use
53
+
54
+ ```python
55
+ from langchain.agents import create_agent
56
+ from langgraph.checkpoint.memory import InMemorySaver
57
+ from evenkeel import EvenKeelMiddleware, ModelPrice, PriceTable
58
+
59
+ agent = create_agent(
60
+ "openai:gpt-4o-mini",
61
+ tools=[...],
62
+ middleware=[
63
+ EvenKeelMiddleware(
64
+ budget_usd=0.50,
65
+ prices=PriceTable({"gpt-4o-mini": ModelPrice(input_per_mtok=0.15, output_per_mtok=0.60)}),
66
+ )
67
+ ],
68
+ checkpointer=InMemorySaver(), # needed for pauses
69
+ )
70
+ ```
71
+
72
+ With Jev (needs `TYPESAFE_API_KEY`):
73
+
74
+ ```python
75
+ from typesafe_sdk import TypeSafeClient
76
+ from evenkeel import EvenKeelMiddleware, JevClassifier
77
+
78
+ EvenKeelMiddleware(classifier=JevClassifier(TypeSafeClient()))
79
+ ```
80
+
81
+ ## How failures are handled
82
+
83
+ | Kind | Example | Action |
84
+ |---|---|---|
85
+ | `transient` | 429, 503, timeout, dropped connection | Back off and retry the same call |
86
+ | `bad_output` | Malformed output, a call that breaks the schema (recognised by Jev, or raised as `BadOutputError`) | Retry with a hint saying what was wrong |
87
+ | `bad_input` | Input a tool rejects (`BadInputError`, HTTP 400 or 422, a pydantic `ValidationError`) | Model call: fail. Tool call: return the error to the model so it fixes its arguments |
88
+ | `needs_human` | Refund over a limit, a policy block | Pause with `interrupt()` and wait for a decision |
89
+ | `fatal` | A bug, bad credentials | Re-raise the original error at once |
90
+
91
+ Retries and escalations stop at `max_attempts` (default 4). Each failure is recorded (error, kind, confidence,
92
+ which classifier decided, action): in the agent's state under `evenkeel_attempts` when the call ends in a
93
+ result, and on the raised error as `error.evenkeel_attempts` when it finally fails.
94
+
95
+ Some failures never reach the middleware, because LangChain turns them into messages for the model itself:
96
+ a call to a tool that does not exist, and tool arguments that fail the tool's schema.
97
+
98
+ ## Pauses and how to resume
99
+
100
+ A failure that needs a person:
101
+
102
+ ```python
103
+ from langgraph.types import Command
104
+
105
+ agent.invoke(Command(resume={"decision": "retry", "hint": "approved by finance"}), config)
106
+ agent.invoke(Command(resume={"decision": "fail"}), config)
107
+ agent.invoke(Command(resume={"decision": "skip", "message": "Refunded by hand."}), config) # tool calls only
108
+ ```
109
+
110
+ A `retry` with a `hint` on a tool call sends the reviewer's words to the model, which can call the tool again; a
111
+ tool cannot take a hint itself. A call that still needs a person after one approved retry fails. An answer
112
+ that is not one of these values does not raise: the run pauses again, with an `error` field in the payload
113
+ saying what was wrong.
114
+
115
+ Over budget:
116
+
117
+ ```python
118
+ agent.invoke(Command(resume={"approve": True, "budget_usd": 1.00}), config)
119
+ agent.invoke(Command(resume={"approve": False}), config) # ends the run
120
+ ```
121
+
122
+ The budget applies to each run: every new `invoke` on a thread starts with zero spend and the configured
123
+ budget.
124
+
125
+ ## Limitations
126
+
127
+ - **Budgets can overshoot by one call.** The check runs before a model call, but a call's cost is known only
128
+ after it. keel has the same limitation.
129
+ - **No built-in prices.** You supply them; a model without a price, or a response without token usage, is an
130
+ error rather than a silent $0. Dated or versioned names (`gpt-4o-mini-2024-07-18`, `claude-3-5-sonnet-20241022`, `gpt-4-0613`, `...-latest`)
131
+ use the price of their base name; any other variant (`gpt-4o-mini` against `gpt-4`, `gpt-4-turbo`) needs its own
132
+ entry.
133
+ - **On resume after an escalation, the failing call runs once more** before the decision applies, because
134
+ LangGraph re-runs the paused node. For a tool with side effects, make it idempotent or use `skip`.
135
+ - **Classification is only as good as the classifier.** Rules only recognise explicit signals (status codes,
136
+ network errors, evenkeel's error classes): on keel's 75 error messages from public SDK issues they pick the
137
+ right kind 29 times and the right action 39 times. Jev reads the message itself; see
138
+ [keel's results](https://github.com/yodablocks/keel#failure-classification-rules-vs-jev) for its accuracy on
139
+ real error messages, and `evals/` to measure it here.
140
+
141
+ ## Develop
142
+
143
+ uv sync --all-extras
144
+ uv run pytest
@@ -0,0 +1,118 @@
1
+ # evenkeel
2
+
3
+ Keeps a LangChain agent on an even keel when things fail: failures are classified and handled by kind,
4
+ bad output is retried with a hint, spend is tracked in dollars with a budget, and failures that need a person
5
+ pause the run for approval.
6
+
7
+ Ported from [keel](https://github.com/yodablocks/keel), a durable execution engine for AI agents.
8
+ An independent project, not affiliated with LangChain. Experimental: 0.x, expect changes.
9
+
10
+ ## What it adds to LangChain's built-in middleware
11
+
12
+ | Need | Built in | EvenKeelMiddleware |
13
+ |---|---|---|
14
+ | Retries | `ModelRetryMiddleware`, `ToolRetryMiddleware` retry with the same input | Classifies the failure first: back off, retry **with a hint**, fail at once, or escalate |
15
+ | Limits | `ModelCallLimitMiddleware` counts calls | Tracks **dollars** from token usage and pauses at a budget |
16
+ | Approvals | `HumanInTheLoopMiddleware` asks before chosen tools | Escalates **when a failure needs a person**, from any model or tool call |
17
+
18
+ Use it alongside the built-ins: `ModelFallbackMiddleware` for switching models, `HumanInTheLoopMiddleware`
19
+ for approvals fixed in advance.
20
+
21
+ ## Install
22
+
23
+ pip install evenkeel # rule-based classification
24
+ pip install "evenkeel[jev]" # adds TypeSafe's Jev classifier
25
+
26
+ ## Use
27
+
28
+ ```python
29
+ from langchain.agents import create_agent
30
+ from langgraph.checkpoint.memory import InMemorySaver
31
+ from evenkeel import EvenKeelMiddleware, ModelPrice, PriceTable
32
+
33
+ agent = create_agent(
34
+ "openai:gpt-4o-mini",
35
+ tools=[...],
36
+ middleware=[
37
+ EvenKeelMiddleware(
38
+ budget_usd=0.50,
39
+ prices=PriceTable({"gpt-4o-mini": ModelPrice(input_per_mtok=0.15, output_per_mtok=0.60)}),
40
+ )
41
+ ],
42
+ checkpointer=InMemorySaver(), # needed for pauses
43
+ )
44
+ ```
45
+
46
+ With Jev (needs `TYPESAFE_API_KEY`):
47
+
48
+ ```python
49
+ from typesafe_sdk import TypeSafeClient
50
+ from evenkeel import EvenKeelMiddleware, JevClassifier
51
+
52
+ EvenKeelMiddleware(classifier=JevClassifier(TypeSafeClient()))
53
+ ```
54
+
55
+ ## How failures are handled
56
+
57
+ | Kind | Example | Action |
58
+ |---|---|---|
59
+ | `transient` | 429, 503, timeout, dropped connection | Back off and retry the same call |
60
+ | `bad_output` | Malformed output, a call that breaks the schema (recognised by Jev, or raised as `BadOutputError`) | Retry with a hint saying what was wrong |
61
+ | `bad_input` | Input a tool rejects (`BadInputError`, HTTP 400 or 422, a pydantic `ValidationError`) | Model call: fail. Tool call: return the error to the model so it fixes its arguments |
62
+ | `needs_human` | Refund over a limit, a policy block | Pause with `interrupt()` and wait for a decision |
63
+ | `fatal` | A bug, bad credentials | Re-raise the original error at once |
64
+
65
+ Retries and escalations stop at `max_attempts` (default 4). Each failure is recorded (error, kind, confidence,
66
+ which classifier decided, action): in the agent's state under `evenkeel_attempts` when the call ends in a
67
+ result, and on the raised error as `error.evenkeel_attempts` when it finally fails.
68
+
69
+ Some failures never reach the middleware, because LangChain turns them into messages for the model itself:
70
+ a call to a tool that does not exist, and tool arguments that fail the tool's schema.
71
+
72
+ ## Pauses and how to resume
73
+
74
+ A failure that needs a person:
75
+
76
+ ```python
77
+ from langgraph.types import Command
78
+
79
+ agent.invoke(Command(resume={"decision": "retry", "hint": "approved by finance"}), config)
80
+ agent.invoke(Command(resume={"decision": "fail"}), config)
81
+ agent.invoke(Command(resume={"decision": "skip", "message": "Refunded by hand."}), config) # tool calls only
82
+ ```
83
+
84
+ A `retry` with a `hint` on a tool call sends the reviewer's words to the model, which can call the tool again; a
85
+ tool cannot take a hint itself. A call that still needs a person after one approved retry fails. An answer
86
+ that is not one of these values does not raise: the run pauses again, with an `error` field in the payload
87
+ saying what was wrong.
88
+
89
+ Over budget:
90
+
91
+ ```python
92
+ agent.invoke(Command(resume={"approve": True, "budget_usd": 1.00}), config)
93
+ agent.invoke(Command(resume={"approve": False}), config) # ends the run
94
+ ```
95
+
96
+ The budget applies to each run: every new `invoke` on a thread starts with zero spend and the configured
97
+ budget.
98
+
99
+ ## Limitations
100
+
101
+ - **Budgets can overshoot by one call.** The check runs before a model call, but a call's cost is known only
102
+ after it. keel has the same limitation.
103
+ - **No built-in prices.** You supply them; a model without a price, or a response without token usage, is an
104
+ error rather than a silent $0. Dated or versioned names (`gpt-4o-mini-2024-07-18`, `claude-3-5-sonnet-20241022`, `gpt-4-0613`, `...-latest`)
105
+ use the price of their base name; any other variant (`gpt-4o-mini` against `gpt-4`, `gpt-4-turbo`) needs its own
106
+ entry.
107
+ - **On resume after an escalation, the failing call runs once more** before the decision applies, because
108
+ LangGraph re-runs the paused node. For a tool with side effects, make it idempotent or use `skip`.
109
+ - **Classification is only as good as the classifier.** Rules only recognise explicit signals (status codes,
110
+ network errors, evenkeel's error classes): on keel's 75 error messages from public SDK issues they pick the
111
+ right kind 29 times and the right action 39 times. Jev reads the message itself; see
112
+ [keel's results](https://github.com/yodablocks/keel#failure-classification-rules-vs-jev) for its accuracy on
113
+ real error messages, and `evals/` to measure it here.
114
+
115
+ ## Develop
116
+
117
+ uv sync --all-extras
118
+ uv run pytest
@@ -0,0 +1,34 @@
1
+ """Failure classification, retries with hints, dollar budgets and escalation for LangChain agents."""
2
+
3
+ from ._core import AttemptRecord
4
+ from .budget import ModelPrice, PriceTable
5
+ from .failure import (
6
+ BadInputError,
7
+ BadOutputError,
8
+ FailureContext,
9
+ FailureVerdict,
10
+ NeedsHumanError,
11
+ OverBudgetError,
12
+ RuleClassifier,
13
+ default_policy,
14
+ )
15
+ from .jev import JevClassifier
16
+ from .middleware import EvenKeelMiddleware
17
+
18
+ __version__ = "0.1.0"
19
+
20
+ __all__ = [
21
+ "AttemptRecord",
22
+ "BadInputError",
23
+ "BadOutputError",
24
+ "EvenKeelMiddleware",
25
+ "FailureContext",
26
+ "FailureVerdict",
27
+ "JevClassifier",
28
+ "ModelPrice",
29
+ "NeedsHumanError",
30
+ "OverBudgetError",
31
+ "PriceTable",
32
+ "RuleClassifier",
33
+ "default_policy",
34
+ ]
@@ -0,0 +1,21 @@
1
+ """The one place that imports LangChain and LangGraph internals. If an import path moves in a new release,
2
+ change it here only."""
3
+
4
+ from langchain.agents.middleware import AgentMiddleware, AgentState, ModelRequest, ModelResponse, hook_config
5
+ from langchain.agents.middleware.types import ExtendedModelResponse
6
+ from langgraph.errors import GraphBubbleUp
7
+ from langgraph.prebuilt.tool_node import ToolCallRequest
8
+ from langgraph.types import Command, interrupt
9
+
10
+ __all__ = [
11
+ "AgentMiddleware",
12
+ "AgentState",
13
+ "Command",
14
+ "ExtendedModelResponse",
15
+ "GraphBubbleUp",
16
+ "ModelRequest",
17
+ "ModelResponse",
18
+ "ToolCallRequest",
19
+ "hook_config",
20
+ "interrupt",
21
+ ]
@@ -0,0 +1,126 @@
1
+ """Decision logic shared by every integration: apply the policy, enforce the attempt limit, build records,
2
+ and read reviewers' resume values. No LangChain imports, so other integrations can reuse it."""
3
+
4
+ from __future__ import annotations
5
+
6
+ from dataclasses import dataclass
7
+ from typing import Any, Literal
8
+
9
+ from typing_extensions import TypedDict
10
+
11
+ from .failure import (
12
+ Escalate,
13
+ Fail,
14
+ FailureAction,
15
+ FailureContext,
16
+ FailurePolicy,
17
+ FailureVerdict,
18
+ Retry,
19
+ RetryWithHint,
20
+ error_message,
21
+ )
22
+
23
+ _RECORD_ERROR_CHARS = 500
24
+
25
+
26
+ class AttemptRecord(TypedDict):
27
+ attempt: int
28
+ call: str
29
+ error_type: str
30
+ error: str
31
+ kind: str
32
+ confidence: float
33
+ source: str
34
+ action: str
35
+
36
+
37
+ @dataclass(frozen=True)
38
+ class Decision:
39
+ action: FailureAction
40
+ record: AttemptRecord
41
+
42
+
43
+ def decide(ctx: FailureContext, verdict: FailureVerdict, policy: FailurePolicy) -> Decision:
44
+ action = policy(verdict, ctx)
45
+ if isinstance(action, (Retry, RetryWithHint, Escalate)) and ctx.attempt >= ctx.max_attempts:
46
+ action = Fail(reason=f"attempt limit reached ({ctx.max_attempts}): {error_message(ctx.error)}")
47
+ record = AttemptRecord(
48
+ attempt=ctx.attempt,
49
+ call=ctx.step or "",
50
+ error_type=type(ctx.error).__name__,
51
+ error=error_message(ctx.error)[:_RECORD_ERROR_CHARS],
52
+ kind=verdict.kind,
53
+ confidence=round(verdict.confidence, 3),
54
+ source=verdict.source,
55
+ action=action.type,
56
+ )
57
+ return Decision(action=action, record=record)
58
+
59
+
60
+ def escalation_payload(ctx: FailureContext, verdict: FailureVerdict, reason: str) -> dict[str, Any]:
61
+ is_tool = (ctx.step or "").startswith("tool:")
62
+ return {
63
+ "type": "evenkeel_escalation",
64
+ "call": ctx.step or "",
65
+ "error_type": type(ctx.error).__name__,
66
+ "error": error_message(ctx.error)[:_RECORD_ERROR_CHARS],
67
+ "kind": verdict.kind,
68
+ "confidence": round(verdict.confidence, 3),
69
+ "reason": reason,
70
+ "accepted": ["retry", "fail", "skip"] if is_tool else ["retry", "fail"],
71
+ }
72
+
73
+
74
+ @dataclass(frozen=True)
75
+ class Resume:
76
+ decision: Literal["retry", "fail", "skip"]
77
+ hint: str | None
78
+ message: str | None
79
+
80
+
81
+ _RESUME_HELP = (
82
+ 'Resume with {"decision": "retry"} (optionally with "hint": "..."), {"decision": "fail"}, '
83
+ 'or, for tool calls only, {"decision": "skip", "message": "..."}'
84
+ )
85
+
86
+
87
+ def parse_resume(value: Any, *, call_type: Literal["model", "tool"]) -> Resume:
88
+ if not isinstance(value, dict) or value.get("decision") not in ("retry", "fail", "skip"):
89
+ raise ValueError(f"Unrecognised escalation decision {value!r}. {_RESUME_HELP}")
90
+ decision = value["decision"]
91
+ hint = value.get("hint")
92
+ message = value.get("message")
93
+ if hint is not None and not isinstance(hint, str):
94
+ raise ValueError(f"The decision's hint must be a string, got {hint!r}. {_RESUME_HELP}")
95
+ if decision == "skip":
96
+ if call_type == "model":
97
+ raise ValueError(f"A model call cannot be skipped: it has no stand-in result. {_RESUME_HELP}")
98
+ if not isinstance(message, str):
99
+ raise ValueError(f"A skip decision needs a message for the model. {_RESUME_HELP}")
100
+ return Resume("skip", None, message)
101
+ return Resume(decision, hint, None)
102
+
103
+
104
+ _BUDGET_HELP = 'Resume with {"approve": true, "budget_usd": <more than spent>} or {"approve": false}'
105
+
106
+
107
+ def parse_budget_resume(value: Any, *, spent_usd: float) -> float | None:
108
+ if not isinstance(value, dict) or not isinstance(value.get("approve"), bool):
109
+ raise ValueError(f"Unrecognised budget decision {value!r}. {_BUDGET_HELP}")
110
+ if not value["approve"]:
111
+ return None
112
+ budget = value.get("budget_usd")
113
+ if not isinstance(budget, (int, float)) or isinstance(budget, bool) or budget <= spent_usd:
114
+ raise ValueError(
115
+ f"An approval needs a budget_usd above the ${spent_usd:.4f} already spent, "
116
+ f"got {budget!r}. {_BUDGET_HELP}"
117
+ )
118
+ return float(budget)
119
+
120
+
121
+ def attach_attempts(error: BaseException, records: list[AttemptRecord]) -> None:
122
+ try:
123
+ error.evenkeel_attempts = list(records) # type: ignore[attr-defined]
124
+ except AttributeError: # some built-in exceptions refuse new attributes
125
+ if hasattr(error, "add_note"): # Python 3.11+
126
+ error.add_note(f"evenkeel attempts: {records}")
@@ -0,0 +1,74 @@
1
+ """Turns a model response's token usage into dollars, from prices the user supplies."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import re
6
+ from collections.abc import Mapping
7
+ from dataclasses import dataclass
8
+
9
+ from langchain_core.messages import AIMessage
10
+
11
+
12
+ class UnknownModelPriceError(ValueError):
13
+ """The response's model is not in the price table."""
14
+
15
+
16
+ class MissingUsageError(ValueError):
17
+ """The response reported no token usage, so its cost is unknown."""
18
+
19
+
20
+ @dataclass(frozen=True)
21
+ class ModelPrice:
22
+ input_per_mtok: float
23
+ """Dollars per million input tokens."""
24
+ output_per_mtok: float
25
+ """Dollars per million output tokens."""
26
+
27
+
28
+ def model_name_of(message: AIMessage) -> str | None:
29
+ metadata = message.response_metadata or {}
30
+ name = metadata.get("model_name") or metadata.get("model")
31
+ return name if isinstance(name, str) else None
32
+
33
+
34
+ # What may follow a priced name: only a date or version number (-2024-07-18, -20241022, -0613, -001) or the
35
+ # -latest alias. Anything else (-5-sonnet, -turbo, -1106-preview) is a different model with its own price.
36
+ _VERSION_SUFFIX = re.compile(r"-(\d{4}-\d{2}-\d{2}|\d{3,8}|latest)")
37
+
38
+
39
+ class PriceTable:
40
+ """Prices per model name. A name with a version or date suffix uses its base price
41
+ (gpt-4o-mini-2024-07-18 -> gpt-4o-mini); a different model never borrows another's price
42
+ (gpt-4o-mini does not use gpt-4, and gpt-4-turbo does not use gpt-4)."""
43
+
44
+ def __init__(self, prices: Mapping[str, ModelPrice]) -> None:
45
+ self._prices = dict(prices)
46
+
47
+ def price_for(self, model_name: str) -> ModelPrice:
48
+ if model_name in self._prices:
49
+ return self._prices[model_name]
50
+ matches = [
51
+ key
52
+ for key in self._prices
53
+ if model_name.startswith(key) and _VERSION_SUFFIX.fullmatch(model_name[len(key) :])
54
+ ]
55
+ if not matches:
56
+ raise UnknownModelPriceError(
57
+ f"No price for model {model_name!r}. Add it to PriceTable; known: {sorted(self._prices)}"
58
+ )
59
+ return self._prices[max(matches, key=len)]
60
+
61
+ def cost_usd(self, message: AIMessage) -> float:
62
+ name = model_name_of(message)
63
+ if name is None:
64
+ raise UnknownModelPriceError(
65
+ "The response does not say which model produced it; response_metadata keys: "
66
+ f"{sorted((message.response_metadata or {}).keys())}"
67
+ )
68
+ price = self.price_for(name)
69
+ usage = message.usage_metadata
70
+ if not usage:
71
+ raise MissingUsageError(f"Model {name!r} reported no token usage, so this call's cost is unknown")
72
+ return (
73
+ usage["input_tokens"] * price.input_per_mtok + usage["output_tokens"] * price.output_per_mtok
74
+ ) / 1e6