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.
Files changed (29) hide show
  1. argus_gateway-0.1.0/PKG-INFO +113 -0
  2. argus_gateway-0.1.0/README.md +92 -0
  3. argus_gateway-0.1.0/argus_gateway.egg-info/PKG-INFO +113 -0
  4. argus_gateway-0.1.0/argus_gateway.egg-info/SOURCES.txt +27 -0
  5. argus_gateway-0.1.0/argus_gateway.egg-info/dependency_links.txt +1 -0
  6. argus_gateway-0.1.0/argus_gateway.egg-info/entry_points.txt +2 -0
  7. argus_gateway-0.1.0/argus_gateway.egg-info/requires.txt +8 -0
  8. argus_gateway-0.1.0/argus_gateway.egg-info/top_level.txt +1 -0
  9. argus_gateway-0.1.0/gateway/__init__.py +7 -0
  10. argus_gateway-0.1.0/gateway/auth.py +108 -0
  11. argus_gateway-0.1.0/gateway/config.py +52 -0
  12. argus_gateway-0.1.0/gateway/ledger.py +106 -0
  13. argus_gateway-0.1.0/gateway/providers/__init__.py +5 -0
  14. argus_gateway-0.1.0/gateway/providers/anthropic.py +85 -0
  15. argus_gateway-0.1.0/gateway/providers/base.py +55 -0
  16. argus_gateway-0.1.0/gateway/providers/generic.py +29 -0
  17. argus_gateway-0.1.0/gateway/providers/mock.py +28 -0
  18. argus_gateway-0.1.0/gateway/providers/openai.py +65 -0
  19. argus_gateway-0.1.0/gateway/router.py +62 -0
  20. argus_gateway-0.1.0/gateway/security/__init__.py +29 -0
  21. argus_gateway-0.1.0/gateway/security/audit.py +72 -0
  22. argus_gateway-0.1.0/gateway/security/detectors.py +110 -0
  23. argus_gateway-0.1.0/gateway/security/policy.py +91 -0
  24. argus_gateway-0.1.0/gateway/server.py +222 -0
  25. argus_gateway-0.1.0/pyproject.toml +35 -0
  26. argus_gateway-0.1.0/setup.cfg +4 -0
  27. argus_gateway-0.1.0/tests/test_router.py +51 -0
  28. argus_gateway-0.1.0/tests/test_security.py +79 -0
  29. 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,2 @@
1
+ [console_scripts]
2
+ argus = gateway.server:run
@@ -0,0 +1,8 @@
1
+ fastapi>=0.110
2
+ uvicorn>=0.29
3
+ httpx>=0.27
4
+ tokenecon>=0.2.0
5
+
6
+ [test]
7
+ pytest>=8
8
+ httpx>=0.27
@@ -0,0 +1,7 @@
1
+ """Argus: the security layer for all AI models.
2
+
3
+ An OpenAI-compatible AI gateway with pluggable threat detectors, policy
4
+ enforcement, audit logging, and cost intelligence.
5
+ """
6
+
7
+ __version__ = "0.1.0"
@@ -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]}
@@ -0,0 +1,5 @@
1
+ """Provider adapters: one normalized interface per LLM provider."""
2
+
3
+ from .base import ChatMessage, ChatRequest, ChatResponse, Provider
4
+
5
+ __all__ = ["ChatMessage", "ChatRequest", "ChatResponse", "Provider"]