mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
feat(terraform): expose key type on virtual keys
This commit is contained in:
parent
bc3b5b1d5b
commit
fe5e8298bb
7 changed files with 64 additions and 1 deletions
|
|
@ -17,6 +17,7 @@ longer signal it.
|
|||
### Added
|
||||
|
||||
- **key**: Computed `server_metadata` attribute on `litellm_key` exposing every metadata entry the proxy stores, so metadata created outside Terraform is visible in state and drift on it shows on refresh, while `metadata` keeps tracking only the declared entries and updates keep preserving undeclared ones
|
||||
- **key**: Optional `key_type` argument on `litellm_key` for selecting the key's default route access; changing it creates a new key
|
||||
- **team_member_add**: `tpm_limit`, `rpm_limit`, `budget_duration`, and `allowed_models` attributes on `litellm_team_member_add`, applied to every member of the resource; `budget_duration` and `allowed_models` ride on `/team/member_add`, while the limits are sent through `/team/member_update`, which is where the proxy accepts them
|
||||
- **team**: Optional `team_id` argument on `litellm_team`, so teams can be created with a stable, human-readable ID instead of a provider-generated UUID; changing it forces replacement
|
||||
- `litellm_jwt_key_mapping` accepts `token_id` as an alternative to `key`, so a
|
||||
|
|
|
|||
|
|
@ -79,6 +79,7 @@ Here's an example of creating an API key with various options:
|
|||
|
||||
```hcl
|
||||
resource "litellm_key" "example_key" {
|
||||
key_type = "llm_api"
|
||||
models = ["gpt-4", "claude-3.5-sonnet"]
|
||||
max_budget = 100.0
|
||||
user_id = "user123"
|
||||
|
|
@ -123,6 +124,7 @@ resource "litellm_key" "example_key" {
|
|||
|
||||
The <code>litellm_key</code> resource supports the following options:
|
||||
|
||||
- <code>key_type</code>: Choose the key's default route access
|
||||
- <code>models</code>: List of allowed models for this key
|
||||
- <code>max_budget</code>: Maximum budget for the key
|
||||
- <code>user_id</code> and <code>team_id</code>: Associate the key with a user and team
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ Manages a LiteLLM API key.
|
|||
|
||||
```hcl
|
||||
resource "litellm_key" "example" {
|
||||
key_type = "llm_api"
|
||||
models = ["gpt-3.5-turbo", "gpt-4"]
|
||||
max_budget = 100.0
|
||||
user_id = "user123"
|
||||
|
|
@ -52,6 +53,8 @@ resource "litellm_key" "example" {
|
|||
|
||||
The following arguments are supported:
|
||||
|
||||
* `key_type` - (Optional) Type of key that determines its default allowed routes. One of `llm_api`, `management`, `read_only` or `default`. Changing it creates a new key.
|
||||
|
||||
* `models` - (Optional) List of models that can be used with this key. This restricts the key to only use the specified models.
|
||||
|
||||
* `max_budget` - (Optional) Maximum budget for this key. This sets an upper limit on the total spend allowed for this key.
|
||||
|
|
|
|||
|
|
@ -242,6 +242,10 @@ func (c *Client) parseKeyResponse(resp map[string]interface{}) (*Key, error) {
|
|||
if s, ok := v.(string); ok {
|
||||
createdKey.TokenID = s
|
||||
}
|
||||
case "key_type":
|
||||
if s, ok := v.(string); ok {
|
||||
createdKey.KeyType = s
|
||||
}
|
||||
case "models":
|
||||
if models, ok := v.([]interface{}); ok {
|
||||
createdKey.Models = make([]string, len(models))
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ import (
|
|||
"github.com/hashicorp/go-cty/cty"
|
||||
"github.com/hashicorp/terraform-plugin-sdk/v2/diag"
|
||||
"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
|
||||
"github.com/hashicorp/terraform-plugin-sdk/v2/helper/validation"
|
||||
)
|
||||
|
||||
func resourceKey() *schema.Resource {
|
||||
|
|
@ -34,6 +35,14 @@ func resourceKey() *schema.Resource {
|
|||
Type: schema.TypeString,
|
||||
Computed: true,
|
||||
},
|
||||
"key_type": {
|
||||
Type: schema.TypeString,
|
||||
Optional: true,
|
||||
Computed: true,
|
||||
ForceNew: true,
|
||||
ValidateFunc: validation.StringInSlice([]string{"llm_api", "management", "read_only", "default"}, false),
|
||||
Description: "Type of key that determines its default allowed routes. Changing it creates a new key",
|
||||
},
|
||||
"models": {
|
||||
Type: schema.TypeList,
|
||||
Optional: true,
|
||||
|
|
@ -449,6 +458,7 @@ func resourceKeyDelete(ctx context.Context, d *schema.ResourceData, m interface{
|
|||
}
|
||||
|
||||
func mapResourceDataToKey(d *schema.ResourceData, key *Key) {
|
||||
key.KeyType = d.Get("key_type").(string)
|
||||
key.Models = expandStringList(d.Get("models").([]interface{}))
|
||||
if v, ok := d.GetOk("max_budget"); ok {
|
||||
val := v.(float64)
|
||||
|
|
@ -504,6 +514,9 @@ func mapKeyToResourceData(d *schema.ResourceData, key *Key) {
|
|||
|
||||
// Note: "key" is write-only and must not be set here (Read operations).
|
||||
// It is only set during Create so it is available during apply.
|
||||
if key.KeyType != "" {
|
||||
d.Set("key_type", key.KeyType)
|
||||
}
|
||||
|
||||
if len(key.Models) > 0 {
|
||||
d.Set("models", key.Models)
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ func newKeyResourceData(t *testing.T, raw map[string]interface{}) *schema.Resour
|
|||
|
||||
func TestMapResourceDataToKeyNewFields(t *testing.T) {
|
||||
d := newKeyResourceData(t, map[string]interface{}{
|
||||
"key_type": "llm_api",
|
||||
"budget_id": "budget-1",
|
||||
"enforced_params": []interface{}{"user"},
|
||||
"allowed_routes": []interface{}{"/chat/completions"},
|
||||
|
|
@ -36,6 +37,9 @@ func TestMapResourceDataToKeyNewFields(t *testing.T) {
|
|||
key := &Key{}
|
||||
mapResourceDataToKey(d, key)
|
||||
|
||||
if key.KeyType != "llm_api" {
|
||||
t.Errorf("KeyType = %q, want llm_api", key.KeyType)
|
||||
}
|
||||
if key.BudgetID != "budget-1" {
|
||||
t.Errorf("BudgetID = %q, want budget-1", key.BudgetID)
|
||||
}
|
||||
|
|
@ -139,6 +143,7 @@ func TestParseKeyResponseNewFields(t *testing.T) {
|
|||
client := NewClient("http://localhost:4000", "test-key", true)
|
||||
resp := map[string]interface{}{
|
||||
"key": "sk-test",
|
||||
"key_type": "llm_api",
|
||||
"budget_id": "budget-1",
|
||||
"enforced_params": []interface{}{"user"},
|
||||
"allowed_routes": []interface{}{"/chat/completions"},
|
||||
|
|
@ -154,6 +159,9 @@ func TestParseKeyResponseNewFields(t *testing.T) {
|
|||
if err != nil {
|
||||
t.Fatalf("parseKeyResponse returned error: %v", err)
|
||||
}
|
||||
if key.KeyType != "llm_api" {
|
||||
t.Errorf("KeyType = %q, want llm_api", key.KeyType)
|
||||
}
|
||||
if key.BudgetID != "budget-1" || key.OrganizationID != "org-1" || key.ProjectID != "proj-1" {
|
||||
t.Errorf("string fields not parsed: %+v", key)
|
||||
}
|
||||
|
|
@ -183,7 +191,10 @@ func TestCreateKeySendsConfigSuppliedKey(t *testing.T) {
|
|||
defer srv.Close()
|
||||
|
||||
client := NewClient(srv.URL, "test-key", true)
|
||||
d := newKeyResourceData(t, map[string]interface{}{"key": "sk-custom"})
|
||||
d := newKeyResourceData(t, map[string]interface{}{
|
||||
"key": "sk-custom",
|
||||
"key_type": "llm_api",
|
||||
})
|
||||
|
||||
diags := resourceKeyCreate(context.Background(), d, client)
|
||||
if diags.HasError() {
|
||||
|
|
@ -192,11 +203,35 @@ func TestCreateKeySendsConfigSuppliedKey(t *testing.T) {
|
|||
if captured["key"] != "sk-custom" {
|
||||
t.Errorf("create payload key = %v, want sk-custom", captured["key"])
|
||||
}
|
||||
if captured["key_type"] != "llm_api" {
|
||||
t.Errorf("create payload key_type = %v, want llm_api", captured["key_type"])
|
||||
}
|
||||
if d.Id() != "hash-1" {
|
||||
t.Errorf("resource ID = %q, want hash-1", d.Id())
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeyTypeRejectsUnknownValue(t *testing.T) {
|
||||
_, errs := resourceKey().Schema["key_type"].ValidateFunc("unrestricted", "key_type")
|
||||
if len(errs) == 0 {
|
||||
t.Fatal("key_type accepted an unknown value")
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeyTypeChangeForcesReplacement(t *testing.T) {
|
||||
res := resourceKey()
|
||||
priorData := newKeyResourceData(t, map[string]interface{}{"key_type": "default"})
|
||||
priorData.SetId("hash-1")
|
||||
config := terraform.NewResourceConfigRaw(map[string]interface{}{"key_type": "llm_api"})
|
||||
diff, err := res.Diff(context.Background(), priorData.State(), config, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("diff failed: %v", err)
|
||||
}
|
||||
if diff == nil || !diff.RequiresNew() {
|
||||
t.Fatalf("changing key_type must force replacement, diff = %+v", diff)
|
||||
}
|
||||
}
|
||||
|
||||
// The proxy validates each model_max_budget entry as a BudgetConfig object and
|
||||
// 500s on a bare number, so the JSON string must reach /key/generate as nested
|
||||
// objects and the proxy's response must map back to equivalent JSON in state.
|
||||
|
|
@ -370,6 +405,7 @@ func TestGetKeyUnwrapsInfoEnvelope(t *testing.T) {
|
|||
w.Write([]byte(`{
|
||||
"key": "hash-1",
|
||||
"info": {
|
||||
"key_type": "llm_api",
|
||||
"key_alias": "envelope-alias",
|
||||
"models": ["gpt-4o-mini"],
|
||||
"budget_id": "budget-1",
|
||||
|
|
@ -388,6 +424,9 @@ func TestGetKeyUnwrapsInfoEnvelope(t *testing.T) {
|
|||
if key.KeyAlias != "envelope-alias" {
|
||||
t.Errorf("KeyAlias = %q, want envelope-alias (info envelope not unwrapped)", key.KeyAlias)
|
||||
}
|
||||
if key.KeyType != "llm_api" {
|
||||
t.Errorf("KeyType = %q, want llm_api", key.KeyType)
|
||||
}
|
||||
if key.BudgetID != "budget-1" || key.TeamID != "team-1" {
|
||||
t.Errorf("nested fields not parsed: %+v", key)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -120,6 +120,7 @@ type ModelInfo struct {
|
|||
type Key struct {
|
||||
Key string `json:"key,omitempty"`
|
||||
TokenID string `json:"token_id,omitempty"`
|
||||
KeyType string `json:"key_type,omitempty"`
|
||||
Models []string `json:"models"`
|
||||
Spend float64 `json:"spend,omitempty"`
|
||||
MaxBudget *float64 `json:"max_budget,omitempty"`
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue