argus-gateway 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.
- argus_gateway-0.1.0/PKG-INFO +113 -0
- argus_gateway-0.1.0/README.md +92 -0
- argus_gateway-0.1.0/argus_gateway.egg-info/PKG-INFO +113 -0
- argus_gateway-0.1.0/argus_gateway.egg-info/SOURCES.txt +27 -0
- argus_gateway-0.1.0/argus_gateway.egg-info/dependency_links.txt +1 -0
- argus_gateway-0.1.0/argus_gateway.egg-info/entry_points.txt +2 -0
- argus_gateway-0.1.0/argus_gateway.egg-info/requires.txt +8 -0
- argus_gateway-0.1.0/argus_gateway.egg-info/top_level.txt +1 -0
- argus_gateway-0.1.0/gateway/__init__.py +7 -0
- argus_gateway-0.1.0/gateway/auth.py +108 -0
- argus_gateway-0.1.0/gateway/config.py +52 -0
- argus_gateway-0.1.0/gateway/ledger.py +106 -0
- argus_gateway-0.1.0/gateway/providers/__init__.py +5 -0
- argus_gateway-0.1.0/gateway/providers/anthropic.py +85 -0
- argus_gateway-0.1.0/gateway/providers/base.py +55 -0
- argus_gateway-0.1.0/gateway/providers/generic.py +29 -0
- argus_gateway-0.1.0/gateway/providers/mock.py +28 -0
- argus_gateway-0.1.0/gateway/providers/openai.py +65 -0
- argus_gateway-0.1.0/gateway/router.py +62 -0
- argus_gateway-0.1.0/gateway/security/__init__.py +29 -0
- argus_gateway-0.1.0/gateway/security/audit.py +72 -0
- argus_gateway-0.1.0/gateway/security/detectors.py +110 -0
- argus_gateway-0.1.0/gateway/security/policy.py +91 -0
- argus_gateway-0.1.0/gateway/server.py +222 -0
- argus_gateway-0.1.0/pyproject.toml +35 -0
- argus_gateway-0.1.0/setup.cfg +4 -0
- argus_gateway-0.1.0/tests/test_router.py +51 -0
- argus_gateway-0.1.0/tests/test_security.py +79 -0
- argus_gateway-0.1.0/tests/test_server.py +162 -0
|
@@ -0,0 +1,113 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: argus-gateway
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Argus: the security layer for all AI models — an OpenAI-compatible AI gateway with pluggable threat detectors, policy enforcement, audit, and cost intelligence
|
|
5
|
+
Author: Karmendra Pandey
|
|
6
|
+
License: MIT
|
|
7
|
+
Keywords: llm,ai-gateway,proxy,cost,routing,agents
|
|
8
|
+
Classifier: Programming Language :: Python :: 3
|
|
9
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
10
|
+
Classifier: Intended Audience :: Developers
|
|
11
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
12
|
+
Requires-Python: >=3.9
|
|
13
|
+
Description-Content-Type: text/markdown
|
|
14
|
+
Requires-Dist: fastapi>=0.110
|
|
15
|
+
Requires-Dist: uvicorn>=0.29
|
|
16
|
+
Requires-Dist: httpx>=0.27
|
|
17
|
+
Requires-Dist: tokenecon>=0.2.0
|
|
18
|
+
Provides-Extra: test
|
|
19
|
+
Requires-Dist: pytest>=8; extra == "test"
|
|
20
|
+
Requires-Dist: httpx>=0.27; extra == "test"
|
|
21
|
+
|
|
22
|
+
# Argus — the security layer for all AI models
|
|
23
|
+
|
|
24
|
+
An OpenAI-compatible gateway that sits in front of **any** LLM — OpenAI,
|
|
25
|
+
Anthropic, Bedrock, Gemini, Mistral, vLLM, Ollama, anything with an
|
|
26
|
+
OpenAI-shaped endpoint — and enforces security policy before a model ever
|
|
27
|
+
sees your data.
|
|
28
|
+
|
|
29
|
+
Every request passes through the policy engine on the way in and on the way
|
|
30
|
+
out. Cost intelligence (tiered routing, graceful budgets, per-task
|
|
31
|
+
attribution) rides along via [tokenecon](https://pypi.org/project/tokenecon/),
|
|
32
|
+
but security is the product.
|
|
33
|
+
|
|
34
|
+
## The security layer
|
|
35
|
+
|
|
36
|
+
```
|
|
37
|
+
request -> auth -> INPUT policy -> budget -> route -> model
|
|
38
|
+
-> OUTPUT policy -> audit -> response
|
|
39
|
+
```
|
|
40
|
+
|
|
41
|
+
**Pluggable detectors** (each scans, the policy decides):
|
|
42
|
+
|
|
43
|
+
| Detector | Catches | Default action |
|
|
44
|
+
|---|---|---|
|
|
45
|
+
| Secrets | AWS keys, private keys, GitHub/OpenAI tokens, `password = ...` | **block** — never reaches a model |
|
|
46
|
+
| PII | emails, phones, SSNs, credit cards, IPs | **redact** — `[REDACTED:EMAIL]`, request continues |
|
|
47
|
+
| Prompt injection | "ignore previous instructions", role overrides, system-prompt probes, known jailbreaks | **flag** — allowed but marked (one flag away from block) |
|
|
48
|
+
| Blocklist | your own regexes | **block** |
|
|
49
|
+
|
|
50
|
+
**Policy engine:** per-direction rules (`in` / `out` / `both`) with four
|
|
51
|
+
actions — `allow`, `redact`, `block`, `flag`. Add an ML-based detector later
|
|
52
|
+
by subclassing `Detector`; no other code changes.
|
|
53
|
+
|
|
54
|
+
**Audit log:** every decision recorded — timestamp, key, task, detector,
|
|
55
|
+
action, redacted snippet. Raw PII and secrets never touch the audit trail.
|
|
56
|
+
`GET /admin/audit` for compliance review.
|
|
57
|
+
|
|
58
|
+
## Quickstart
|
|
59
|
+
|
|
60
|
+
```bash
|
|
61
|
+
pip install argus-gateway
|
|
62
|
+
export AIGW_MASTER_KEY="sk-admin-secret"
|
|
63
|
+
export OPENAI_API_KEY="..."
|
|
64
|
+
argus # serves on 127.0.0.1:4000
|
|
65
|
+
```
|
|
66
|
+
|
|
67
|
+
Point any OpenAI client at it — one `base_url` change:
|
|
68
|
+
|
|
69
|
+
```python
|
|
70
|
+
import openai
|
|
71
|
+
client = openai.OpenAI(api_key="<virtual-key>", base_url="http://localhost:4000/v1")
|
|
72
|
+
```
|
|
73
|
+
|
|
74
|
+
Self-hosted models sit behind the same layer:
|
|
75
|
+
|
|
76
|
+
```python
|
|
77
|
+
from gateway.providers.generic import GenericProvider
|
|
78
|
+
# vLLM, Ollama, llama.cpp — any OpenAI-shaped endpoint
|
|
79
|
+
```
|
|
80
|
+
|
|
81
|
+
Every response carries what happened:
|
|
82
|
+
|
|
83
|
+
```json
|
|
84
|
+
"gateway": {
|
|
85
|
+
"provider": "openai", "tier": "small", "difficulty": 0.17,
|
|
86
|
+
"cost_usd": 0.000004, "degraded": false,
|
|
87
|
+
"security": {"redactions": ["pii:EMAIL"], "flags": ["prompt_injection:INSTRUCTION_OVERRIDE"]}
|
|
88
|
+
}
|
|
89
|
+
```
|
|
90
|
+
|
|
91
|
+
## Also inside: cost intelligence
|
|
92
|
+
|
|
93
|
+
- **Tiered routing** — requests scored for difficulty, cheapest capable tier wins
|
|
94
|
+
- **Graceful budgets** — over budget? Falls back to the cheapest tier and flags
|
|
95
|
+
`degraded: true` instead of failing (hard-stop mode available)
|
|
96
|
+
- **Per-task attribution** — `X-Task-Id` groups an agentic run's turns into one
|
|
97
|
+
costed task: `GET /admin/tasks/{id}`
|
|
98
|
+
|
|
99
|
+
## Admin API (`X-Admin-Key` required)
|
|
100
|
+
|
|
101
|
+
- `POST /admin/keys` — issue virtual keys with budgets (`degrade`/`hard`, daily/monthly)
|
|
102
|
+
- `GET /admin/spend`, `GET /admin/keys/{id}/spend` — spend tracking
|
|
103
|
+
- `GET /admin/tasks/{task_id}` — per-task cost breakdown
|
|
104
|
+
- `GET /admin/audit?limit=&key_id=` — security decisions
|
|
105
|
+
|
|
106
|
+
## Roadmap
|
|
107
|
+
|
|
108
|
+
Streaming (SSE), Bedrock/Gemini/Mistral adapters, semantic caching,
|
|
109
|
+
OTEL/Prometheus export, admin UI, SSO/RBAC, ML-based injection detector.
|
|
110
|
+
|
|
111
|
+
## License
|
|
112
|
+
|
|
113
|
+
MIT — Karmendra Pandey.
|
|
@@ -0,0 +1,92 @@
|
|
|
1
|
+
# Argus — the security layer for all AI models
|
|
2
|
+
|
|
3
|
+
An OpenAI-compatible gateway that sits in front of **any** LLM — OpenAI,
|
|
4
|
+
Anthropic, Bedrock, Gemini, Mistral, vLLM, Ollama, anything with an
|
|
5
|
+
OpenAI-shaped endpoint — and enforces security policy before a model ever
|
|
6
|
+
sees your data.
|
|
7
|
+
|
|
8
|
+
Every request passes through the policy engine on the way in and on the way
|
|
9
|
+
out. Cost intelligence (tiered routing, graceful budgets, per-task
|
|
10
|
+
attribution) rides along via [tokenecon](https://pypi.org/project/tokenecon/),
|
|
11
|
+
but security is the product.
|
|
12
|
+
|
|
13
|
+
## The security layer
|
|
14
|
+
|
|
15
|
+
```
|
|
16
|
+
request -> auth -> INPUT policy -> budget -> route -> model
|
|
17
|
+
-> OUTPUT policy -> audit -> response
|
|
18
|
+
```
|
|
19
|
+
|
|
20
|
+
**Pluggable detectors** (each scans, the policy decides):
|
|
21
|
+
|
|
22
|
+
| Detector | Catches | Default action |
|
|
23
|
+
|---|---|---|
|
|
24
|
+
| Secrets | AWS keys, private keys, GitHub/OpenAI tokens, `password = ...` | **block** — never reaches a model |
|
|
25
|
+
| PII | emails, phones, SSNs, credit cards, IPs | **redact** — `[REDACTED:EMAIL]`, request continues |
|
|
26
|
+
| Prompt injection | "ignore previous instructions", role overrides, system-prompt probes, known jailbreaks | **flag** — allowed but marked (one flag away from block) |
|
|
27
|
+
| Blocklist | your own regexes | **block** |
|
|
28
|
+
|
|
29
|
+
**Policy engine:** per-direction rules (`in` / `out` / `both`) with four
|
|
30
|
+
actions — `allow`, `redact`, `block`, `flag`. Add an ML-based detector later
|
|
31
|
+
by subclassing `Detector`; no other code changes.
|
|
32
|
+
|
|
33
|
+
**Audit log:** every decision recorded — timestamp, key, task, detector,
|
|
34
|
+
action, redacted snippet. Raw PII and secrets never touch the audit trail.
|
|
35
|
+
`GET /admin/audit` for compliance review.
|
|
36
|
+
|
|
37
|
+
## Quickstart
|
|
38
|
+
|
|
39
|
+
```bash
|
|
40
|
+
pip install argus-gateway
|
|
41
|
+
export AIGW_MASTER_KEY="sk-admin-secret"
|
|
42
|
+
export OPENAI_API_KEY="..."
|
|
43
|
+
argus # serves on 127.0.0.1:4000
|
|
44
|
+
```
|
|
45
|
+
|
|
46
|
+
Point any OpenAI client at it — one `base_url` change:
|
|
47
|
+
|
|
48
|
+
```python
|
|
49
|
+
import openai
|
|
50
|
+
client = openai.OpenAI(api_key="<virtual-key>", base_url="http://localhost:4000/v1")
|
|
51
|
+
```
|
|
52
|
+
|
|
53
|
+
Self-hosted models sit behind the same layer:
|
|
54
|
+
|
|
55
|
+
```python
|
|
56
|
+
from gateway.providers.generic import GenericProvider
|
|
57
|
+
# vLLM, Ollama, llama.cpp — any OpenAI-shaped endpoint
|
|
58
|
+
```
|
|
59
|
+
|
|
60
|
+
Every response carries what happened:
|
|
61
|
+
|
|
62
|
+
```json
|
|
63
|
+
"gateway": {
|
|
64
|
+
"provider": "openai", "tier": "small", "difficulty": 0.17,
|
|
65
|
+
"cost_usd": 0.000004, "degraded": false,
|
|
66
|
+
"security": {"redactions": ["pii:EMAIL"], "flags": ["prompt_injection:INSTRUCTION_OVERRIDE"]}
|
|
67
|
+
}
|
|
68
|
+
```
|
|
69
|
+
|
|
70
|
+
## Also inside: cost intelligence
|
|
71
|
+
|
|
72
|
+
- **Tiered routing** — requests scored for difficulty, cheapest capable tier wins
|
|
73
|
+
- **Graceful budgets** — over budget? Falls back to the cheapest tier and flags
|
|
74
|
+
`degraded: true` instead of failing (hard-stop mode available)
|
|
75
|
+
- **Per-task attribution** — `X-Task-Id` groups an agentic run's turns into one
|
|
76
|
+
costed task: `GET /admin/tasks/{id}`
|
|
77
|
+
|
|
78
|
+
## Admin API (`X-Admin-Key` required)
|
|
79
|
+
|
|
80
|
+
- `POST /admin/keys` — issue virtual keys with budgets (`degrade`/`hard`, daily/monthly)
|
|
81
|
+
- `GET /admin/spend`, `GET /admin/keys/{id}/spend` — spend tracking
|
|
82
|
+
- `GET /admin/tasks/{task_id}` — per-task cost breakdown
|
|
83
|
+
- `GET /admin/audit?limit=&key_id=` — security decisions
|
|
84
|
+
|
|
85
|
+
## Roadmap
|
|
86
|
+
|
|
87
|
+
Streaming (SSE), Bedrock/Gemini/Mistral adapters, semantic caching,
|
|
88
|
+
OTEL/Prometheus export, admin UI, SSO/RBAC, ML-based injection detector.
|
|
89
|
+
|
|
90
|
+
## License
|
|
91
|
+
|
|
92
|
+
MIT — Karmendra Pandey.
|
|
@@ -0,0 +1,113 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: argus-gateway
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Argus: the security layer for all AI models — an OpenAI-compatible AI gateway with pluggable threat detectors, policy enforcement, audit, and cost intelligence
|
|
5
|
+
Author: Karmendra Pandey
|
|
6
|
+
License: MIT
|
|
7
|
+
Keywords: llm,ai-gateway,proxy,cost,routing,agents
|
|
8
|
+
Classifier: Programming Language :: Python :: 3
|
|
9
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
10
|
+
Classifier: Intended Audience :: Developers
|
|
11
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
12
|
+
Requires-Python: >=3.9
|
|
13
|
+
Description-Content-Type: text/markdown
|
|
14
|
+
Requires-Dist: fastapi>=0.110
|
|
15
|
+
Requires-Dist: uvicorn>=0.29
|
|
16
|
+
Requires-Dist: httpx>=0.27
|
|
17
|
+
Requires-Dist: tokenecon>=0.2.0
|
|
18
|
+
Provides-Extra: test
|
|
19
|
+
Requires-Dist: pytest>=8; extra == "test"
|
|
20
|
+
Requires-Dist: httpx>=0.27; extra == "test"
|
|
21
|
+
|
|
22
|
+
# Argus — the security layer for all AI models
|
|
23
|
+
|
|
24
|
+
An OpenAI-compatible gateway that sits in front of **any** LLM — OpenAI,
|
|
25
|
+
Anthropic, Bedrock, Gemini, Mistral, vLLM, Ollama, anything with an
|
|
26
|
+
OpenAI-shaped endpoint — and enforces security policy before a model ever
|
|
27
|
+
sees your data.
|
|
28
|
+
|
|
29
|
+
Every request passes through the policy engine on the way in and on the way
|
|
30
|
+
out. Cost intelligence (tiered routing, graceful budgets, per-task
|
|
31
|
+
attribution) rides along via [tokenecon](https://pypi.org/project/tokenecon/),
|
|
32
|
+
but security is the product.
|
|
33
|
+
|
|
34
|
+
## The security layer
|
|
35
|
+
|
|
36
|
+
```
|
|
37
|
+
request -> auth -> INPUT policy -> budget -> route -> model
|
|
38
|
+
-> OUTPUT policy -> audit -> response
|
|
39
|
+
```
|
|
40
|
+
|
|
41
|
+
**Pluggable detectors** (each scans, the policy decides):
|
|
42
|
+
|
|
43
|
+
| Detector | Catches | Default action |
|
|
44
|
+
|---|---|---|
|
|
45
|
+
| Secrets | AWS keys, private keys, GitHub/OpenAI tokens, `password = ...` | **block** — never reaches a model |
|
|
46
|
+
| PII | emails, phones, SSNs, credit cards, IPs | **redact** — `[REDACTED:EMAIL]`, request continues |
|
|
47
|
+
| Prompt injection | "ignore previous instructions", role overrides, system-prompt probes, known jailbreaks | **flag** — allowed but marked (one flag away from block) |
|
|
48
|
+
| Blocklist | your own regexes | **block** |
|
|
49
|
+
|
|
50
|
+
**Policy engine:** per-direction rules (`in` / `out` / `both`) with four
|
|
51
|
+
actions — `allow`, `redact`, `block`, `flag`. Add an ML-based detector later
|
|
52
|
+
by subclassing `Detector`; no other code changes.
|
|
53
|
+
|
|
54
|
+
**Audit log:** every decision recorded — timestamp, key, task, detector,
|
|
55
|
+
action, redacted snippet. Raw PII and secrets never touch the audit trail.
|
|
56
|
+
`GET /admin/audit` for compliance review.
|
|
57
|
+
|
|
58
|
+
## Quickstart
|
|
59
|
+
|
|
60
|
+
```bash
|
|
61
|
+
pip install argus-gateway
|
|
62
|
+
export AIGW_MASTER_KEY="sk-admin-secret"
|
|
63
|
+
export OPENAI_API_KEY="..."
|
|
64
|
+
argus # serves on 127.0.0.1:4000
|
|
65
|
+
```
|
|
66
|
+
|
|
67
|
+
Point any OpenAI client at it — one `base_url` change:
|
|
68
|
+
|
|
69
|
+
```python
|
|
70
|
+
import openai
|
|
71
|
+
client = openai.OpenAI(api_key="<virtual-key>", base_url="http://localhost:4000/v1")
|
|
72
|
+
```
|
|
73
|
+
|
|
74
|
+
Self-hosted models sit behind the same layer:
|
|
75
|
+
|
|
76
|
+
```python
|
|
77
|
+
from gateway.providers.generic import GenericProvider
|
|
78
|
+
# vLLM, Ollama, llama.cpp — any OpenAI-shaped endpoint
|
|
79
|
+
```
|
|
80
|
+
|
|
81
|
+
Every response carries what happened:
|
|
82
|
+
|
|
83
|
+
```json
|
|
84
|
+
"gateway": {
|
|
85
|
+
"provider": "openai", "tier": "small", "difficulty": 0.17,
|
|
86
|
+
"cost_usd": 0.000004, "degraded": false,
|
|
87
|
+
"security": {"redactions": ["pii:EMAIL"], "flags": ["prompt_injection:INSTRUCTION_OVERRIDE"]}
|
|
88
|
+
}
|
|
89
|
+
```
|
|
90
|
+
|
|
91
|
+
## Also inside: cost intelligence
|
|
92
|
+
|
|
93
|
+
- **Tiered routing** — requests scored for difficulty, cheapest capable tier wins
|
|
94
|
+
- **Graceful budgets** — over budget? Falls back to the cheapest tier and flags
|
|
95
|
+
`degraded: true` instead of failing (hard-stop mode available)
|
|
96
|
+
- **Per-task attribution** — `X-Task-Id` groups an agentic run's turns into one
|
|
97
|
+
costed task: `GET /admin/tasks/{id}`
|
|
98
|
+
|
|
99
|
+
## Admin API (`X-Admin-Key` required)
|
|
100
|
+
|
|
101
|
+
- `POST /admin/keys` — issue virtual keys with budgets (`degrade`/`hard`, daily/monthly)
|
|
102
|
+
- `GET /admin/spend`, `GET /admin/keys/{id}/spend` — spend tracking
|
|
103
|
+
- `GET /admin/tasks/{task_id}` — per-task cost breakdown
|
|
104
|
+
- `GET /admin/audit?limit=&key_id=` — security decisions
|
|
105
|
+
|
|
106
|
+
## Roadmap
|
|
107
|
+
|
|
108
|
+
Streaming (SSE), Bedrock/Gemini/Mistral adapters, semantic caching,
|
|
109
|
+
OTEL/Prometheus export, admin UI, SSO/RBAC, ML-based injection detector.
|
|
110
|
+
|
|
111
|
+
## License
|
|
112
|
+
|
|
113
|
+
MIT — Karmendra Pandey.
|
|
@@ -0,0 +1,27 @@
|
|
|
1
|
+
README.md
|
|
2
|
+
pyproject.toml
|
|
3
|
+
argus_gateway.egg-info/PKG-INFO
|
|
4
|
+
argus_gateway.egg-info/SOURCES.txt
|
|
5
|
+
argus_gateway.egg-info/dependency_links.txt
|
|
6
|
+
argus_gateway.egg-info/entry_points.txt
|
|
7
|
+
argus_gateway.egg-info/requires.txt
|
|
8
|
+
argus_gateway.egg-info/top_level.txt
|
|
9
|
+
gateway/__init__.py
|
|
10
|
+
gateway/auth.py
|
|
11
|
+
gateway/config.py
|
|
12
|
+
gateway/ledger.py
|
|
13
|
+
gateway/router.py
|
|
14
|
+
gateway/server.py
|
|
15
|
+
gateway/providers/__init__.py
|
|
16
|
+
gateway/providers/anthropic.py
|
|
17
|
+
gateway/providers/base.py
|
|
18
|
+
gateway/providers/generic.py
|
|
19
|
+
gateway/providers/mock.py
|
|
20
|
+
gateway/providers/openai.py
|
|
21
|
+
gateway/security/__init__.py
|
|
22
|
+
gateway/security/audit.py
|
|
23
|
+
gateway/security/detectors.py
|
|
24
|
+
gateway/security/policy.py
|
|
25
|
+
tests/test_router.py
|
|
26
|
+
tests/test_security.py
|
|
27
|
+
tests/test_server.py
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
gateway
|
|
@@ -0,0 +1,108 @@
|
|
|
1
|
+
"""Virtual API keys: issuance, verification, and budget state.
|
|
2
|
+
|
|
3
|
+
Keys are stored as SHA-256 hashes (the raw key is shown once at creation).
|
|
4
|
+
Budgets are enforced per key; spend is read from the ledger.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
from __future__ import annotations
|
|
8
|
+
|
|
9
|
+
import hashlib
|
|
10
|
+
import secrets
|
|
11
|
+
import sqlite3
|
|
12
|
+
import time
|
|
13
|
+
from dataclasses import dataclass
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
@dataclass
|
|
17
|
+
class KeyRecord:
|
|
18
|
+
id: int
|
|
19
|
+
name: str
|
|
20
|
+
key_hash: str
|
|
21
|
+
budget_usd: float
|
|
22
|
+
period: str # "daily" | "monthly"
|
|
23
|
+
mode: str # "degrade" | "hard"
|
|
24
|
+
created_at: float
|
|
25
|
+
revoked: bool = False
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
_SCHEMA = """
|
|
29
|
+
CREATE TABLE IF NOT EXISTS keys (
|
|
30
|
+
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
31
|
+
name TEXT NOT NULL,
|
|
32
|
+
key_hash TEXT NOT NULL UNIQUE,
|
|
33
|
+
budget_usd REAL NOT NULL,
|
|
34
|
+
period TEXT NOT NULL DEFAULT 'monthly',
|
|
35
|
+
mode TEXT NOT NULL DEFAULT 'degrade',
|
|
36
|
+
created_at REAL NOT NULL,
|
|
37
|
+
revoked INTEGER NOT NULL DEFAULT 0
|
|
38
|
+
);
|
|
39
|
+
"""
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
class KeyStore:
|
|
43
|
+
def __init__(self, db_path: str):
|
|
44
|
+
self.db_path = db_path
|
|
45
|
+
with sqlite3.connect(db_path) as c:
|
|
46
|
+
c.executescript(_SCHEMA)
|
|
47
|
+
|
|
48
|
+
@staticmethod
|
|
49
|
+
def _hash(raw: str) -> str:
|
|
50
|
+
return hashlib.sha256(raw.encode()).hexdigest()
|
|
51
|
+
|
|
52
|
+
def create(
|
|
53
|
+
self,
|
|
54
|
+
name: str,
|
|
55
|
+
budget_usd: float = 10.0,
|
|
56
|
+
period: str = "monthly",
|
|
57
|
+
mode: str = "degrade",
|
|
58
|
+
) -> tuple[KeyRecord, str]:
|
|
59
|
+
"""Create a key; returns (record, raw_key). Raw key is shown once."""
|
|
60
|
+
raw = "sk-gw-" + secrets.token_urlsafe(32)
|
|
61
|
+
rec = KeyRecord(
|
|
62
|
+
id=0,
|
|
63
|
+
name=name,
|
|
64
|
+
key_hash=self._hash(raw),
|
|
65
|
+
budget_usd=budget_usd,
|
|
66
|
+
period=period,
|
|
67
|
+
mode=mode,
|
|
68
|
+
created_at=time.time(),
|
|
69
|
+
)
|
|
70
|
+
with sqlite3.connect(self.db_path) as c:
|
|
71
|
+
cur = c.execute(
|
|
72
|
+
"INSERT INTO keys (name, key_hash, budget_usd, period, mode, created_at)"
|
|
73
|
+
" VALUES (?, ?, ?, ?, ?, ?)",
|
|
74
|
+
(rec.name, rec.key_hash, rec.budget_usd, rec.period, rec.mode, rec.created_at),
|
|
75
|
+
)
|
|
76
|
+
rec.id = cur.lastrowid
|
|
77
|
+
return rec, raw
|
|
78
|
+
|
|
79
|
+
def verify(self, raw: str) -> KeyRecord | None:
|
|
80
|
+
"""Return the key record for a valid, non-revoked key, else None."""
|
|
81
|
+
with sqlite3.connect(self.db_path) as c:
|
|
82
|
+
row = c.execute(
|
|
83
|
+
"SELECT id, name, key_hash, budget_usd, period, mode, created_at, revoked"
|
|
84
|
+
" FROM keys WHERE key_hash = ?",
|
|
85
|
+
(self._hash(raw),),
|
|
86
|
+
).fetchone()
|
|
87
|
+
if not row or row[7]:
|
|
88
|
+
return None
|
|
89
|
+
return KeyRecord(
|
|
90
|
+
id=row[0], name=row[1], key_hash=row[2], budget_usd=row[3],
|
|
91
|
+
period=row[4], mode=row[5], created_at=row[6], revoked=bool(row[7]),
|
|
92
|
+
)
|
|
93
|
+
|
|
94
|
+
def revoke(self, key_id: int) -> None:
|
|
95
|
+
with sqlite3.connect(self.db_path) as c:
|
|
96
|
+
c.execute("UPDATE keys SET revoked = 1 WHERE id = ?", (key_id,))
|
|
97
|
+
|
|
98
|
+
def list(self) -> list[KeyRecord]:
|
|
99
|
+
with sqlite3.connect(self.db_path) as c:
|
|
100
|
+
rows = c.execute(
|
|
101
|
+
"SELECT id, name, key_hash, budget_usd, period, mode, created_at, revoked"
|
|
102
|
+
" FROM keys ORDER BY id"
|
|
103
|
+
).fetchall()
|
|
104
|
+
return [
|
|
105
|
+
KeyRecord(id=r[0], name=r[1], key_hash=r[2], budget_usd=r[3], period=r[4],
|
|
106
|
+
mode=r[5], created_at=r[6], revoked=bool(r[7]))
|
|
107
|
+
for r in rows
|
|
108
|
+
]
|
|
@@ -0,0 +1,52 @@
|
|
|
1
|
+
"""Gateway configuration: providers, tiers, and runtime settings.
|
|
2
|
+
|
|
3
|
+
Providers are declared in code (or loaded from a dict); API keys come from
|
|
4
|
+
the environment. Prices are tokenecon ModelTiers so routing and costing share
|
|
5
|
+
one source of truth.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
from __future__ import annotations
|
|
9
|
+
|
|
10
|
+
import os
|
|
11
|
+
from dataclasses import dataclass, field
|
|
12
|
+
from typing import Dict, List
|
|
13
|
+
|
|
14
|
+
from tokenecon import ModelTier
|
|
15
|
+
|
|
16
|
+
from .providers.base import Provider
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
@dataclass
|
|
20
|
+
class ProviderEntry:
|
|
21
|
+
"""One backend behind the gateway."""
|
|
22
|
+
|
|
23
|
+
provider: Provider
|
|
24
|
+
tier: ModelTier
|
|
25
|
+
models: List[str] # model names this backend serves, e.g. ["gpt-4o-mini"]
|
|
26
|
+
enabled: bool = True
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
@dataclass
|
|
30
|
+
class GatewayConfig:
|
|
31
|
+
db_path: str = os.environ.get("AIGW_DB", "aigw.db")
|
|
32
|
+
master_key: str = os.environ.get("AIGW_MASTER_KEY", "sk-admin-changeme")
|
|
33
|
+
default_budget_usd: float = float(os.environ.get("AIGW_DEFAULT_BUDGET", "10.0"))
|
|
34
|
+
default_budget_mode: str = os.environ.get("AIGW_BUDGET_MODE", "degrade") # degrade | hard
|
|
35
|
+
safety_margin: float = 0.1 # tier-selection margin, mirrors tokenecon router
|
|
36
|
+
providers: List[ProviderEntry] = field(default_factory=list)
|
|
37
|
+
|
|
38
|
+
def models(self) -> List[Dict[str, str]]:
|
|
39
|
+
"""Model catalog in OpenAI /v1/models shape."""
|
|
40
|
+
out = []
|
|
41
|
+
for entry in self.providers:
|
|
42
|
+
if not entry.enabled:
|
|
43
|
+
continue
|
|
44
|
+
for m in entry.models:
|
|
45
|
+
out.append({"id": m, "object": "model", "owned_by": entry.provider.name})
|
|
46
|
+
return out
|
|
47
|
+
|
|
48
|
+
def entry_for_model(self, model: str) -> ProviderEntry | None:
|
|
49
|
+
for entry in self.providers:
|
|
50
|
+
if entry.enabled and model in entry.models:
|
|
51
|
+
return entry
|
|
52
|
+
return None
|
|
@@ -0,0 +1,106 @@
|
|
|
1
|
+
"""Cost ledger: every proxied call recorded; spend queryable per key and per task.
|
|
2
|
+
|
|
3
|
+
A task groups the turns of one agentic run (client passes X-Task-Id). This is
|
|
4
|
+
the per-task attribution LiteLLM-style gateways don't do.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
from __future__ import annotations
|
|
8
|
+
|
|
9
|
+
import sqlite3
|
|
10
|
+
import time
|
|
11
|
+
from dataclasses import dataclass
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
@dataclass
|
|
15
|
+
class CallEntry:
|
|
16
|
+
key_id: int
|
|
17
|
+
task_id: str | None
|
|
18
|
+
provider: str
|
|
19
|
+
model: str
|
|
20
|
+
tier: str
|
|
21
|
+
input_tokens: int
|
|
22
|
+
output_tokens: int
|
|
23
|
+
cost_usd: float
|
|
24
|
+
degraded: bool = False
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
_SCHEMA = """
|
|
28
|
+
CREATE TABLE IF NOT EXISTS calls (
|
|
29
|
+
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
30
|
+
ts REAL NOT NULL,
|
|
31
|
+
key_id INTEGER NOT NULL,
|
|
32
|
+
task_id TEXT,
|
|
33
|
+
provider TEXT NOT NULL,
|
|
34
|
+
model TEXT NOT NULL,
|
|
35
|
+
tier TEXT NOT NULL,
|
|
36
|
+
input_tokens INTEGER NOT NULL,
|
|
37
|
+
output_tokens INTEGER NOT NULL,
|
|
38
|
+
cost_usd REAL NOT NULL,
|
|
39
|
+
degraded INTEGER NOT NULL DEFAULT 0
|
|
40
|
+
);
|
|
41
|
+
CREATE INDEX IF NOT EXISTS idx_calls_key_ts ON calls (key_id, ts);
|
|
42
|
+
CREATE INDEX IF NOT EXISTS idx_calls_task ON calls (task_id);
|
|
43
|
+
"""
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
class Ledger:
|
|
47
|
+
def __init__(self, db_path: str):
|
|
48
|
+
self.db_path = db_path
|
|
49
|
+
with sqlite3.connect(db_path) as c:
|
|
50
|
+
c.executescript(_SCHEMA)
|
|
51
|
+
|
|
52
|
+
def record(self, entry: CallEntry) -> int:
|
|
53
|
+
with sqlite3.connect(self.db_path) as c:
|
|
54
|
+
cur = c.execute(
|
|
55
|
+
"INSERT INTO calls (ts, key_id, task_id, provider, model, tier,"
|
|
56
|
+
" input_tokens, output_tokens, cost_usd, degraded)"
|
|
57
|
+
" VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
|
|
58
|
+
(time.time(), entry.key_id, entry.task_id, entry.provider,
|
|
59
|
+
entry.model, entry.tier, entry.input_tokens, entry.output_tokens,
|
|
60
|
+
entry.cost_usd, int(entry.degraded)),
|
|
61
|
+
)
|
|
62
|
+
return cur.lastrowid
|
|
63
|
+
|
|
64
|
+
def _period_start(self, period: str) -> float:
|
|
65
|
+
now = time.time()
|
|
66
|
+
if period == "daily":
|
|
67
|
+
return now - 86400
|
|
68
|
+
return now - 30 * 86400 # monthly ~= 30d
|
|
69
|
+
|
|
70
|
+
def key_spend(self, key_id: int, period: str = "monthly") -> float:
|
|
71
|
+
with sqlite3.connect(self.db_path) as c:
|
|
72
|
+
row = c.execute(
|
|
73
|
+
"SELECT COALESCE(SUM(cost_usd), 0) FROM calls WHERE key_id = ? AND ts >= ?",
|
|
74
|
+
(key_id, self._period_start(period)),
|
|
75
|
+
).fetchone()
|
|
76
|
+
return float(row[0])
|
|
77
|
+
|
|
78
|
+
def task_spend(self, task_id: str) -> dict:
|
|
79
|
+
"""Totals + per-call breakdown for one task id."""
|
|
80
|
+
with sqlite3.connect(self.db_path) as c:
|
|
81
|
+
rows = c.execute(
|
|
82
|
+
"SELECT ts, provider, model, tier, input_tokens, output_tokens,"
|
|
83
|
+
" cost_usd, degraded FROM calls WHERE task_id = ? ORDER BY ts",
|
|
84
|
+
(task_id,),
|
|
85
|
+
).fetchall()
|
|
86
|
+
calls = [
|
|
87
|
+
{"ts": r[0], "provider": r[1], "model": r[2], "tier": r[3],
|
|
88
|
+
"input_tokens": r[4], "output_tokens": r[5],
|
|
89
|
+
"cost_usd": r[6], "degraded": bool(r[7])}
|
|
90
|
+
for r in rows
|
|
91
|
+
]
|
|
92
|
+
return {
|
|
93
|
+
"task_id": task_id,
|
|
94
|
+
"calls": len(calls),
|
|
95
|
+
"total_cost_usd": round(sum(c["cost_usd"] for c in calls), 6),
|
|
96
|
+
"total_tokens": sum(c["input_tokens"] + c["output_tokens"] for c in calls),
|
|
97
|
+
"breakdown": calls,
|
|
98
|
+
}
|
|
99
|
+
|
|
100
|
+
def total_spend(self) -> dict:
|
|
101
|
+
with sqlite3.connect(self.db_path) as c:
|
|
102
|
+
row = c.execute(
|
|
103
|
+
"SELECT COUNT(*), COALESCE(SUM(cost_usd), 0),"
|
|
104
|
+
" COALESCE(SUM(input_tokens + output_tokens), 0) FROM calls"
|
|
105
|
+
).fetchone()
|
|
106
|
+
return {"calls": row[0], "total_cost_usd": round(row[1], 6), "total_tokens": row[2]}
|