citadel-predict 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.
- citadel_predict-0.1.0/.gitignore +36 -0
- citadel_predict-0.1.0/PKG-INFO +215 -0
- citadel_predict-0.1.0/README.md +190 -0
- citadel_predict-0.1.0/pyproject.toml +56 -0
- citadel_predict-0.1.0/src/citadel_predict/__init__.py +30 -0
- citadel_predict-0.1.0/src/citadel_predict/cli.py +179 -0
- citadel_predict-0.1.0/src/citadel_predict/client.py +243 -0
- citadel_predict-0.1.0/src/citadel_predict/config.py +115 -0
- citadel_predict-0.1.0/src/citadel_predict/errors.py +109 -0
- citadel_predict-0.1.0/src/citadel_predict/py.typed +1 -0
- citadel_predict-0.1.0/tests/__init__.py +3 -0
- citadel_predict-0.1.0/tests/test_cli.py +118 -0
- citadel_predict-0.1.0/tests/test_client.py +110 -0
- citadel_predict-0.1.0/tests/test_config.py +75 -0
- citadel_predict-0.1.0/tests/test_errors.py +130 -0
- citadel_predict-0.1.0/tests/test_integration.py +90 -0
- citadel_predict-0.1.0/tests/test_safeguards.py +65 -0
|
@@ -0,0 +1,36 @@
|
|
|
1
|
+
# Virtual Environments
|
|
2
|
+
.venv/
|
|
3
|
+
**/.venv/
|
|
4
|
+
|
|
5
|
+
# Byte-compiled / cache files
|
|
6
|
+
__pycache__/
|
|
7
|
+
**/*.pyc
|
|
8
|
+
**/*.pyo
|
|
9
|
+
**/*.pyd
|
|
10
|
+
|
|
11
|
+
# Testing & Linting Caches
|
|
12
|
+
.pytest_cache/
|
|
13
|
+
.ruff_cache/
|
|
14
|
+
.mypy_cache/
|
|
15
|
+
|
|
16
|
+
# Node & Frontend
|
|
17
|
+
node_modules/
|
|
18
|
+
**/node_modules/
|
|
19
|
+
frontend/.next/
|
|
20
|
+
frontend/out/
|
|
21
|
+
|
|
22
|
+
# Environment Variables & Secrets
|
|
23
|
+
.env
|
|
24
|
+
.env.*
|
|
25
|
+
!.env.example
|
|
26
|
+
|
|
27
|
+
# IDE & Agents
|
|
28
|
+
.agents/
|
|
29
|
+
.gemini/
|
|
30
|
+
.vscode/
|
|
31
|
+
.idea/
|
|
32
|
+
|
|
33
|
+
# Distribution & Build
|
|
34
|
+
build/
|
|
35
|
+
dist/
|
|
36
|
+
*.egg-info/
|
|
@@ -0,0 +1,215 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: citadel-predict
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Pre-execution LLM token budget and cost prediction client for AI agents
|
|
5
|
+
Author: Citadel Predict Team
|
|
6
|
+
License: MIT
|
|
7
|
+
Keywords: agent,ai,budget,cost-estimation,governance,guardrails,llm,tokens
|
|
8
|
+
Classifier: Development Status :: 4 - Beta
|
|
9
|
+
Classifier: Intended Audience :: Developers
|
|
10
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
11
|
+
Classifier: Programming Language :: Python :: 3
|
|
12
|
+
Classifier: Programming Language :: Python :: 3.9
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
16
|
+
Classifier: Topic :: Software Development :: Libraries :: Python Modules
|
|
17
|
+
Requires-Python: >=3.9
|
|
18
|
+
Requires-Dist: httpx>=0.24.0
|
|
19
|
+
Provides-Extra: dev
|
|
20
|
+
Requires-Dist: mypy>=1.10.0; extra == 'dev'
|
|
21
|
+
Requires-Dist: pytest-cov>=4.1.0; extra == 'dev'
|
|
22
|
+
Requires-Dist: pytest>=8.0.0; extra == 'dev'
|
|
23
|
+
Requires-Dist: ruff>=0.4.0; extra == 'dev'
|
|
24
|
+
Description-Content-Type: text/markdown
|
|
25
|
+
|
|
26
|
+
# citadel-predict
|
|
27
|
+
|
|
28
|
+
[](https://pypi.org/project/citadel-predict/)
|
|
29
|
+
[](https://pypi.org/project/citadel-predict/)
|
|
30
|
+
[](https://opensource.org/licenses/MIT)
|
|
31
|
+
|
|
32
|
+
**Pre-execution token budget and cost predictor client for AI agents.**
|
|
33
|
+
|
|
34
|
+
`citadel-predict` is a lightweight, pure HTTP Python client and CLI for the Citadel Predict API. It allows developers, CI pipelines, and autonomous agent loops to estimate LLM token consumption and cost ranges *before* initiating expensive agent runs.
|
|
35
|
+
|
|
36
|
+
---
|
|
37
|
+
|
|
38
|
+
## Installation
|
|
39
|
+
|
|
40
|
+
```bash
|
|
41
|
+
pip install citadel-predict
|
|
42
|
+
```
|
|
43
|
+
|
|
44
|
+
---
|
|
45
|
+
|
|
46
|
+
## Quickstart (3 Lines of Code)
|
|
47
|
+
|
|
48
|
+
```python
|
|
49
|
+
from citadel_predict import predict_cost
|
|
50
|
+
|
|
51
|
+
result = predict_cost(
|
|
52
|
+
task_text="Research competitor pricing across 3 sources and draft report",
|
|
53
|
+
tools=["web_search", "draft_document"]
|
|
54
|
+
)
|
|
55
|
+
|
|
56
|
+
print(f"Expected: {result['expected_tokens']:,} tokens (Range: {result['low_tokens']:,} – {result['high_tokens']:,})")
|
|
57
|
+
```
|
|
58
|
+
|
|
59
|
+
Output:
|
|
60
|
+
```text
|
|
61
|
+
Expected: 3,200 tokens (Range: 1,500 – 5,800)
|
|
62
|
+
```
|
|
63
|
+
|
|
64
|
+
---
|
|
65
|
+
|
|
66
|
+
## CLI Usage
|
|
67
|
+
|
|
68
|
+
`citadel-predict` includes a full-featured CLI for terminal workflows and CI/CD cost checks:
|
|
69
|
+
|
|
70
|
+
```bash
|
|
71
|
+
# Pretty terminal card output
|
|
72
|
+
citadel-predict --task "Audit repository and write migration guide" --tools list_files,read_document,draft_document
|
|
73
|
+
|
|
74
|
+
# Scripting / CI mode (JSON output)
|
|
75
|
+
citadel-predict --task "Calculate statistical metrics" --tools calculator --json
|
|
76
|
+
|
|
77
|
+
# Override API key or URL
|
|
78
|
+
citadel-predict --task "..." --api-key "cp_live_12345" --api-url "https://api.citadel.dev"
|
|
79
|
+
```
|
|
80
|
+
|
|
81
|
+
### CLI Exit Codes
|
|
82
|
+
- `0`: Success
|
|
83
|
+
- `2`: Validation Error / Bad Request (HTTP 400 / 422 or missing task)
|
|
84
|
+
- `3`: Authentication Failure (HTTP 401)
|
|
85
|
+
- `4`: Rate Limit Exceeded (HTTP 429)
|
|
86
|
+
- `5`: Server Error (HTTP 5xx)
|
|
87
|
+
- `6`: Network / Timeout Error
|
|
88
|
+
|
|
89
|
+
---
|
|
90
|
+
|
|
91
|
+
## Real Agent Integration: Pre-Execution Guardrails
|
|
92
|
+
|
|
93
|
+
Existing agent governance tools (e.g., Portkey, Langfuse, LiteLLM) are **reactive**—they record costs during or after an execution. `citadel-predict` is **predictive**—enabling pre-flight budget checks and dynamic routing before running reasoning loops.
|
|
94
|
+
|
|
95
|
+
### LangGraph / CrewAI Pre-Flight Cost Guardrail Example
|
|
96
|
+
|
|
97
|
+
```python
|
|
98
|
+
from typing import TypedDict, List
|
|
99
|
+
from citadel_predict import predict_cost, CitadelError
|
|
100
|
+
|
|
101
|
+
class AgentState(TypedDict):
|
|
102
|
+
task: str
|
|
103
|
+
tools: List[str]
|
|
104
|
+
budget_tokens: int
|
|
105
|
+
approved: bool
|
|
106
|
+
|
|
107
|
+
def pre_flight_budget_guardrail(state: AgentState) -> AgentState:
|
|
108
|
+
"""
|
|
109
|
+
Evaluates token budget before dispatching tools or multi-agent loops.
|
|
110
|
+
"""
|
|
111
|
+
try:
|
|
112
|
+
prediction = predict_cost(
|
|
113
|
+
task_text=state["task"],
|
|
114
|
+
tools=state["tools"],
|
|
115
|
+
model_id="claude-sonnet"
|
|
116
|
+
)
|
|
117
|
+
except CitadelError as e:
|
|
118
|
+
print(f"Cost prediction unavailable: {e}. Falling back to default budget.")
|
|
119
|
+
return state
|
|
120
|
+
|
|
121
|
+
expected = prediction["expected_tokens"]
|
|
122
|
+
high = prediction["high_tokens"]
|
|
123
|
+
is_ood = prediction["out_of_distribution"]
|
|
124
|
+
|
|
125
|
+
print(f"Pre-flight estimate: ~{expected:,} tokens (Upper bound: {high:,})")
|
|
126
|
+
if is_ood:
|
|
127
|
+
print(f"Warning: Out-of-Distribution task ({prediction['ood_reasons']})")
|
|
128
|
+
|
|
129
|
+
# Guardrail Policy: Escalate if upper bound exceeds budget
|
|
130
|
+
if high > state["budget_tokens"]:
|
|
131
|
+
print(f"[BLOCKED] High-estimate ({high:,}) exceeds budget ({state['budget_tokens']:,})")
|
|
132
|
+
# In a real agent: switch to smaller model, ask human for approval, or prune tool access
|
|
133
|
+
state["approved"] = False
|
|
134
|
+
else:
|
|
135
|
+
state["approved"] = True
|
|
136
|
+
|
|
137
|
+
return state
|
|
138
|
+
|
|
139
|
+
# Example usage in workflow
|
|
140
|
+
initial_state: AgentState = {
|
|
141
|
+
"task": "Perform exhaustive market research across 20 industry filings",
|
|
142
|
+
"tools": ["web_search", "fetch_url", "draft_document"],
|
|
143
|
+
"budget_tokens": 10000,
|
|
144
|
+
"approved": False
|
|
145
|
+
}
|
|
146
|
+
|
|
147
|
+
state = pre_flight_budget_guardrail(initial_state)
|
|
148
|
+
if not state["approved"]:
|
|
149
|
+
print("Action required: Human-in-the-loop approval or task reformulation needed.")
|
|
150
|
+
```
|
|
151
|
+
|
|
152
|
+
---
|
|
153
|
+
|
|
154
|
+
## Authentication & Configuration
|
|
155
|
+
|
|
156
|
+
The client resolves your API key and base URL according to the following priority:
|
|
157
|
+
|
|
158
|
+
1. **Explicit argument**: `predict_cost(..., api_key="...", api_url="...")` or CLI `--api-key` / `--api-url`
|
|
159
|
+
2. **Environment variables**: `CITADEL_API_KEY` and `CITADEL_API_URL`
|
|
160
|
+
3. **Configuration file**: `~/.citadel/config.toml`
|
|
161
|
+
|
|
162
|
+
### Example `~/.citadel/config.toml`
|
|
163
|
+
```toml
|
|
164
|
+
api_key = "cp_live_your_api_key_here"
|
|
165
|
+
api_url = "https://api.citadel.dev"
|
|
166
|
+
```
|
|
167
|
+
|
|
168
|
+
---
|
|
169
|
+
|
|
170
|
+
## Error Handling
|
|
171
|
+
|
|
172
|
+
`citadel-predict` surfaces typed, catchable exceptions:
|
|
173
|
+
|
|
174
|
+
```python
|
|
175
|
+
from citadel_predict import (
|
|
176
|
+
predict_cost,
|
|
177
|
+
CitadelAuthError,
|
|
178
|
+
CitadelRateLimitError,
|
|
179
|
+
CitadelValidationError,
|
|
180
|
+
CitadelServerError,
|
|
181
|
+
CitadelNetworkError,
|
|
182
|
+
)
|
|
183
|
+
|
|
184
|
+
try:
|
|
185
|
+
result = predict_cost("Analyze dataset", tools=["calculator"])
|
|
186
|
+
except CitadelAuthError:
|
|
187
|
+
# 401: Missing or invalid API key
|
|
188
|
+
...
|
|
189
|
+
except CitadelRateLimitError as e:
|
|
190
|
+
# 429: Rate limited; check e.retry_after
|
|
191
|
+
print(f"Retry after {e.retry_after} seconds")
|
|
192
|
+
except CitadelValidationError as e:
|
|
193
|
+
# 422: Input validation bounds exceeded (e.g. task > 4000 chars)
|
|
194
|
+
...
|
|
195
|
+
except CitadelNetworkError as e:
|
|
196
|
+
# Timeout or connection failure
|
|
197
|
+
...
|
|
198
|
+
```
|
|
199
|
+
|
|
200
|
+
---
|
|
201
|
+
|
|
202
|
+
## Honest Limitations
|
|
203
|
+
|
|
204
|
+
`citadel-predict` is a thin client wrapping the hosted calibration model. It directly inherits the current system characteristics:
|
|
205
|
+
|
|
206
|
+
1. **Single-Model Calibration**: Calibration is currently tuned specifically for **Claude Sonnet** (`claude-sonnet`). Future releases will introduce multi-model support via `model_id`.
|
|
207
|
+
2. **Calibration Dataset Scale**: Calibrated on $N=20$ diverse task archetypes across 80 benchmarked runs.
|
|
208
|
+
3. **Synthetic Tool Sizing**: Ground-truth data was collected using deterministic mock tool outputs with representative context expansion. Real-world tools with unbounded payload returns (e.g., massive scraped DOMs) may exhibit higher variance.
|
|
209
|
+
4. **Pre-execution Estimation**: Token predictions represent calibrated statistical ranges $[low, expected, high]$, not runtime guarantees against infinite loops or divergent agent reasoning.
|
|
210
|
+
|
|
211
|
+
---
|
|
212
|
+
|
|
213
|
+
## License
|
|
214
|
+
|
|
215
|
+
MIT
|
|
@@ -0,0 +1,190 @@
|
|
|
1
|
+
# citadel-predict
|
|
2
|
+
|
|
3
|
+
[](https://pypi.org/project/citadel-predict/)
|
|
4
|
+
[](https://pypi.org/project/citadel-predict/)
|
|
5
|
+
[](https://opensource.org/licenses/MIT)
|
|
6
|
+
|
|
7
|
+
**Pre-execution token budget and cost predictor client for AI agents.**
|
|
8
|
+
|
|
9
|
+
`citadel-predict` is a lightweight, pure HTTP Python client and CLI for the Citadel Predict API. It allows developers, CI pipelines, and autonomous agent loops to estimate LLM token consumption and cost ranges *before* initiating expensive agent runs.
|
|
10
|
+
|
|
11
|
+
---
|
|
12
|
+
|
|
13
|
+
## Installation
|
|
14
|
+
|
|
15
|
+
```bash
|
|
16
|
+
pip install citadel-predict
|
|
17
|
+
```
|
|
18
|
+
|
|
19
|
+
---
|
|
20
|
+
|
|
21
|
+
## Quickstart (3 Lines of Code)
|
|
22
|
+
|
|
23
|
+
```python
|
|
24
|
+
from citadel_predict import predict_cost
|
|
25
|
+
|
|
26
|
+
result = predict_cost(
|
|
27
|
+
task_text="Research competitor pricing across 3 sources and draft report",
|
|
28
|
+
tools=["web_search", "draft_document"]
|
|
29
|
+
)
|
|
30
|
+
|
|
31
|
+
print(f"Expected: {result['expected_tokens']:,} tokens (Range: {result['low_tokens']:,} – {result['high_tokens']:,})")
|
|
32
|
+
```
|
|
33
|
+
|
|
34
|
+
Output:
|
|
35
|
+
```text
|
|
36
|
+
Expected: 3,200 tokens (Range: 1,500 – 5,800)
|
|
37
|
+
```
|
|
38
|
+
|
|
39
|
+
---
|
|
40
|
+
|
|
41
|
+
## CLI Usage
|
|
42
|
+
|
|
43
|
+
`citadel-predict` includes a full-featured CLI for terminal workflows and CI/CD cost checks:
|
|
44
|
+
|
|
45
|
+
```bash
|
|
46
|
+
# Pretty terminal card output
|
|
47
|
+
citadel-predict --task "Audit repository and write migration guide" --tools list_files,read_document,draft_document
|
|
48
|
+
|
|
49
|
+
# Scripting / CI mode (JSON output)
|
|
50
|
+
citadel-predict --task "Calculate statistical metrics" --tools calculator --json
|
|
51
|
+
|
|
52
|
+
# Override API key or URL
|
|
53
|
+
citadel-predict --task "..." --api-key "cp_live_12345" --api-url "https://api.citadel.dev"
|
|
54
|
+
```
|
|
55
|
+
|
|
56
|
+
### CLI Exit Codes
|
|
57
|
+
- `0`: Success
|
|
58
|
+
- `2`: Validation Error / Bad Request (HTTP 400 / 422 or missing task)
|
|
59
|
+
- `3`: Authentication Failure (HTTP 401)
|
|
60
|
+
- `4`: Rate Limit Exceeded (HTTP 429)
|
|
61
|
+
- `5`: Server Error (HTTP 5xx)
|
|
62
|
+
- `6`: Network / Timeout Error
|
|
63
|
+
|
|
64
|
+
---
|
|
65
|
+
|
|
66
|
+
## Real Agent Integration: Pre-Execution Guardrails
|
|
67
|
+
|
|
68
|
+
Existing agent governance tools (e.g., Portkey, Langfuse, LiteLLM) are **reactive**—they record costs during or after an execution. `citadel-predict` is **predictive**—enabling pre-flight budget checks and dynamic routing before running reasoning loops.
|
|
69
|
+
|
|
70
|
+
### LangGraph / CrewAI Pre-Flight Cost Guardrail Example
|
|
71
|
+
|
|
72
|
+
```python
|
|
73
|
+
from typing import TypedDict, List
|
|
74
|
+
from citadel_predict import predict_cost, CitadelError
|
|
75
|
+
|
|
76
|
+
class AgentState(TypedDict):
|
|
77
|
+
task: str
|
|
78
|
+
tools: List[str]
|
|
79
|
+
budget_tokens: int
|
|
80
|
+
approved: bool
|
|
81
|
+
|
|
82
|
+
def pre_flight_budget_guardrail(state: AgentState) -> AgentState:
|
|
83
|
+
"""
|
|
84
|
+
Evaluates token budget before dispatching tools or multi-agent loops.
|
|
85
|
+
"""
|
|
86
|
+
try:
|
|
87
|
+
prediction = predict_cost(
|
|
88
|
+
task_text=state["task"],
|
|
89
|
+
tools=state["tools"],
|
|
90
|
+
model_id="claude-sonnet"
|
|
91
|
+
)
|
|
92
|
+
except CitadelError as e:
|
|
93
|
+
print(f"Cost prediction unavailable: {e}. Falling back to default budget.")
|
|
94
|
+
return state
|
|
95
|
+
|
|
96
|
+
expected = prediction["expected_tokens"]
|
|
97
|
+
high = prediction["high_tokens"]
|
|
98
|
+
is_ood = prediction["out_of_distribution"]
|
|
99
|
+
|
|
100
|
+
print(f"Pre-flight estimate: ~{expected:,} tokens (Upper bound: {high:,})")
|
|
101
|
+
if is_ood:
|
|
102
|
+
print(f"Warning: Out-of-Distribution task ({prediction['ood_reasons']})")
|
|
103
|
+
|
|
104
|
+
# Guardrail Policy: Escalate if upper bound exceeds budget
|
|
105
|
+
if high > state["budget_tokens"]:
|
|
106
|
+
print(f"[BLOCKED] High-estimate ({high:,}) exceeds budget ({state['budget_tokens']:,})")
|
|
107
|
+
# In a real agent: switch to smaller model, ask human for approval, or prune tool access
|
|
108
|
+
state["approved"] = False
|
|
109
|
+
else:
|
|
110
|
+
state["approved"] = True
|
|
111
|
+
|
|
112
|
+
return state
|
|
113
|
+
|
|
114
|
+
# Example usage in workflow
|
|
115
|
+
initial_state: AgentState = {
|
|
116
|
+
"task": "Perform exhaustive market research across 20 industry filings",
|
|
117
|
+
"tools": ["web_search", "fetch_url", "draft_document"],
|
|
118
|
+
"budget_tokens": 10000,
|
|
119
|
+
"approved": False
|
|
120
|
+
}
|
|
121
|
+
|
|
122
|
+
state = pre_flight_budget_guardrail(initial_state)
|
|
123
|
+
if not state["approved"]:
|
|
124
|
+
print("Action required: Human-in-the-loop approval or task reformulation needed.")
|
|
125
|
+
```
|
|
126
|
+
|
|
127
|
+
---
|
|
128
|
+
|
|
129
|
+
## Authentication & Configuration
|
|
130
|
+
|
|
131
|
+
The client resolves your API key and base URL according to the following priority:
|
|
132
|
+
|
|
133
|
+
1. **Explicit argument**: `predict_cost(..., api_key="...", api_url="...")` or CLI `--api-key` / `--api-url`
|
|
134
|
+
2. **Environment variables**: `CITADEL_API_KEY` and `CITADEL_API_URL`
|
|
135
|
+
3. **Configuration file**: `~/.citadel/config.toml`
|
|
136
|
+
|
|
137
|
+
### Example `~/.citadel/config.toml`
|
|
138
|
+
```toml
|
|
139
|
+
api_key = "cp_live_your_api_key_here"
|
|
140
|
+
api_url = "https://api.citadel.dev"
|
|
141
|
+
```
|
|
142
|
+
|
|
143
|
+
---
|
|
144
|
+
|
|
145
|
+
## Error Handling
|
|
146
|
+
|
|
147
|
+
`citadel-predict` surfaces typed, catchable exceptions:
|
|
148
|
+
|
|
149
|
+
```python
|
|
150
|
+
from citadel_predict import (
|
|
151
|
+
predict_cost,
|
|
152
|
+
CitadelAuthError,
|
|
153
|
+
CitadelRateLimitError,
|
|
154
|
+
CitadelValidationError,
|
|
155
|
+
CitadelServerError,
|
|
156
|
+
CitadelNetworkError,
|
|
157
|
+
)
|
|
158
|
+
|
|
159
|
+
try:
|
|
160
|
+
result = predict_cost("Analyze dataset", tools=["calculator"])
|
|
161
|
+
except CitadelAuthError:
|
|
162
|
+
# 401: Missing or invalid API key
|
|
163
|
+
...
|
|
164
|
+
except CitadelRateLimitError as e:
|
|
165
|
+
# 429: Rate limited; check e.retry_after
|
|
166
|
+
print(f"Retry after {e.retry_after} seconds")
|
|
167
|
+
except CitadelValidationError as e:
|
|
168
|
+
# 422: Input validation bounds exceeded (e.g. task > 4000 chars)
|
|
169
|
+
...
|
|
170
|
+
except CitadelNetworkError as e:
|
|
171
|
+
# Timeout or connection failure
|
|
172
|
+
...
|
|
173
|
+
```
|
|
174
|
+
|
|
175
|
+
---
|
|
176
|
+
|
|
177
|
+
## Honest Limitations
|
|
178
|
+
|
|
179
|
+
`citadel-predict` is a thin client wrapping the hosted calibration model. It directly inherits the current system characteristics:
|
|
180
|
+
|
|
181
|
+
1. **Single-Model Calibration**: Calibration is currently tuned specifically for **Claude Sonnet** (`claude-sonnet`). Future releases will introduce multi-model support via `model_id`.
|
|
182
|
+
2. **Calibration Dataset Scale**: Calibrated on $N=20$ diverse task archetypes across 80 benchmarked runs.
|
|
183
|
+
3. **Synthetic Tool Sizing**: Ground-truth data was collected using deterministic mock tool outputs with representative context expansion. Real-world tools with unbounded payload returns (e.g., massive scraped DOMs) may exhibit higher variance.
|
|
184
|
+
4. **Pre-execution Estimation**: Token predictions represent calibrated statistical ranges $[low, expected, high]$, not runtime guarantees against infinite loops or divergent agent reasoning.
|
|
185
|
+
|
|
186
|
+
---
|
|
187
|
+
|
|
188
|
+
## License
|
|
189
|
+
|
|
190
|
+
MIT
|
|
@@ -0,0 +1,56 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "citadel-predict"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "Pre-execution LLM token budget and cost prediction client for AI agents"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.9"
|
|
11
|
+
license = { text = "MIT" }
|
|
12
|
+
authors = [
|
|
13
|
+
{ name = "Citadel Predict Team" }
|
|
14
|
+
]
|
|
15
|
+
keywords = [
|
|
16
|
+
"ai",
|
|
17
|
+
"agent",
|
|
18
|
+
"cost-estimation",
|
|
19
|
+
"tokens",
|
|
20
|
+
"llm",
|
|
21
|
+
"budget",
|
|
22
|
+
"guardrails",
|
|
23
|
+
"governance"
|
|
24
|
+
]
|
|
25
|
+
classifiers = [
|
|
26
|
+
"Development Status :: 4 - Beta",
|
|
27
|
+
"Intended Audience :: Developers",
|
|
28
|
+
"License :: OSI Approved :: MIT License",
|
|
29
|
+
"Programming Language :: Python :: 3",
|
|
30
|
+
"Programming Language :: Python :: 3.9",
|
|
31
|
+
"Programming Language :: Python :: 3.10",
|
|
32
|
+
"Programming Language :: Python :: 3.11",
|
|
33
|
+
"Programming Language :: Python :: 3.12",
|
|
34
|
+
"Topic :: Software Development :: Libraries :: Python Modules",
|
|
35
|
+
]
|
|
36
|
+
dependencies = [
|
|
37
|
+
"httpx>=0.24.0",
|
|
38
|
+
]
|
|
39
|
+
|
|
40
|
+
[project.optional-dependencies]
|
|
41
|
+
dev = [
|
|
42
|
+
"pytest>=8.0.0",
|
|
43
|
+
"pytest-cov>=4.1.0",
|
|
44
|
+
"ruff>=0.4.0",
|
|
45
|
+
"mypy>=1.10.0",
|
|
46
|
+
]
|
|
47
|
+
|
|
48
|
+
[project.scripts]
|
|
49
|
+
citadel-predict = "citadel_predict.cli:main"
|
|
50
|
+
|
|
51
|
+
[tool.hatch.build.targets.wheel]
|
|
52
|
+
packages = ["src/citadel_predict"]
|
|
53
|
+
|
|
54
|
+
[tool.pytest.ini_options]
|
|
55
|
+
testpaths = ["tests"]
|
|
56
|
+
pythonpath = ["src"]
|
|
@@ -0,0 +1,30 @@
|
|
|
1
|
+
"""
|
|
2
|
+
citadel-predict: Developer client for AI agent pre-execution cost prediction.
|
|
3
|
+
"""
|
|
4
|
+
|
|
5
|
+
from .client import CitadelClient, predict_cost
|
|
6
|
+
from .config import resolve_api_key, resolve_api_url
|
|
7
|
+
from .errors import (
|
|
8
|
+
CitadelAuthError,
|
|
9
|
+
CitadelBadRequestError,
|
|
10
|
+
CitadelError,
|
|
11
|
+
CitadelNetworkError,
|
|
12
|
+
CitadelRateLimitError,
|
|
13
|
+
CitadelServerError,
|
|
14
|
+
CitadelValidationError,
|
|
15
|
+
)
|
|
16
|
+
|
|
17
|
+
__version__ = "0.1.0"
|
|
18
|
+
__all__ = [
|
|
19
|
+
"predict_cost",
|
|
20
|
+
"CitadelClient",
|
|
21
|
+
"resolve_api_key",
|
|
22
|
+
"resolve_api_url",
|
|
23
|
+
"CitadelError",
|
|
24
|
+
"CitadelAuthError",
|
|
25
|
+
"CitadelRateLimitError",
|
|
26
|
+
"CitadelValidationError",
|
|
27
|
+
"CitadelBadRequestError",
|
|
28
|
+
"CitadelServerError",
|
|
29
|
+
"CitadelNetworkError",
|
|
30
|
+
]
|
|
@@ -0,0 +1,179 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Command-line interface for Citadel Predict.
|
|
3
|
+
|
|
4
|
+
Usage:
|
|
5
|
+
citadel-predict --task "Research 3 competitors" --tools web_search,draft_document
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
import argparse
|
|
9
|
+
import json
|
|
10
|
+
import sys
|
|
11
|
+
from typing import Optional, Sequence
|
|
12
|
+
|
|
13
|
+
from .client import CitadelClient
|
|
14
|
+
from .errors import (
|
|
15
|
+
CitadelAuthError,
|
|
16
|
+
CitadelBadRequestError,
|
|
17
|
+
CitadelError,
|
|
18
|
+
CitadelNetworkError,
|
|
19
|
+
CitadelRateLimitError,
|
|
20
|
+
CitadelServerError,
|
|
21
|
+
CitadelValidationError,
|
|
22
|
+
)
|
|
23
|
+
|
|
24
|
+
EXIT_SUCCESS = 0
|
|
25
|
+
EXIT_GENERAL_ERROR = 1
|
|
26
|
+
EXIT_VALIDATION_ERROR = 2
|
|
27
|
+
EXIT_AUTH_ERROR = 3
|
|
28
|
+
EXIT_RATE_LIMIT_ERROR = 4
|
|
29
|
+
EXIT_SERVER_ERROR = 5
|
|
30
|
+
EXIT_NETWORK_ERROR = 6
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
def parse_args(args: Optional[Sequence[str]] = None) -> argparse.Namespace:
|
|
34
|
+
parser = argparse.ArgumentParser(
|
|
35
|
+
prog="citadel-predict",
|
|
36
|
+
description="Predict AI agent token budgets and costs before execution.",
|
|
37
|
+
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
|
|
38
|
+
)
|
|
39
|
+
parser.add_argument(
|
|
40
|
+
"--task",
|
|
41
|
+
"-t",
|
|
42
|
+
type=str,
|
|
43
|
+
help="Task description string for the agent run.",
|
|
44
|
+
)
|
|
45
|
+
parser.add_argument(
|
|
46
|
+
"positional_task",
|
|
47
|
+
nargs="?",
|
|
48
|
+
type=str,
|
|
49
|
+
help="Task description (if --task is not passed).",
|
|
50
|
+
)
|
|
51
|
+
parser.add_argument(
|
|
52
|
+
"--tools",
|
|
53
|
+
type=str,
|
|
54
|
+
default="",
|
|
55
|
+
help="Comma-separated list of tool names (e.g. 'web_search,fetch_url,draft_document').",
|
|
56
|
+
)
|
|
57
|
+
parser.add_argument(
|
|
58
|
+
"--model-id",
|
|
59
|
+
"-m",
|
|
60
|
+
type=str,
|
|
61
|
+
default="claude-sonnet",
|
|
62
|
+
help="Model calibration identifier.",
|
|
63
|
+
)
|
|
64
|
+
parser.add_argument(
|
|
65
|
+
"--api-key",
|
|
66
|
+
"-k",
|
|
67
|
+
type=str,
|
|
68
|
+
default=None,
|
|
69
|
+
help="Citadel API key (overrides CITADEL_API_KEY env var and config file).",
|
|
70
|
+
)
|
|
71
|
+
parser.add_argument(
|
|
72
|
+
"--api-url",
|
|
73
|
+
"-u",
|
|
74
|
+
type=str,
|
|
75
|
+
default=None,
|
|
76
|
+
help="Citadel API base URL (default: http://localhost:8000 or CITADEL_API_URL).",
|
|
77
|
+
)
|
|
78
|
+
parser.add_argument(
|
|
79
|
+
"--timeout",
|
|
80
|
+
type=float,
|
|
81
|
+
default=10.0,
|
|
82
|
+
help="Request timeout in seconds.",
|
|
83
|
+
)
|
|
84
|
+
parser.add_argument(
|
|
85
|
+
"--json",
|
|
86
|
+
action="store_true",
|
|
87
|
+
help="Output raw JSON instead of formatted text.",
|
|
88
|
+
)
|
|
89
|
+
return parser.parse_args(args)
|
|
90
|
+
|
|
91
|
+
|
|
92
|
+
def format_pretty_output(result: dict) -> str:
|
|
93
|
+
lines = [
|
|
94
|
+
"=" * 50,
|
|
95
|
+
" CITADEL PREDICT — TOKEN BUDGET ESTIMATE",
|
|
96
|
+
"=" * 50,
|
|
97
|
+
f"Model: {result.get('model_id')}",
|
|
98
|
+
f"Expected Tokens: {result.get('expected_tokens', 0):,} tokens",
|
|
99
|
+
f"Predicted Range: {result.get('low_tokens', 0):,} – {result.get('high_tokens', 0):,} tokens",
|
|
100
|
+
f"Confidence: {result.get('confidence', 'unknown').upper()}",
|
|
101
|
+
]
|
|
102
|
+
|
|
103
|
+
if result.get("out_of_distribution"):
|
|
104
|
+
lines.append(f"OOD Warning: YES ({', '.join(result.get('ood_reasons', []))})")
|
|
105
|
+
else:
|
|
106
|
+
lines.append("OOD Warning: No (In-Distribution)")
|
|
107
|
+
|
|
108
|
+
factors = result.get("driving_factors", [])
|
|
109
|
+
if factors:
|
|
110
|
+
lines.append("Driving Factors:")
|
|
111
|
+
for factor in factors:
|
|
112
|
+
lines.append(f" • {factor}")
|
|
113
|
+
|
|
114
|
+
lines.append("=" * 50)
|
|
115
|
+
return "\n".join(lines)
|
|
116
|
+
|
|
117
|
+
|
|
118
|
+
def main(args: Optional[Sequence[str]] = None) -> int:
|
|
119
|
+
parsed = parse_args(args)
|
|
120
|
+
|
|
121
|
+
task = parsed.task or parsed.positional_task
|
|
122
|
+
if not task:
|
|
123
|
+
sys.stderr.write("Error: Task description is required. Use --task '...' or pass as argument.\n")
|
|
124
|
+
return EXIT_VALIDATION_ERROR
|
|
125
|
+
|
|
126
|
+
tools = [t.strip() for t in parsed.tools.split(",") if t.strip()] if parsed.tools else []
|
|
127
|
+
|
|
128
|
+
try:
|
|
129
|
+
with CitadelClient(
|
|
130
|
+
api_key=parsed.api_key,
|
|
131
|
+
base_url=parsed.api_url,
|
|
132
|
+
timeout=parsed.timeout,
|
|
133
|
+
) as client:
|
|
134
|
+
result = client.predict(
|
|
135
|
+
task_text=task,
|
|
136
|
+
tools=tools,
|
|
137
|
+
model_id=parsed.model_id,
|
|
138
|
+
)
|
|
139
|
+
|
|
140
|
+
if parsed.json:
|
|
141
|
+
print(json.dumps(result, indent=2))
|
|
142
|
+
else:
|
|
143
|
+
print(format_pretty_output(result))
|
|
144
|
+
return EXIT_SUCCESS
|
|
145
|
+
|
|
146
|
+
except CitadelAuthError as exc:
|
|
147
|
+
sys.stderr.write(f"Authentication Error: {exc.message}\n")
|
|
148
|
+
return EXIT_AUTH_ERROR
|
|
149
|
+
|
|
150
|
+
except CitadelRateLimitError as exc:
|
|
151
|
+
msg = f"Rate Limit Error: {exc.message}"
|
|
152
|
+
if exc.retry_after is not None:
|
|
153
|
+
msg += f" (Retry after {exc.retry_after}s)"
|
|
154
|
+
sys.stderr.write(f"{msg}\n")
|
|
155
|
+
return EXIT_RATE_LIMIT_ERROR
|
|
156
|
+
|
|
157
|
+
except (CitadelValidationError, CitadelBadRequestError) as exc:
|
|
158
|
+
sys.stderr.write(f"Validation Error: {exc.message}\n")
|
|
159
|
+
return EXIT_VALIDATION_ERROR
|
|
160
|
+
|
|
161
|
+
except CitadelServerError as exc:
|
|
162
|
+
sys.stderr.write(f"Server Error: {exc.message}\n")
|
|
163
|
+
return EXIT_SERVER_ERROR
|
|
164
|
+
|
|
165
|
+
except CitadelNetworkError as exc:
|
|
166
|
+
sys.stderr.write(f"Network Error: {exc.message}\n")
|
|
167
|
+
return EXIT_NETWORK_ERROR
|
|
168
|
+
|
|
169
|
+
except CitadelError as exc:
|
|
170
|
+
sys.stderr.write(f"Citadel Error: {exc.message}\n")
|
|
171
|
+
return EXIT_GENERAL_ERROR
|
|
172
|
+
|
|
173
|
+
except Exception as exc:
|
|
174
|
+
sys.stderr.write(f"Unexpected Error: {exc}\n")
|
|
175
|
+
return EXIT_GENERAL_ERROR
|
|
176
|
+
|
|
177
|
+
|
|
178
|
+
if __name__ == "__main__":
|
|
179
|
+
sys.exit(main())
|