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.
@@ -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