invokelens-sdk 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.
Files changed (29) hide show
  1. invokelens_sdk-0.1.0/LICENSE +21 -0
  2. invokelens_sdk-0.1.0/PKG-INFO +157 -0
  3. invokelens_sdk-0.1.0/README.md +123 -0
  4. invokelens_sdk-0.1.0/pyproject.toml +63 -0
  5. invokelens_sdk-0.1.0/setup.cfg +4 -0
  6. invokelens_sdk-0.1.0/src/invokelens_sdk/__init__.py +22 -0
  7. invokelens_sdk-0.1.0/src/invokelens_sdk/_version.py +1 -0
  8. invokelens_sdk-0.1.0/src/invokelens_sdk/client.py +118 -0
  9. invokelens_sdk-0.1.0/src/invokelens_sdk/config.py +14 -0
  10. invokelens_sdk-0.1.0/src/invokelens_sdk/cost.py +36 -0
  11. invokelens_sdk-0.1.0/src/invokelens_sdk/decorators.py +345 -0
  12. invokelens_sdk-0.1.0/src/invokelens_sdk/exceptions.py +41 -0
  13. invokelens_sdk-0.1.0/src/invokelens_sdk/fingerprint.py +93 -0
  14. invokelens_sdk-0.1.0/src/invokelens_sdk/py.typed +0 -0
  15. invokelens_sdk-0.1.0/src/invokelens_sdk/schema.py +67 -0
  16. invokelens_sdk-0.1.0/src/invokelens_sdk/status.py +162 -0
  17. invokelens_sdk-0.1.0/src/invokelens_sdk/tracing.py +191 -0
  18. invokelens_sdk-0.1.0/src/invokelens_sdk/transport.py +218 -0
  19. invokelens_sdk-0.1.0/src/invokelens_sdk.egg-info/PKG-INFO +157 -0
  20. invokelens_sdk-0.1.0/src/invokelens_sdk.egg-info/SOURCES.txt +27 -0
  21. invokelens_sdk-0.1.0/src/invokelens_sdk.egg-info/dependency_links.txt +1 -0
  22. invokelens_sdk-0.1.0/src/invokelens_sdk.egg-info/requires.txt +12 -0
  23. invokelens_sdk-0.1.0/src/invokelens_sdk.egg-info/top_level.txt +1 -0
  24. invokelens_sdk-0.1.0/tests/test_cost.py +31 -0
  25. invokelens_sdk-0.1.0/tests/test_decorators.py +319 -0
  26. invokelens_sdk-0.1.0/tests/test_fingerprint.py +108 -0
  27. invokelens_sdk-0.1.0/tests/test_status.py +114 -0
  28. invokelens_sdk-0.1.0/tests/test_tracing.py +199 -0
  29. invokelens_sdk-0.1.0/tests/test_transport.py +364 -0
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 InvokeLens
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,157 @@
1
+ Metadata-Version: 2.4
2
+ Name: invokelens-sdk
3
+ Version: 0.1.0
4
+ Summary: SDK for InvokeLens — AI Agent Guardrails & Observability Platform
5
+ Author-email: InvokeLens <support@invokelens.com>
6
+ License: MIT
7
+ Project-URL: Homepage, https://invokelens.com
8
+ Project-URL: Repository, https://github.com/InvokeLens/invokelens-sdk
9
+ Project-URL: Changelog, https://github.com/InvokeLens/invokelens-sdk/blob/main/CHANGELOG.md
10
+ Project-URL: Documentation, https://docs.invokelens.com
11
+ Keywords: aws,bedrock,agents,observability,monitoring,guardrails,telemetry,llm
12
+ Classifier: Development Status :: 4 - Beta
13
+ Classifier: Intended Audience :: Developers
14
+ Classifier: Programming Language :: Python :: 3
15
+ Classifier: Programming Language :: Python :: 3.11
16
+ Classifier: Programming Language :: Python :: 3.12
17
+ Classifier: Topic :: Software Development :: Libraries :: Python Modules
18
+ Classifier: Topic :: System :: Monitoring
19
+ Classifier: Typing :: Typed
20
+ Requires-Python: >=3.11
21
+ Description-Content-Type: text/markdown
22
+ License-File: LICENSE
23
+ Requires-Dist: pydantic>=2.0
24
+ Requires-Dist: httpx>=0.24.0
25
+ Provides-Extra: eventbridge
26
+ Requires-Dist: boto3>=1.28.0; extra == "eventbridge"
27
+ Provides-Extra: dev
28
+ Requires-Dist: pytest>=7.0; extra == "dev"
29
+ Requires-Dist: pytest-asyncio; extra == "dev"
30
+ Requires-Dist: moto[dynamodb,events]>=4.0; extra == "dev"
31
+ Requires-Dist: mypy; extra == "dev"
32
+ Requires-Dist: ruff; extra == "dev"
33
+ Dynamic: license-file
34
+
35
+ # InvokeLens SDK
36
+
37
+ **Observability & Guardrails for AWS Bedrock Agents**
38
+
39
+ InvokeLens captures telemetry from your AWS Bedrock agent invocations — cost, latency, token usage, tool calls, and errors — and sends it to the InvokeLens platform for monitoring, alerting, and analysis.
40
+
41
+ ## Installation
42
+
43
+ ```bash
44
+ pip install invokelens-sdk
45
+ ```
46
+
47
+ ## Quick Start
48
+
49
+ ```python
50
+ from invokelens_sdk import InvokeLensClient
51
+
52
+ # Initialize the client
53
+ client = InvokeLensClient(
54
+ api_key="your-api-key",
55
+ endpoint_url="https://your-invokelens-api.com/v1",
56
+ )
57
+
58
+ # Decorate your Bedrock agent function
59
+ @client.observe(agent_id="my-agent", agent_name="Customer Support Bot")
60
+ def invoke_agent(prompt: str):
61
+ import boto3
62
+ bedrock = boto3.client("bedrock-agent-runtime")
63
+ response = bedrock.invoke_agent(
64
+ agentId="ABCDEFGHIJ",
65
+ agentAliasId="TSTALIASID",
66
+ sessionId="session-123",
67
+ inputText=prompt,
68
+ )
69
+ return response
70
+
71
+ # Call your function as normal — telemetry is captured automatically
72
+ result = invoke_agent("What is the status of order #1234?")
73
+
74
+ # Flush remaining events on shutdown
75
+ client.shutdown()
76
+ ```
77
+
78
+ ## Configuration
79
+
80
+ | Parameter | Default | Description |
81
+ |-----------|---------|-------------|
82
+ | `api_key` | *(required)* | Your InvokeLens API key |
83
+ | `endpoint_url` | `https://api.invokelens.com` | InvokeLens ingest endpoint URL |
84
+ | `transport_mode` | `"http"` | Transport backend: `"http"` or `"eventbridge"` |
85
+ | `event_bus_name` | `None` | EventBridge bus name (required if transport_mode is `"eventbridge"`) |
86
+ | `batch_size` | `10` | Number of events per batch flush |
87
+ | `flush_interval` | `5.0` | Seconds between automatic flushes |
88
+
89
+ ## What Gets Captured
90
+
91
+ The `@client.observe()` decorator automatically captures:
92
+
93
+ - **Timing** — invocation start, end, and duration
94
+ - **Token usage** — input and output token counts (auto-detected from Bedrock response)
95
+ - **Model ID** — which Bedrock model was used (auto-detected)
96
+ - **Cost estimate** — computed from token usage and model pricing
97
+ - **Status** — SUCCESS, FAILURE, or TIMEOUT
98
+ - **Errors** — exception type and message (truncated to 500 chars)
99
+ - **Tool calls** — names of tools invoked during execution
100
+
101
+ ### Optional Fields
102
+
103
+ You can enrich events with additional context:
104
+
105
+ ```python
106
+ @client.observe(
107
+ agent_id="my-agent",
108
+ agent_name="Customer Support Bot",
109
+ model_id="anthropic.claude-3-sonnet", # override auto-detection
110
+ )
111
+ def invoke_agent(prompt: str):
112
+ ...
113
+ ```
114
+
115
+ ## Cost Estimation
116
+
117
+ The SDK includes built-in pricing for common Bedrock models:
118
+
119
+ ```python
120
+ from invokelens_sdk import estimate_cost
121
+
122
+ cost = estimate_cost(
123
+ model_id="anthropic.claude-3-sonnet",
124
+ input_tokens=1000,
125
+ output_tokens=500,
126
+ )
127
+ print(f"Estimated cost: ${cost:.4f}")
128
+ ```
129
+
130
+ ## Transport Modes
131
+
132
+ ### HTTP (Default)
133
+
134
+ Sends batched events to the InvokeLens API via HTTPS. Includes automatic retry with exponential backoff (3 attempts).
135
+
136
+ ### EventBridge
137
+
138
+ Publishes events to an Amazon EventBridge bus. Useful for AWS-native architectures where you want to process events through EventBridge rules.
139
+
140
+ ```python
141
+ client = InvokeLensClient(
142
+ api_key="your-api-key",
143
+ transport_mode="eventbridge",
144
+ event_bus_name="invokelens-events",
145
+ )
146
+ ```
147
+
148
+ ## Requirements
149
+
150
+ - Python 3.11+
151
+ - `boto3` >= 1.28.0
152
+ - `pydantic` >= 2.0
153
+ - `httpx` >= 0.24.0
154
+
155
+ ## License
156
+
157
+ MIT License. See [LICENSE](LICENSE) for details.
@@ -0,0 +1,123 @@
1
+ # InvokeLens SDK
2
+
3
+ **Observability & Guardrails for AWS Bedrock Agents**
4
+
5
+ InvokeLens captures telemetry from your AWS Bedrock agent invocations — cost, latency, token usage, tool calls, and errors — and sends it to the InvokeLens platform for monitoring, alerting, and analysis.
6
+
7
+ ## Installation
8
+
9
+ ```bash
10
+ pip install invokelens-sdk
11
+ ```
12
+
13
+ ## Quick Start
14
+
15
+ ```python
16
+ from invokelens_sdk import InvokeLensClient
17
+
18
+ # Initialize the client
19
+ client = InvokeLensClient(
20
+ api_key="your-api-key",
21
+ endpoint_url="https://your-invokelens-api.com/v1",
22
+ )
23
+
24
+ # Decorate your Bedrock agent function
25
+ @client.observe(agent_id="my-agent", agent_name="Customer Support Bot")
26
+ def invoke_agent(prompt: str):
27
+ import boto3
28
+ bedrock = boto3.client("bedrock-agent-runtime")
29
+ response = bedrock.invoke_agent(
30
+ agentId="ABCDEFGHIJ",
31
+ agentAliasId="TSTALIASID",
32
+ sessionId="session-123",
33
+ inputText=prompt,
34
+ )
35
+ return response
36
+
37
+ # Call your function as normal — telemetry is captured automatically
38
+ result = invoke_agent("What is the status of order #1234?")
39
+
40
+ # Flush remaining events on shutdown
41
+ client.shutdown()
42
+ ```
43
+
44
+ ## Configuration
45
+
46
+ | Parameter | Default | Description |
47
+ |-----------|---------|-------------|
48
+ | `api_key` | *(required)* | Your InvokeLens API key |
49
+ | `endpoint_url` | `https://api.invokelens.com` | InvokeLens ingest endpoint URL |
50
+ | `transport_mode` | `"http"` | Transport backend: `"http"` or `"eventbridge"` |
51
+ | `event_bus_name` | `None` | EventBridge bus name (required if transport_mode is `"eventbridge"`) |
52
+ | `batch_size` | `10` | Number of events per batch flush |
53
+ | `flush_interval` | `5.0` | Seconds between automatic flushes |
54
+
55
+ ## What Gets Captured
56
+
57
+ The `@client.observe()` decorator automatically captures:
58
+
59
+ - **Timing** — invocation start, end, and duration
60
+ - **Token usage** — input and output token counts (auto-detected from Bedrock response)
61
+ - **Model ID** — which Bedrock model was used (auto-detected)
62
+ - **Cost estimate** — computed from token usage and model pricing
63
+ - **Status** — SUCCESS, FAILURE, or TIMEOUT
64
+ - **Errors** — exception type and message (truncated to 500 chars)
65
+ - **Tool calls** — names of tools invoked during execution
66
+
67
+ ### Optional Fields
68
+
69
+ You can enrich events with additional context:
70
+
71
+ ```python
72
+ @client.observe(
73
+ agent_id="my-agent",
74
+ agent_name="Customer Support Bot",
75
+ model_id="anthropic.claude-3-sonnet", # override auto-detection
76
+ )
77
+ def invoke_agent(prompt: str):
78
+ ...
79
+ ```
80
+
81
+ ## Cost Estimation
82
+
83
+ The SDK includes built-in pricing for common Bedrock models:
84
+
85
+ ```python
86
+ from invokelens_sdk import estimate_cost
87
+
88
+ cost = estimate_cost(
89
+ model_id="anthropic.claude-3-sonnet",
90
+ input_tokens=1000,
91
+ output_tokens=500,
92
+ )
93
+ print(f"Estimated cost: ${cost:.4f}")
94
+ ```
95
+
96
+ ## Transport Modes
97
+
98
+ ### HTTP (Default)
99
+
100
+ Sends batched events to the InvokeLens API via HTTPS. Includes automatic retry with exponential backoff (3 attempts).
101
+
102
+ ### EventBridge
103
+
104
+ Publishes events to an Amazon EventBridge bus. Useful for AWS-native architectures where you want to process events through EventBridge rules.
105
+
106
+ ```python
107
+ client = InvokeLensClient(
108
+ api_key="your-api-key",
109
+ transport_mode="eventbridge",
110
+ event_bus_name="invokelens-events",
111
+ )
112
+ ```
113
+
114
+ ## Requirements
115
+
116
+ - Python 3.11+
117
+ - `boto3` >= 1.28.0
118
+ - `pydantic` >= 2.0
119
+ - `httpx` >= 0.24.0
120
+
121
+ ## License
122
+
123
+ MIT License. See [LICENSE](LICENSE) for details.
@@ -0,0 +1,63 @@
1
+ [build-system]
2
+ requires = ["setuptools>=68.0"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "invokelens-sdk"
7
+ version = "0.1.0"
8
+ description = "SDK for InvokeLens — AI Agent Guardrails & Observability Platform"
9
+ readme = "README.md"
10
+ license = {text = "MIT"}
11
+ requires-python = ">=3.11"
12
+ authors = [
13
+ { name = "InvokeLens", email = "support@invokelens.com" },
14
+ ]
15
+ keywords = [
16
+ "aws",
17
+ "bedrock",
18
+ "agents",
19
+ "observability",
20
+ "monitoring",
21
+ "guardrails",
22
+ "telemetry",
23
+ "llm",
24
+ ]
25
+ classifiers = [
26
+ "Development Status :: 4 - Beta",
27
+ "Intended Audience :: Developers",
28
+ "Programming Language :: Python :: 3",
29
+ "Programming Language :: Python :: 3.11",
30
+ "Programming Language :: Python :: 3.12",
31
+ "Topic :: Software Development :: Libraries :: Python Modules",
32
+ "Topic :: System :: Monitoring",
33
+ "Typing :: Typed",
34
+ ]
35
+ dependencies = [
36
+ "pydantic>=2.0",
37
+ "httpx>=0.24.0",
38
+ ]
39
+
40
+ [project.optional-dependencies]
41
+ eventbridge = ["boto3>=1.28.0"]
42
+ dev = [
43
+ "pytest>=7.0",
44
+ "pytest-asyncio",
45
+ "moto[dynamodb,events]>=4.0",
46
+ "mypy",
47
+ "ruff",
48
+ ]
49
+
50
+ [project.urls]
51
+ Homepage = "https://invokelens.com"
52
+ Repository = "https://github.com/InvokeLens/invokelens-sdk"
53
+ Changelog = "https://github.com/InvokeLens/invokelens-sdk/blob/main/CHANGELOG.md"
54
+ Documentation = "https://docs.invokelens.com"
55
+
56
+ [tool.setuptools.packages.find]
57
+ where = ["src"]
58
+
59
+ [tool.pytest.ini_options]
60
+ testpaths = ["tests"]
61
+
62
+ [tool.ruff]
63
+ line-length = 100
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,22 @@
1
+ """InvokeLens SDK — AI Agent Observability & Guardrails Platform."""
2
+
3
+ from .client import InvokeLensClient
4
+ from .schema import TelemetryEvent
5
+ from .cost import estimate_cost, set_custom_pricing
6
+ from .exceptions import AgentBlockedError, PolicyViolationError
7
+ from .tracing import Span, TraceContext
8
+ from .fingerprint import compute_fingerprint
9
+ from ._version import __version__
10
+
11
+ __all__ = [
12
+ "InvokeLensClient",
13
+ "TelemetryEvent",
14
+ "estimate_cost",
15
+ "set_custom_pricing",
16
+ "AgentBlockedError",
17
+ "PolicyViolationError",
18
+ "Span",
19
+ "TraceContext",
20
+ "compute_fingerprint",
21
+ "__version__",
22
+ ]
@@ -0,0 +1 @@
1
+ __version__ = "0.1.0"
@@ -0,0 +1,118 @@
1
+ """Main entry point for the InvokeLens SDK."""
2
+
3
+ import functools
4
+ from typing import Optional
5
+
6
+ from .decorators import ObserveDecorator
7
+ from .transport import EventTransport
8
+ from .config import SDKConfig
9
+ from .status import AgentStatusChecker
10
+ from .tracing import TraceContext
11
+ from ._version import __version__
12
+
13
+
14
+ class InvokeLensClient:
15
+ """Client for sending agent telemetry to the InvokeLens platform.
16
+
17
+ Usage:
18
+ client = InvokeLensClient(api_key="il_live_abc123")
19
+
20
+ @client.observe(agent_id="my-agent", agent_name="Support Bot")
21
+ def ask_agent(prompt: str):
22
+ response = bedrock.invoke_agent(...)
23
+ return response
24
+
25
+ # On app shutdown:
26
+ client.shutdown()
27
+ """
28
+
29
+ def __init__(
30
+ self,
31
+ api_key: str,
32
+ endpoint_url: Optional[str] = None,
33
+ transport_mode: str = "http",
34
+ event_bus_name: Optional[str] = None,
35
+ batch_size: int = 10,
36
+ flush_interval: float = 5.0,
37
+ enable_kill_switch: bool = True,
38
+ status_check_ttl: float = 10.0,
39
+ ):
40
+ self.api_key = api_key
41
+ self.config = SDKConfig(
42
+ api_key=api_key,
43
+ endpoint_url=endpoint_url or "https://api.invokelens.com",
44
+ transport_mode=transport_mode,
45
+ event_bus_name=event_bus_name,
46
+ )
47
+ self._transport = EventTransport(
48
+ endpoint_url=self.config.endpoint_url,
49
+ mode=transport_mode,
50
+ event_bus_name=event_bus_name,
51
+ api_key=api_key,
52
+ batch_size=batch_size,
53
+ flush_interval_seconds=flush_interval,
54
+ )
55
+
56
+ self._status_checker = None
57
+ if enable_kill_switch:
58
+ self._status_checker = AgentStatusChecker(
59
+ endpoint_url=self.config.endpoint_url,
60
+ api_key=api_key,
61
+ ttl_seconds=status_check_ttl,
62
+ )
63
+
64
+ def observe(
65
+ self,
66
+ agent_id: str,
67
+ agent_name: Optional[str] = None,
68
+ model_id: Optional[str] = None,
69
+ ):
70
+ """Decorator that wraps a function and emits telemetry.
71
+
72
+ Args:
73
+ agent_id: Unique identifier for the agent.
74
+ agent_name: Human-readable name for the agent.
75
+ model_id: Bedrock model ID (auto-detected from response if possible).
76
+ """
77
+ return ObserveDecorator(
78
+ transport=self._transport,
79
+ agent_id=agent_id,
80
+ agent_name=agent_name,
81
+ model_id=model_id,
82
+ api_key=self.api_key,
83
+ sdk_version=__version__,
84
+ status_checker=self._status_checker,
85
+ )
86
+
87
+ def trace_tool(self, name: Optional[str] = None, span_type: str = "tool"):
88
+ """Decorator for tool functions that creates a span around the call.
89
+
90
+ The decorated function must accept a 'trace' keyword argument
91
+ (injected by @observe). If no trace context is present, the function
92
+ runs without tracing.
93
+
94
+ Usage:
95
+ @client.trace_tool(name="web_search")
96
+ def search(query: str, trace: TraceContext = None):
97
+ return do_search(query)
98
+ """
99
+ def decorator(func):
100
+ tool_name = name or func.__name__
101
+
102
+ @functools.wraps(func)
103
+ def wrapper(*args, **kwargs):
104
+ trace = kwargs.get("trace")
105
+ if not isinstance(trace, TraceContext):
106
+ return func(*args, **kwargs)
107
+
108
+ with trace.span(tool_name, span_type=span_type) as s:
109
+ result = func(*args, **kwargs)
110
+ s.output = str(result)[:2000] if result is not None else None
111
+ return result
112
+
113
+ return wrapper
114
+ return decorator
115
+
116
+ def shutdown(self):
117
+ """Flush pending events and clean up. Call on app exit."""
118
+ self._transport.shutdown()
@@ -0,0 +1,14 @@
1
+ """SDK configuration."""
2
+
3
+ from pydantic import BaseModel
4
+ from typing import Optional
5
+
6
+
7
+ class SDKConfig(BaseModel):
8
+ api_key: str
9
+ endpoint_url: str = "https://api.invokelens.com"
10
+ transport_mode: str = "http" # "http" or "eventbridge"
11
+ event_bus_name: Optional[str] = None
12
+ batch_size: int = 10
13
+ flush_interval: float = 5.0
14
+ max_queue_size: int = 1000
@@ -0,0 +1,36 @@
1
+ """Bedrock model pricing lookup for cost estimation."""
2
+
3
+ # Prices in USD per 1,000 tokens (approximate on-demand rates)
4
+ MODEL_PRICING: dict[str, dict[str, float]] = {
5
+ "anthropic.claude-3-5-sonnet": {"input": 0.003, "output": 0.015},
6
+ "anthropic.claude-3-sonnet": {"input": 0.003, "output": 0.015},
7
+ "anthropic.claude-3-haiku": {"input": 0.00025, "output": 0.00125},
8
+ "anthropic.claude-3-opus": {"input": 0.015, "output": 0.075},
9
+ "amazon.titan-text-lite-v1": {"input": 0.0003, "output": 0.0004},
10
+ "amazon.titan-text-express-v1": {"input": 0.0008, "output": 0.0016},
11
+ "meta.llama3-70b-instruct-v1": {"input": 0.00265, "output": 0.0035},
12
+ "meta.llama3-8b-instruct-v1": {"input": 0.0003, "output": 0.0006},
13
+ "mistral.mistral-large": {"input": 0.004, "output": 0.012},
14
+ "mistral.mistral-small": {"input": 0.001, "output": 0.003},
15
+ "cohere.command-r-plus-v1": {"input": 0.003, "output": 0.015},
16
+ "_default": {"input": 0.003, "output": 0.015},
17
+ }
18
+
19
+ _custom_pricing: dict[str, dict[str, float]] = {}
20
+
21
+
22
+ def set_custom_pricing(model_id: str, input_per_1k: float, output_per_1k: float):
23
+ """Override pricing for a specific model."""
24
+ _custom_pricing[model_id] = {"input": input_per_1k, "output": output_per_1k}
25
+
26
+
27
+ def estimate_cost(model_id: str, input_tokens: int, output_tokens: int) -> float:
28
+ """Estimate cost in USD for a given invocation."""
29
+ pricing = (
30
+ _custom_pricing.get(model_id)
31
+ or MODEL_PRICING.get(model_id)
32
+ or MODEL_PRICING["_default"]
33
+ )
34
+ input_cost = (input_tokens / 1000) * pricing["input"]
35
+ output_cost = (output_tokens / 1000) * pricing["output"]
36
+ return round(input_cost + output_cost, 8)