capsule-litellm 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.
- capsule_litellm-0.1.0/LICENSE +32 -0
- capsule_litellm-0.1.0/PKG-INFO +100 -0
- capsule_litellm-0.1.0/README.md +51 -0
- capsule_litellm-0.1.0/pyproject.toml +37 -0
- capsule_litellm-0.1.0/src/capsule_litellm_plugin/__init__.py +21 -0
- capsule_litellm-0.1.0/src/capsule_litellm_plugin/callback.py +213 -0
- capsule_litellm-0.1.0/src/capsule_litellm_plugin/config.py +53 -0
- capsule_litellm-0.1.0/src/capsule_litellm_plugin/transport.py +59 -0
- capsule_litellm-0.1.0/src/capsule_litellm_plugin/types.py +68 -0
|
@@ -0,0 +1,32 @@
|
|
|
1
|
+
Capsule Security Commercial License
|
|
2
|
+
|
|
3
|
+
Copyright (c) Capsule Security, Inc. All rights reserved.
|
|
4
|
+
|
|
5
|
+
This software is licensed, not sold, and is made available solely to customers
|
|
6
|
+
of Capsule Security, Inc. ("Capsule") under the terms of an active commercial
|
|
7
|
+
agreement (the "Subscription Agreement") between the customer and Capsule.
|
|
8
|
+
|
|
9
|
+
Subject to the Subscription Agreement, Capsule grants the customer a
|
|
10
|
+
non-exclusive, non-transferable, non-sublicensable, revocable license to install
|
|
11
|
+
and run this software in connection with the customer's authorized use of the
|
|
12
|
+
Capsule platform, solely for the customer's internal business operations.
|
|
13
|
+
|
|
14
|
+
Except as expressly permitted by the Subscription Agreement, you may not:
|
|
15
|
+
(a) copy, modify, or create derivative works of the software;
|
|
16
|
+
(b) reverse engineer, decompile, or disassemble the software;
|
|
17
|
+
(c) redistribute, sublicense, sell, rent, lease, or otherwise transfer the
|
|
18
|
+
software to any third party;
|
|
19
|
+
(d) remove or alter any proprietary notices contained in the software; or
|
|
20
|
+
(e) use the software except in connection with an active Subscription
|
|
21
|
+
Agreement.
|
|
22
|
+
|
|
23
|
+
Without a valid Subscription Agreement with Capsule, no rights are granted to
|
|
24
|
+
use this software. Unauthorized use is prohibited.
|
|
25
|
+
|
|
26
|
+
THE SOFTWARE IS PROVIDED "AS IS" AND CAPSULE DISCLAIMS ALL WARRANTIES, EXPRESS
|
|
27
|
+
OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY WARRANTIES OF MERCHANTABILITY,
|
|
28
|
+
FITNESS FOR A PARTICULAR PURPOSE, AND NON-INFRINGEMENT. CAPSULE SHALL NOT BE
|
|
29
|
+
LIABLE FOR ANY DAMAGES ARISING FROM OR RELATING TO USE OF THE SOFTWARE EXCEPT
|
|
30
|
+
AS EXPRESSLY SET FORTH IN THE SUBSCRIPTION AGREEMENT.
|
|
31
|
+
|
|
32
|
+
For licensing inquiries, contact: sales@capsulesecurity.io
|
|
@@ -0,0 +1,100 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: capsule-litellm
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Capsule plugin for the LiteLLM gateway — forwards completion requests to the Capsule agentsecurity service for centralized policy enforcement and audit.
|
|
5
|
+
Project-URL: Homepage, https://github.com/capsulesecurity
|
|
6
|
+
Project-URL: Documentation, https://docs.capsulesecurity.io/guides/litellm
|
|
7
|
+
Author: Capsule Security
|
|
8
|
+
License: Capsule Security Commercial License
|
|
9
|
+
|
|
10
|
+
Copyright (c) Capsule Security, Inc. All rights reserved.
|
|
11
|
+
|
|
12
|
+
This software is licensed, not sold, and is made available solely to customers
|
|
13
|
+
of Capsule Security, Inc. ("Capsule") under the terms of an active commercial
|
|
14
|
+
agreement (the "Subscription Agreement") between the customer and Capsule.
|
|
15
|
+
|
|
16
|
+
Subject to the Subscription Agreement, Capsule grants the customer a
|
|
17
|
+
non-exclusive, non-transferable, non-sublicensable, revocable license to install
|
|
18
|
+
and run this software in connection with the customer's authorized use of the
|
|
19
|
+
Capsule platform, solely for the customer's internal business operations.
|
|
20
|
+
|
|
21
|
+
Except as expressly permitted by the Subscription Agreement, you may not:
|
|
22
|
+
(a) copy, modify, or create derivative works of the software;
|
|
23
|
+
(b) reverse engineer, decompile, or disassemble the software;
|
|
24
|
+
(c) redistribute, sublicense, sell, rent, lease, or otherwise transfer the
|
|
25
|
+
software to any third party;
|
|
26
|
+
(d) remove or alter any proprietary notices contained in the software; or
|
|
27
|
+
(e) use the software except in connection with an active Subscription
|
|
28
|
+
Agreement.
|
|
29
|
+
|
|
30
|
+
Without a valid Subscription Agreement with Capsule, no rights are granted to
|
|
31
|
+
use this software. Unauthorized use is prohibited.
|
|
32
|
+
|
|
33
|
+
THE SOFTWARE IS PROVIDED "AS IS" AND CAPSULE DISCLAIMS ALL WARRANTIES, EXPRESS
|
|
34
|
+
OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY WARRANTIES OF MERCHANTABILITY,
|
|
35
|
+
FITNESS FOR A PARTICULAR PURPOSE, AND NON-INFRINGEMENT. CAPSULE SHALL NOT BE
|
|
36
|
+
LIABLE FOR ANY DAMAGES ARISING FROM OR RELATING TO USE OF THE SOFTWARE EXCEPT
|
|
37
|
+
AS EXPRESSLY SET FORTH IN THE SUBSCRIPTION AGREEMENT.
|
|
38
|
+
|
|
39
|
+
For licensing inquiries, contact: sales@capsulesecurity.io
|
|
40
|
+
License-File: LICENSE
|
|
41
|
+
Keywords: agent-security,capsule,litellm,litellm-plugin,llm-gateway,policy,security
|
|
42
|
+
Classifier: License :: Other/Proprietary License
|
|
43
|
+
Classifier: Programming Language :: Python :: 3
|
|
44
|
+
Classifier: Programming Language :: Python :: 3 :: Only
|
|
45
|
+
Requires-Python: >=3.10
|
|
46
|
+
Requires-Dist: httpx>=0.27
|
|
47
|
+
Requires-Dist: litellm[proxy]>=1.40
|
|
48
|
+
Description-Content-Type: text/markdown
|
|
49
|
+
|
|
50
|
+
# capsule-litellm
|
|
51
|
+
|
|
52
|
+
Capsule plugin for the [LiteLLM](https://docs.litellm.ai) LLM gateway. Forwards completion/chat requests to the Capsule `agentsecurity` service for centralized policy enforcement (pre-call allow/deny) and token/latency audit (post-call telemetry).
|
|
53
|
+
|
|
54
|
+
## Install
|
|
55
|
+
|
|
56
|
+
```bash
|
|
57
|
+
pip install capsule-litellm
|
|
58
|
+
```
|
|
59
|
+
|
|
60
|
+
Install it into the same environment that runs your LiteLLM proxy.
|
|
61
|
+
|
|
62
|
+
## Configure
|
|
63
|
+
|
|
64
|
+
Register the callback in your LiteLLM `config.yaml`:
|
|
65
|
+
|
|
66
|
+
```yaml
|
|
67
|
+
litellm_settings:
|
|
68
|
+
callbacks: ['capsule_litellm_plugin.capsule_callback']
|
|
69
|
+
```
|
|
70
|
+
|
|
71
|
+
The callback reads its configuration from environment variables. Set these where the LiteLLM proxy runs:
|
|
72
|
+
|
|
73
|
+
```bash
|
|
74
|
+
export CAPSULE_LITELLM_ENDPOINT="https://agents.capsule.security/v1/litellm/hooks"
|
|
75
|
+
export CAPSULE_LITELLM_TOKEN="<token from the Capsule portal>"
|
|
76
|
+
```
|
|
77
|
+
|
|
78
|
+
The `endpoint` and `token` values are issued by Capsule when you install the LiteLLM integration in the portal (**Integrations → LiteLLM**). The token is a JWT scoped to your tenant and environment — treat it as a secret.
|
|
79
|
+
|
|
80
|
+
### Optional environment variables
|
|
81
|
+
|
|
82
|
+
| Variable | Default | Description |
|
|
83
|
+
| ------------------------------- | ------- | ----------------------------------------------------------------------------------------------------------------- |
|
|
84
|
+
| `CAPSULE_LITELLM_BLOCK_ON_RISK` | `true` | Apply server `deny` verdicts inline. Set to `false` to run the plugin in observe-only mode. |
|
|
85
|
+
| `CAPSULE_LITELLM_FAIL_OPEN` | `true` | Allow the request through if Capsule is unreachable or returns a non-2xx response. Set to `false` to fail closed. |
|
|
86
|
+
| `CAPSULE_LITELLM_TIMEOUT_MS` | `5000` | Per-request HTTP timeout in milliseconds. |
|
|
87
|
+
|
|
88
|
+
## How it works
|
|
89
|
+
|
|
90
|
+
| Hook | Type | Description |
|
|
91
|
+
| ------------------------- | ----------- | ----------------------------------------------------------------------------------------------------- |
|
|
92
|
+
| `async_pre_call_hook` | Blocking | Posts the model + messages to `events/pre` and awaits an allow/deny verdict before the upstream call. |
|
|
93
|
+
| `async_log_success_event` | Observation | Posts model, latency, and token usage to `events/post` after a successful completion. |
|
|
94
|
+
| `async_log_failure_event` | Observation | Posts model, latency, and the error to `events/post` when the upstream call fails. |
|
|
95
|
+
|
|
96
|
+
When `block_on_risk` is enabled and Capsule returns `decision: "deny"`, the pre-call hook raises a `403` so LiteLLM rejects the request before it reaches the upstream provider.
|
|
97
|
+
|
|
98
|
+
## License
|
|
99
|
+
|
|
100
|
+
Capsule Security Commercial License. See [LICENSE](./LICENSE).
|
|
@@ -0,0 +1,51 @@
|
|
|
1
|
+
# capsule-litellm
|
|
2
|
+
|
|
3
|
+
Capsule plugin for the [LiteLLM](https://docs.litellm.ai) LLM gateway. Forwards completion/chat requests to the Capsule `agentsecurity` service for centralized policy enforcement (pre-call allow/deny) and token/latency audit (post-call telemetry).
|
|
4
|
+
|
|
5
|
+
## Install
|
|
6
|
+
|
|
7
|
+
```bash
|
|
8
|
+
pip install capsule-litellm
|
|
9
|
+
```
|
|
10
|
+
|
|
11
|
+
Install it into the same environment that runs your LiteLLM proxy.
|
|
12
|
+
|
|
13
|
+
## Configure
|
|
14
|
+
|
|
15
|
+
Register the callback in your LiteLLM `config.yaml`:
|
|
16
|
+
|
|
17
|
+
```yaml
|
|
18
|
+
litellm_settings:
|
|
19
|
+
callbacks: ['capsule_litellm_plugin.capsule_callback']
|
|
20
|
+
```
|
|
21
|
+
|
|
22
|
+
The callback reads its configuration from environment variables. Set these where the LiteLLM proxy runs:
|
|
23
|
+
|
|
24
|
+
```bash
|
|
25
|
+
export CAPSULE_LITELLM_ENDPOINT="https://agents.capsule.security/v1/litellm/hooks"
|
|
26
|
+
export CAPSULE_LITELLM_TOKEN="<token from the Capsule portal>"
|
|
27
|
+
```
|
|
28
|
+
|
|
29
|
+
The `endpoint` and `token` values are issued by Capsule when you install the LiteLLM integration in the portal (**Integrations → LiteLLM**). The token is a JWT scoped to your tenant and environment — treat it as a secret.
|
|
30
|
+
|
|
31
|
+
### Optional environment variables
|
|
32
|
+
|
|
33
|
+
| Variable | Default | Description |
|
|
34
|
+
| ------------------------------- | ------- | ----------------------------------------------------------------------------------------------------------------- |
|
|
35
|
+
| `CAPSULE_LITELLM_BLOCK_ON_RISK` | `true` | Apply server `deny` verdicts inline. Set to `false` to run the plugin in observe-only mode. |
|
|
36
|
+
| `CAPSULE_LITELLM_FAIL_OPEN` | `true` | Allow the request through if Capsule is unreachable or returns a non-2xx response. Set to `false` to fail closed. |
|
|
37
|
+
| `CAPSULE_LITELLM_TIMEOUT_MS` | `5000` | Per-request HTTP timeout in milliseconds. |
|
|
38
|
+
|
|
39
|
+
## How it works
|
|
40
|
+
|
|
41
|
+
| Hook | Type | Description |
|
|
42
|
+
| ------------------------- | ----------- | ----------------------------------------------------------------------------------------------------- |
|
|
43
|
+
| `async_pre_call_hook` | Blocking | Posts the model + messages to `events/pre` and awaits an allow/deny verdict before the upstream call. |
|
|
44
|
+
| `async_log_success_event` | Observation | Posts model, latency, and token usage to `events/post` after a successful completion. |
|
|
45
|
+
| `async_log_failure_event` | Observation | Posts model, latency, and the error to `events/post` when the upstream call fails. |
|
|
46
|
+
|
|
47
|
+
When `block_on_risk` is enabled and Capsule returns `decision: "deny"`, the pre-call hook raises a `403` so LiteLLM rejects the request before it reaches the upstream provider.
|
|
48
|
+
|
|
49
|
+
## License
|
|
50
|
+
|
|
51
|
+
Capsule Security Commercial License. See [LICENSE](./LICENSE).
|
|
@@ -0,0 +1,37 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling>=1.21"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "capsule-litellm"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "Capsule plugin for the LiteLLM gateway — forwards completion requests to the Capsule agentsecurity service for centralized policy enforcement and audit."
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.10"
|
|
11
|
+
license = { file = "LICENSE" }
|
|
12
|
+
authors = [{ name = "Capsule Security" }]
|
|
13
|
+
keywords = [
|
|
14
|
+
"litellm",
|
|
15
|
+
"litellm-plugin",
|
|
16
|
+
"llm-gateway",
|
|
17
|
+
"capsule",
|
|
18
|
+
"security",
|
|
19
|
+
"policy",
|
|
20
|
+
"agent-security",
|
|
21
|
+
]
|
|
22
|
+
classifiers = [
|
|
23
|
+
"License :: Other/Proprietary License",
|
|
24
|
+
"Programming Language :: Python :: 3",
|
|
25
|
+
"Programming Language :: Python :: 3 :: Only",
|
|
26
|
+
]
|
|
27
|
+
dependencies = [
|
|
28
|
+
"httpx>=0.27",
|
|
29
|
+
"litellm[proxy]>=1.40",
|
|
30
|
+
]
|
|
31
|
+
|
|
32
|
+
[project.urls]
|
|
33
|
+
Homepage = "https://github.com/capsulesecurity"
|
|
34
|
+
Documentation = "https://docs.capsulesecurity.io/guides/litellm"
|
|
35
|
+
|
|
36
|
+
[tool.hatch.build.targets.wheel]
|
|
37
|
+
packages = ["src/capsule_litellm_plugin"]
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
"""Capsule plugin for the LiteLLM LLM gateway.
|
|
2
|
+
|
|
3
|
+
Forwards completion/chat requests to the Capsule agentsecurity service for
|
|
4
|
+
centralized policy enforcement (pre-call allow/deny) and token/latency audit
|
|
5
|
+
(post-call telemetry).
|
|
6
|
+
|
|
7
|
+
LiteLLM loads this via
|
|
8
|
+
`litellm_settings.callbacks: ["capsule_litellm_plugin.capsule_callback"]`,
|
|
9
|
+
which resolves to the module-level `capsule_callback` instance below.
|
|
10
|
+
"""
|
|
11
|
+
|
|
12
|
+
from __future__ import annotations
|
|
13
|
+
|
|
14
|
+
from .callback import CapsuleLiteLlmCallback
|
|
15
|
+
from .config import CapsuleLiteLlmConfig
|
|
16
|
+
|
|
17
|
+
PLUGIN_VERSION = "0.1.0"
|
|
18
|
+
|
|
19
|
+
capsule_callback = CapsuleLiteLlmCallback()
|
|
20
|
+
|
|
21
|
+
__all__ = ["CapsuleLiteLlmCallback", "CapsuleLiteLlmConfig", "capsule_callback", "PLUGIN_VERSION"]
|
|
@@ -0,0 +1,213 @@
|
|
|
1
|
+
"""LiteLLM custom callback that forwards completion requests to Capsule.
|
|
2
|
+
|
|
3
|
+
`async_pre_call_hook` posts the prompt to Capsule before the upstream model
|
|
4
|
+
call and blocks the request when Capsule returns a deny verdict.
|
|
5
|
+
`async_log_success_event` / `async_log_failure_event` post token and latency
|
|
6
|
+
telemetry once the upstream call settles.
|
|
7
|
+
|
|
8
|
+
LiteLLM resolves the `capsule_litellm_plugin.capsule_callback` config string to
|
|
9
|
+
the module-level instance created in `__init__`.
|
|
10
|
+
"""
|
|
11
|
+
|
|
12
|
+
from __future__ import annotations
|
|
13
|
+
|
|
14
|
+
import logging
|
|
15
|
+
import uuid
|
|
16
|
+
from datetime import datetime
|
|
17
|
+
from typing import Any
|
|
18
|
+
|
|
19
|
+
from fastapi import HTTPException
|
|
20
|
+
from litellm.integrations.custom_logger import CustomLogger
|
|
21
|
+
|
|
22
|
+
from .config import CapsuleLiteLlmConfig
|
|
23
|
+
from .transport import CapsuleResponse, CapsuleTransport
|
|
24
|
+
from .types import LiteLlmChatMessage, LiteLlmPostHookInput, LiteLlmPreHookInput
|
|
25
|
+
|
|
26
|
+
logger = logging.getLogger("capsule_litellm_plugin")
|
|
27
|
+
|
|
28
|
+
_PROCESSED_CALL_TYPES = {
|
|
29
|
+
"completion",
|
|
30
|
+
"acompletion",
|
|
31
|
+
"text_completion",
|
|
32
|
+
"atext_completion",
|
|
33
|
+
"anthropic_messages",
|
|
34
|
+
"aresponses",
|
|
35
|
+
}
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
class CapsuleBlockedError(HTTPException):
|
|
39
|
+
def __init__(self, reason: str) -> None:
|
|
40
|
+
super().__init__(status_code=403, detail=f"Blocked by Capsule policy: {reason}")
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
class CapsuleLiteLlmCallback(CustomLogger):
|
|
44
|
+
def __init__(self) -> None:
|
|
45
|
+
super().__init__()
|
|
46
|
+
try:
|
|
47
|
+
self._config: CapsuleLiteLlmConfig | None = CapsuleLiteLlmConfig.from_env()
|
|
48
|
+
self._transport: CapsuleTransport | None = CapsuleTransport(self._config)
|
|
49
|
+
logger.info(
|
|
50
|
+
"capsule_litellm_plugin: enabled (block_on_risk=%s, fail_open=%s)",
|
|
51
|
+
self._config.block_on_risk,
|
|
52
|
+
self._config.fail_open,
|
|
53
|
+
)
|
|
54
|
+
except ValueError as exc:
|
|
55
|
+
self._config = None
|
|
56
|
+
self._transport = None
|
|
57
|
+
logger.warning("capsule_litellm_plugin: disabled — %s", exc)
|
|
58
|
+
|
|
59
|
+
async def async_pre_call_hook(self, user_api_key_dict: Any, cache: Any, data: dict, call_type: str):
|
|
60
|
+
if self._config is None or self._transport is None:
|
|
61
|
+
return data
|
|
62
|
+
if call_type not in _PROCESSED_CALL_TYPES and not data.get("messages"):
|
|
63
|
+
return data
|
|
64
|
+
|
|
65
|
+
request_id = str(data.get("litellm_call_id") or uuid.uuid4())
|
|
66
|
+
payload = LiteLlmPreHookInput(
|
|
67
|
+
request_id=request_id,
|
|
68
|
+
model=str(data.get("model") or "unknown"),
|
|
69
|
+
messages=_extract_messages(data, call_type),
|
|
70
|
+
user=_optional_str(data.get("user")),
|
|
71
|
+
).to_dict()
|
|
72
|
+
|
|
73
|
+
try:
|
|
74
|
+
response = await self._transport.post_pre(payload, request_id)
|
|
75
|
+
except Exception as exc:
|
|
76
|
+
return self._on_transport_failure(exc, data)
|
|
77
|
+
|
|
78
|
+
if not response.ok:
|
|
79
|
+
logger.warning("capsule_litellm_plugin: pre-hook returned status %s", response.status)
|
|
80
|
+
return self._on_transport_failure(None, data)
|
|
81
|
+
|
|
82
|
+
if _is_deny(response):
|
|
83
|
+
reason = _deny_reason(response)
|
|
84
|
+
logger.info("capsule_litellm_plugin: request %s denied — %s", request_id, reason)
|
|
85
|
+
if self._config.block_on_risk:
|
|
86
|
+
raise CapsuleBlockedError(reason)
|
|
87
|
+
return data
|
|
88
|
+
|
|
89
|
+
async def async_log_success_event(self, kwargs: dict, response_obj: Any, start_time: Any, end_time: Any) -> None:
|
|
90
|
+
await self._record_post(kwargs, response_obj, start_time, end_time, success=True)
|
|
91
|
+
|
|
92
|
+
async def async_log_failure_event(self, kwargs: dict, response_obj: Any, start_time: Any, end_time: Any) -> None:
|
|
93
|
+
# LiteLLM also logs a pre-call block as a failed call, although it never reached the upstream model.
|
|
94
|
+
if isinstance(kwargs.get("exception"), CapsuleBlockedError):
|
|
95
|
+
return
|
|
96
|
+
await self._record_post(kwargs, response_obj, start_time, end_time, success=False)
|
|
97
|
+
|
|
98
|
+
async def _record_post(self, kwargs: dict, response_obj: Any, start_time: Any, end_time: Any, *, success: bool) -> None:
|
|
99
|
+
if self._config is None or self._transport is None:
|
|
100
|
+
return
|
|
101
|
+
request_id = str(kwargs.get("litellm_call_id") or uuid.uuid4())
|
|
102
|
+
usage = _extract_usage(response_obj)
|
|
103
|
+
exception = kwargs.get("exception") if not success else None
|
|
104
|
+
payload = LiteLlmPostHookInput(
|
|
105
|
+
request_id=request_id,
|
|
106
|
+
model=str(kwargs.get("model") or "unknown"),
|
|
107
|
+
duration_ms=_duration_ms(start_time, end_time),
|
|
108
|
+
success=success,
|
|
109
|
+
prompt_tokens=usage.get("prompt_tokens"),
|
|
110
|
+
completion_tokens=usage.get("completion_tokens"),
|
|
111
|
+
total_tokens=usage.get("total_tokens"),
|
|
112
|
+
error_code=type(exception).__name__ if exception is not None else None,
|
|
113
|
+
error_message=str(exception) if exception is not None else None,
|
|
114
|
+
).to_dict()
|
|
115
|
+
try:
|
|
116
|
+
await self._transport.post_post(payload, request_id)
|
|
117
|
+
except Exception as exc:
|
|
118
|
+
logger.warning("capsule_litellm_plugin: post-hook delivery failed — %s", exc)
|
|
119
|
+
|
|
120
|
+
def _on_transport_failure(self, exc: Exception | None, data: dict) -> dict:
|
|
121
|
+
if exc is not None:
|
|
122
|
+
logger.warning("capsule_litellm_plugin: pre-hook request failed — %s", exc)
|
|
123
|
+
assert self._config is not None
|
|
124
|
+
if self._config.fail_open:
|
|
125
|
+
return data
|
|
126
|
+
raise CapsuleBlockedError("Capsule policy service unavailable")
|
|
127
|
+
|
|
128
|
+
|
|
129
|
+
def _extract_messages(data: dict, call_type: str) -> list[LiteLlmChatMessage]:
|
|
130
|
+
# Anthropic /v1/messages sends the system prompt as top-level `system`, the Responses API as `instructions`.
|
|
131
|
+
messages = _system_prompt_messages(data.get("system")) + _system_prompt_messages(data.get("instructions"))
|
|
132
|
+
messages += _conversation_messages(data.get("messages"))
|
|
133
|
+
if call_type == "aresponses":
|
|
134
|
+
messages += _conversation_messages(data.get("input"))
|
|
135
|
+
return messages
|
|
136
|
+
|
|
137
|
+
|
|
138
|
+
def _system_prompt_messages(raw: Any) -> list[LiteLlmChatMessage]:
|
|
139
|
+
content = _flatten_content(raw)
|
|
140
|
+
return [LiteLlmChatMessage(role="system", content=content)] if content.strip() else []
|
|
141
|
+
|
|
142
|
+
|
|
143
|
+
def _conversation_messages(raw: Any) -> list[LiteLlmChatMessage]:
|
|
144
|
+
if isinstance(raw, str):
|
|
145
|
+
return [LiteLlmChatMessage(role="user", content=raw)]
|
|
146
|
+
if not isinstance(raw, list):
|
|
147
|
+
return []
|
|
148
|
+
out: list[LiteLlmChatMessage] = []
|
|
149
|
+
for item in raw:
|
|
150
|
+
if not isinstance(item, dict):
|
|
151
|
+
continue
|
|
152
|
+
role = item.get("role")
|
|
153
|
+
if isinstance(role, str):
|
|
154
|
+
out.append(LiteLlmChatMessage(role=role, content=_flatten_content(item.get("content"))))
|
|
155
|
+
elif item.get("type") == "function_call_output":
|
|
156
|
+
out.append(LiteLlmChatMessage(role="tool", content=_flatten_content(item.get("output"))))
|
|
157
|
+
return out
|
|
158
|
+
|
|
159
|
+
|
|
160
|
+
def _flatten_content(content: Any) -> str:
|
|
161
|
+
if isinstance(content, str):
|
|
162
|
+
return content
|
|
163
|
+
if not isinstance(content, list):
|
|
164
|
+
return ""
|
|
165
|
+
return "\n".join(text for text in (_part_text(part) for part in content) if text)
|
|
166
|
+
|
|
167
|
+
|
|
168
|
+
def _part_text(part: Any) -> str:
|
|
169
|
+
if isinstance(part, str):
|
|
170
|
+
return part
|
|
171
|
+
if not isinstance(part, dict):
|
|
172
|
+
return ""
|
|
173
|
+
text = part.get("text")
|
|
174
|
+
if isinstance(text, str):
|
|
175
|
+
return text
|
|
176
|
+
# Anthropic tool_result blocks nest their own content blocks.
|
|
177
|
+
return _flatten_content(part.get("content"))
|
|
178
|
+
|
|
179
|
+
|
|
180
|
+
def _optional_str(value: Any) -> str | None:
|
|
181
|
+
return value if isinstance(value, str) and value else None
|
|
182
|
+
|
|
183
|
+
|
|
184
|
+
def _extract_usage(response_obj: Any) -> dict[str, int]:
|
|
185
|
+
usage = getattr(response_obj, "usage", None)
|
|
186
|
+
if usage is None and isinstance(response_obj, dict):
|
|
187
|
+
usage = response_obj.get("usage")
|
|
188
|
+
result: dict[str, int] = {}
|
|
189
|
+
for key in ("prompt_tokens", "completion_tokens", "total_tokens"):
|
|
190
|
+
value = getattr(usage, key, None)
|
|
191
|
+
if value is None and isinstance(usage, dict):
|
|
192
|
+
value = usage.get(key)
|
|
193
|
+
if isinstance(value, int):
|
|
194
|
+
result[key] = value
|
|
195
|
+
return result
|
|
196
|
+
|
|
197
|
+
|
|
198
|
+
def _duration_ms(start_time: Any, end_time: Any) -> int:
|
|
199
|
+
if isinstance(start_time, datetime) and isinstance(end_time, datetime):
|
|
200
|
+
return max(0, int((end_time - start_time).total_seconds() * 1000))
|
|
201
|
+
return 0
|
|
202
|
+
|
|
203
|
+
|
|
204
|
+
def _is_deny(response: CapsuleResponse) -> bool:
|
|
205
|
+
return isinstance(response.body, dict) and response.body.get("decision") == "deny"
|
|
206
|
+
|
|
207
|
+
|
|
208
|
+
def _deny_reason(response: CapsuleResponse) -> str:
|
|
209
|
+
if isinstance(response.body, dict):
|
|
210
|
+
reason = response.body.get("reason")
|
|
211
|
+
if isinstance(reason, str) and reason:
|
|
212
|
+
return reason
|
|
213
|
+
return "policy violation"
|
|
@@ -0,0 +1,53 @@
|
|
|
1
|
+
"""Configuration for the Capsule LiteLLM plugin.
|
|
2
|
+
|
|
3
|
+
All settings are read from environment variables so the callback instance can
|
|
4
|
+
be constructed at LiteLLM import time, before the proxy config block is
|
|
5
|
+
available to the callback.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
from __future__ import annotations
|
|
9
|
+
|
|
10
|
+
import os
|
|
11
|
+
from dataclasses import dataclass
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
@dataclass(frozen=True)
|
|
15
|
+
class CapsuleLiteLlmConfig:
|
|
16
|
+
endpoint: str
|
|
17
|
+
token: str
|
|
18
|
+
block_on_risk: bool = True
|
|
19
|
+
fail_open: bool = True
|
|
20
|
+
timeout_ms: int = 5000
|
|
21
|
+
|
|
22
|
+
@classmethod
|
|
23
|
+
def from_env(cls) -> "CapsuleLiteLlmConfig":
|
|
24
|
+
endpoint = os.environ.get("CAPSULE_LITELLM_ENDPOINT")
|
|
25
|
+
token = os.environ.get("CAPSULE_LITELLM_TOKEN")
|
|
26
|
+
if not endpoint:
|
|
27
|
+
raise ValueError("CAPSULE_LITELLM_ENDPOINT is not set")
|
|
28
|
+
if not token:
|
|
29
|
+
raise ValueError("CAPSULE_LITELLM_TOKEN is not set")
|
|
30
|
+
return cls(
|
|
31
|
+
endpoint=endpoint,
|
|
32
|
+
token=token,
|
|
33
|
+
block_on_risk=_env_bool("CAPSULE_LITELLM_BLOCK_ON_RISK", True),
|
|
34
|
+
fail_open=_env_bool("CAPSULE_LITELLM_FAIL_OPEN", True),
|
|
35
|
+
timeout_ms=_env_int("CAPSULE_LITELLM_TIMEOUT_MS", 5000),
|
|
36
|
+
)
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
def _env_bool(key: str, default: bool) -> bool:
|
|
40
|
+
raw = os.environ.get(key)
|
|
41
|
+
if raw is None:
|
|
42
|
+
return default
|
|
43
|
+
return raw.strip().lower() in {"1", "true", "yes", "on"}
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
def _env_int(key: str, default: int) -> int:
|
|
47
|
+
raw = os.environ.get(key)
|
|
48
|
+
if raw is None:
|
|
49
|
+
return default
|
|
50
|
+
try:
|
|
51
|
+
return int(raw)
|
|
52
|
+
except ValueError:
|
|
53
|
+
return default
|
|
@@ -0,0 +1,59 @@
|
|
|
1
|
+
"""Async HTTP transport for the Capsule LiteLLM plugin.
|
|
2
|
+
|
|
3
|
+
LiteLLM invokes callback hooks inside its asyncio event loop, so the transport
|
|
4
|
+
uses `httpx.AsyncClient`. Pre-call requests block on Capsule's allow/deny
|
|
5
|
+
verdict; post-call requests are fire-and-forget telemetry. Non-2xx responses
|
|
6
|
+
are surfaced via `CapsuleResponse.ok` rather than raised — the callback's
|
|
7
|
+
dispatch wrapper decides fail-open / fail-closed.
|
|
8
|
+
"""
|
|
9
|
+
|
|
10
|
+
from __future__ import annotations
|
|
11
|
+
|
|
12
|
+
import json
|
|
13
|
+
from dataclasses import dataclass
|
|
14
|
+
from typing import Any
|
|
15
|
+
|
|
16
|
+
import httpx
|
|
17
|
+
|
|
18
|
+
from .config import CapsuleLiteLlmConfig
|
|
19
|
+
|
|
20
|
+
_HOOKS_PATH_SUFFIX = "/v1/litellm/hooks"
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
@dataclass(frozen=True)
|
|
24
|
+
class CapsuleResponse:
|
|
25
|
+
ok: bool
|
|
26
|
+
status: int
|
|
27
|
+
body: Any
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
class CapsuleTransport:
|
|
31
|
+
def __init__(self, config: CapsuleLiteLlmConfig) -> None:
|
|
32
|
+
self._config = config
|
|
33
|
+
base = config.endpoint.rstrip("/")
|
|
34
|
+
if not base.endswith(_HOOKS_PATH_SUFFIX):
|
|
35
|
+
base = f"{base}{_HOOKS_PATH_SUFFIX}"
|
|
36
|
+
self._base = base
|
|
37
|
+
self._client = httpx.AsyncClient(timeout=config.timeout_ms / 1000.0)
|
|
38
|
+
|
|
39
|
+
async def post_pre(self, payload: dict[str, Any], correlation_id: str) -> CapsuleResponse:
|
|
40
|
+
return await self._post("events/pre", payload, correlation_id)
|
|
41
|
+
|
|
42
|
+
async def post_post(self, payload: dict[str, Any], correlation_id: str) -> CapsuleResponse:
|
|
43
|
+
return await self._post("events/post", payload, correlation_id)
|
|
44
|
+
|
|
45
|
+
async def _post(self, path: str, payload: dict[str, Any], correlation_id: str) -> CapsuleResponse:
|
|
46
|
+
headers = {
|
|
47
|
+
"Content-Type": "application/json",
|
|
48
|
+
"Authorization": f"Bearer {self._config.token}",
|
|
49
|
+
"X-Correlation-Id": correlation_id,
|
|
50
|
+
}
|
|
51
|
+
response = await self._client.post(f"{self._base}/{path}", json=payload, headers=headers)
|
|
52
|
+
body: Any = None
|
|
53
|
+
text = response.text
|
|
54
|
+
if text:
|
|
55
|
+
try:
|
|
56
|
+
body = json.loads(text)
|
|
57
|
+
except json.JSONDecodeError:
|
|
58
|
+
body = text
|
|
59
|
+
return CapsuleResponse(ok=response.is_success, status=response.status_code, body=body)
|
|
@@ -0,0 +1,68 @@
|
|
|
1
|
+
"""Wire-format payloads for the Capsule LiteLLM plugin.
|
|
2
|
+
|
|
3
|
+
Mirror the request DTOs the Capsule agentsecurity service accepts at
|
|
4
|
+
`/v1/litellm/hooks/events/{pre,post}`. Keys are camelCase to match the
|
|
5
|
+
NestJS DTOs on the receiving end.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
from __future__ import annotations
|
|
9
|
+
|
|
10
|
+
from dataclasses import dataclass, field
|
|
11
|
+
from typing import Any
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
@dataclass(frozen=True)
|
|
15
|
+
class LiteLlmChatMessage:
|
|
16
|
+
role: str
|
|
17
|
+
content: str
|
|
18
|
+
|
|
19
|
+
def to_dict(self) -> dict[str, Any]:
|
|
20
|
+
return {"role": self.role, "content": self.content}
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
@dataclass(frozen=True)
|
|
24
|
+
class LiteLlmPreHookInput:
|
|
25
|
+
request_id: str
|
|
26
|
+
model: str
|
|
27
|
+
messages: list[LiteLlmChatMessage] = field(default_factory=list)
|
|
28
|
+
user: str | None = None
|
|
29
|
+
|
|
30
|
+
def to_dict(self) -> dict[str, Any]:
|
|
31
|
+
payload: dict[str, Any] = {"requestId": self.request_id, "model": self.model}
|
|
32
|
+
if self.messages:
|
|
33
|
+
payload["messages"] = [m.to_dict() for m in self.messages]
|
|
34
|
+
if self.user:
|
|
35
|
+
payload["user"] = self.user
|
|
36
|
+
return payload
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
@dataclass(frozen=True)
|
|
40
|
+
class LiteLlmPostHookInput:
|
|
41
|
+
request_id: str
|
|
42
|
+
model: str
|
|
43
|
+
duration_ms: int
|
|
44
|
+
success: bool
|
|
45
|
+
prompt_tokens: int | None = None
|
|
46
|
+
completion_tokens: int | None = None
|
|
47
|
+
total_tokens: int | None = None
|
|
48
|
+
error_code: str | None = None
|
|
49
|
+
error_message: str | None = None
|
|
50
|
+
|
|
51
|
+
def to_dict(self) -> dict[str, Any]:
|
|
52
|
+
payload: dict[str, Any] = {
|
|
53
|
+
"requestId": self.request_id,
|
|
54
|
+
"model": self.model,
|
|
55
|
+
"durationMs": self.duration_ms,
|
|
56
|
+
"success": self.success,
|
|
57
|
+
}
|
|
58
|
+
if self.prompt_tokens is not None:
|
|
59
|
+
payload["promptTokens"] = self.prompt_tokens
|
|
60
|
+
if self.completion_tokens is not None:
|
|
61
|
+
payload["completionTokens"] = self.completion_tokens
|
|
62
|
+
if self.total_tokens is not None:
|
|
63
|
+
payload["totalTokens"] = self.total_tokens
|
|
64
|
+
if self.error_code:
|
|
65
|
+
payload["errorCode"] = self.error_code
|
|
66
|
+
if self.error_message:
|
|
67
|
+
payload["errorMessage"] = self.error_message
|
|
68
|
+
return payload
|