Add litellm service (#2)

* Start litellm service

* Add Azure OpenAI

* Move litellm to main compose.yaml

* Add litellm DB

Enables cost tracking

README.md

@@ -1 +1,16 @@
 # Trail of Bits AIxCC Finals CRS
+
+# Local Development
+We use `docker compose` to test CRS components locally during development.
+
+Copy `env.template` to `.env` and set variables.
+
+Start the services with
+```
+docker compose up -d
+```
+
+Stop the services with:
+```
+docker compose down
+```

compose.yaml

@@ -9,3 +9,36 @@ services:
       - "127.0.0.1:8000:8000"
     environment:
       - PYTHONUNBUFFERED=1
+
+  litellm:
+    image: ghcr.io/berriai/litellm:litellm_stable_release_branch-v1.57.8-stable
+    configs:
+      - source: litellm_config
+        target: /app/config.yaml
+    env_file: .env
+    ports:
+      - "127.0.0.1:8080:8080"
+    environment:
+        DATABASE_URL: "postgresql://litellm_user:litellm_password11@litellm-db:5432/litellm"
+
+    command: ["--config", "/app/config.yaml", "--port", "8080"]
+    depends_on:
+      litellm-db:
+        condition: service_healthy
+
+  litellm-db:
+    image: postgres:17.2
+    restart: always
+    environment:
+      POSTGRES_DB: litellm
+      POSTGRES_USER: litellm_user
+      POSTGRES_PASSWORD: litellm_password11
+    healthcheck:
+      test: ["CMD-SHELL", "pg_isready -d $${POSTGRES_DB} -U $${POSTGRES_USER}"]
+      interval: 5s
+      timeout: 5s
+      retries: 10
+
+configs:
+  litellm_config:
+    file: ./litellm/litellm_config.yaml

env.template

@@ -0,0 +1,7 @@
+# docker compose .env template
+
+# LiteLLM variables
+# You only need to set variables for APIs you will use.
+AZURE_API_BASE=<INSERT_HOST>
+AZURE_API_KEY=<INSERT_KEY>
+OPENAI_API_KEY=<INSERT_KEY>

litellm/README.md

@@ -0,0 +1,29 @@
+# litellm
+The litellm service proxies all LLMs.
+
+After deploying LiteLLM, you can test it with:
+```
+curl --location 'http://127.0.0.1:8080/chat/completions' \
+--header 'Content-Type: application/json' \
+--header "Authorization: Bearer sk-1234" \
+--data ' {
+      "model": "azure-gpt-4o-mini",
+      "messages": [
+        {
+          "role": "user",
+          "content": "explain the color red"
+        }
+      ]
+    }
+'
+```
+
+Create a virtual key with a budget of 0.01:
+```
+curl 'http://127.0.0.1:8080/key/generate' \
+--header 'Authorization: Bearer sk-1234' \
+--header 'Content-Type: application/json' \
+--data-raw '{"max_budget": 0.01}' | jq ".key"
+```
+
+Use the virtual key in normal requests as the bearer token.

litellm/litellm_config.yaml

@@ -0,0 +1,29 @@
+model_list:
+  - model_name: azure-gpt-4o
+    litellm_params:
+      model: azure/gpt-4o-2024-08-06 # Name defined in Azure Deployment
+      api_base: os.environ/AZURE_API_BASE
+      api_key: os.environ/AZURE_API_KEY
+      tpm: 1000000 # TPM of deployment is 10M
+      rpm: 6000 # RPM of deployment is 60K
+
+  - model_name: azure-gpt-4o-mini
+    litellm_params:
+      model: azure/gpt-4o-mini-2024-07-18 # Name defined in Azure Deployment
+      api_base: os.environ/AZURE_API_BASE
+      api_key: os.environ/AZURE_API_KEY
+      tpm: 2000000 # TPM of deployment is 20M
+      rpm: 20000 # RPM of deployment is 200K
+
+  - model_name: openai-gpt-4o
+    litellm_params:
+      model: openai/gpt-4o
+      api_key: os.environ/OPENAI_API_KEY
+
+  - model_name: openai-gpt-4o-mini
+    litellm_params:
+      model: openai/gpt-4o-mini
+      api_key: os.environ/OPENAI_API_KEY
+
+general_settings:
+  master_key: sk-1234