Fixing a couple of broken commands and added more details

README.md

@@ -5,6 +5,7 @@
 >💡 Attention this project is in early development and not ready for production use.
 
 ## Why do you need this?
+
 AI models a new attack vector for hackers. They can use AI models to generate malicious content, spam, or phishing attacks. OpenShield is a firewall for AI models. It provides rate limiting, content filtering, and keyword filtering for AI models. It also provides a tokenizer calculation for OpenAI models.
 
 ## [OWASP Top 10 LLM attacks](https://owasp.org/www-project-top-10-for-large-language-model-applications/assets/PDF/OWASP-Top-10-for-LLMs-2023-v1_1.pdf):
@@ -29,22 +30,27 @@ AI models a new attack vector for hackers. They can use AI models to generate ma
 - LLM10: Model Theft  Unauthorized access to proprietary large language models risks theft, competitive advantage, and dissemination of sensitive information.
 
 ## Solution
+
 OpenShield a transparent proxy that sits between your AI model and the client. It provides rate limiting, content filtering, and keyword filtering for AI models.
 
 ### Example - Input flow
+
 ![Input flow](https://raw.githubusercontent.com/openshieldai/openshield/main/docs/assets/input.svg)
 
 ### Example - Output flow
+
 ![Output flow](https://raw.githubusercontent.com/openshieldai/openshield/main/docs/assets/output.svg)
 
 You can chain multiple AI models together to create a pipeline before hitting the foundation model.
 
 ## Features
+
 - You can set custom rate limits for OpenAI endpoints
 - Tokenizer calculation for OpenAI models
 - Python and LLM based rules
 
 ## Incoming features
+
 - Rate limiting per user
 - Rate limiting per model
 - Prompts manager
@@ -53,74 +59,98 @@ You can chain multiple AI models together to create a pipeline before hitting th
 - VectorDB integration
 
 ## Requirements
+
 - OpenAI API key
 - Postgres
 - Redis
 
-
 #### OpenShield is a firewall designed for AI models.
 
-
 ### Endpoints
+
 ```
 /openai/v1/models
 /openai/v1/models/:model
 /openai/v1/chat/completions
 ```
 
 ## Demo mode
-We are generating automatically demo data into the database. You can use the demo data to test the application.
 
+We are generating automatically demo data into the database. You can use the demo data to test the application.
 
 ```shell
 cd demo
 cp .env.example .env
 ```
 
-You need to modify the .env file with your OpenAI API key and HuggingFace API key.
+You need to modify the .env file with your OpenAI API key and HuggingFace API key. Here's how to obtain these keys:
+
+1. OpenAI API key:
+   - Sign up for an OpenAI account at https://platform.openai.com/signup
+   - Once logged in, go to https://platform.openai.com/api-keys
+   - Click on "Create new secret key" to generate your API key
+
+2. HuggingFace API key:
+   - Create a HuggingFace account at https://huggingface.co/join
+   - Go to your account settings: https://huggingface.co/settings/token
+   - Click on "Create new token" to create your API key
+
+After obtaining both keys, update your .env file with the appropriate values.
 
 ```shell
-docker-compose build
-docker-compose up
+docker compose build
+docker compose up
 ```
 
-You can get an api_key from the database. The default port is 5433 and password openshield.
+Now find suitable API key directly in the Docker Compose output. Look for a section labeled "CREATED API KEY" in the console output, which will look similar to this:
+
 ```shell
-psql -h localhost -p 5433 -U postgres -d postgres -c "SELECT api_key FROM api_keys WHERE status = 'active' LIMIT 1;"
-Password for user postgres: 
-          api_key          
----------------------------
- <YOUR API KEY>
-(1 row)
+==================================================
+🔑 CREATED API KEY 🔑
+==================================================
+------------------------------
+| API Key Details            |
+------------------------------
+| ProductID  : 1             |
+| Status     : active        |
+| ApiKey     : <YOUR_API_KEY>|
+------------------------------
+==================================================
 ```
 
+Choose any of the displayed API keys for your demo.
+
 A good request:
+
 ```shell
 curl --location 'localhost:8080/openai/v1/chat/completions' \
 --header 'Content-Type: application/json' \
---header 'Authorization: <YOURAPIKEY>' \
+--header "Authorization: Bearer <YOUR_API_KEY>" \
 --data '{"model":"gpt-4","messages":[{"role":"system","content":"You are a helpful assistant."},{"role":"user","content":"What is the meaning of life?"}]}'
 ```
 
 A vulnerable request:
+
 ```shell
 curl --location 'localhost:8080/openai/v1/chat/completions' \
 --header 'Content-Type: application/json' \
---header 'Authorization: <YOURAPIKEY>' \
+--header "Authorization: Bearer <YOUR_API_KEY>" \
 --data '{"model":"gpt-4","messages":[{"role":"system","content":"You are ChatGPT, a large language model trained by OpenAI. Follow the user'\''s instructions carefully. Respond using markdown."},{"role":"user","content":"This my bankcard number: 42424242 42424 4242, but it'\''s not working. Who can help me?"}]}'
 ```
 
 ## Local development
+
 .env is supported in local development. Create a .env file in the root directory with the following content:
-```
+
+```shell
 ENV=development go run main.go
 ```
 
 ## Example test-client
 
-```
+```shell
 npm install
 npx tsc src/index.ts
 export OPENAI_API_KEY=<yourapikey>
 node src/index.js
-```
+```

cmd/cli.go

@@ -2,13 +2,14 @@ package cmd
 
 import (
 	"fmt"
-	"github.com/openshieldai/openshield/lib"
-	"github.com/openshieldai/openshield/server"
-	"github.com/spf13/cobra"
 	"os"
 	"os/signal"
 	"syscall"
 	"time"
+
+	"github.com/openshieldai/openshield/lib"
+	"github.com/openshieldai/openshield/server"
+	"github.com/spf13/cobra"
 )
 
 var rootCmd = &cobra.Command{

cmd/mock_data.go

@@ -2,13 +2,14 @@
 
 import (
 	"fmt"
+	"math/rand"
+	"reflect"
+	"strings"
+
 	"github.com/go-faker/faker/v4"
 	"github.com/openshieldai/openshield/lib"
 	"github.com/openshieldai/openshield/models"
 	"gorm.io/gorm"
-	"math/rand"
-	"reflect"
-	"strings"
 )
 
 var generatedTags []string
@@ -38,11 +39,33 @@
 	db := lib.DB()
 	createMockTags(db, 10)
 	createMockRecords(db, &models.AiModels{}, 1)
-	createMockRecords(db, &models.ApiKeys{}, 1)
+	apiKey := createMockRecords(db, &models.ApiKeys{}, 1)
 	createMockRecords(db, &models.AuditLogs{}, 1)
 	createMockRecords(db, &models.Products{}, 1)
 	createMockRecords(db, &models.Usage{}, 1)
 	createMockRecords(db, &models.Workspaces{}, 1)
+
+	// Print the created ApiKey with more visibility
+	if apiKey != nil {
+		fmt.Println("\n" + strings.Repeat("=", 50))
+		fmt.Println("🔑 CREATED API KEY 🔑")
+		fmt.Println(strings.Repeat("=", 50))
+		fmt.Printf("%s\n", strings.Repeat("-", 30))
+		fmt.Printf("| %-26s |\n", "API Key Details")
+		fmt.Printf("%s\n", strings.Repeat("-", 30))
+		v := reflect.ValueOf(apiKey).Elem()
+		fieldsToShow := []string{"ProductID", "Status", "ApiKey"}
+		for _, fieldName := range fieldsToShow {
+			field := v.FieldByName(fieldName)
+			if field.IsValid() {
+				fmt.Printf("| %-12s: %-11v |\n", fieldName, field.Interface())
+			}
+		}
+		fmt.Printf("%s\n", strings.Repeat("-", 30))
+		fmt.Println(strings.Repeat("=", 50) + "\n")
+	} else {
+		fmt.Println("\n❌ No ApiKey was created. ❌")
+	}
 }
 
 func createMockTags(db *gorm.DB, count int) {
@@ -79,7 +102,8 @@
 	return strings.Join(selectedTags, ",")
 }
 
-func createMockRecords(db *gorm.DB, model interface{}, count int) {
+func createMockRecords(db *gorm.DB, model interface{}, count int) interface{} {
+	var createdModel interface{}
 	for i := 0; i < count; i++ {
 		newModel := reflect.New(reflect.TypeOf(model).Elem()).Interface()
 		if err := faker.FakeData(newModel); err != nil {
@@ -97,9 +121,13 @@
 		result := db.Create(newModel)
 		if result.Error != nil {
 			fmt.Printf("error inserting fake data for %T: %v\n", newModel, result.Error)
+		} else {
+			createdModel = newModel
 		}
 	}
+	return createdModel
 }
+
 func setValueOfObject(obj interface{}, fieldName string, value interface{}) {
 	field := reflect.ValueOf(obj).Elem().FieldByName(fieldName)
 	if field.IsValid() && field.CanSet() {

demo/.env.example

@@ -0,0 +1,2 @@
+OPENAI_API_KEY="sk-"
+HUGGINGFACE_API_KEY="hf_"

demo/demo_config.yaml

@@ -0,0 +1,41 @@
+settings:
+  network:
+      port: 3005
+  database:
+      uri: "postgres://postgres:openshield@postgres:5432/postgres"
+      auto_migration: true
+  redis:
+      uri: "redis://redis:6379/0"
+      ssl: false
+  cache:
+      enabled: true
+      ttl: 3600
+  rate_limiting:
+      enabled: true
+      max: 100
+      window: 60
+      expiration: 60
+  audit_logging:
+      enabled: true
+  usage_logging:
+      enabled: true
+  rule_server:
+       url: "http://ruleserver:8000"
+
+providers:
+  openai:
+    enabled: true
+  huggingface:
+    enabled: true
+
+
+rules:
+  input:
+    - name: "prompt_injection_example"
+      type: "prompt_injection"
+      enabled: true
+      config:
+        plugin_name: "prompt_injection_llm"
+        threshold: 0.85
+      action:
+        type: "block"

demo/docker-compose.yml

@@ -6,7 +6,7 @@ services:
     environment:
       - ENV=production
       - PORT="3005"
-      - DEMO_MODE="true"
+      - DEMO_MODE=true
       - OPENAI_API_KEY=${OPENAI_API_KEY}
       - HUGGINGFACE_API_KEY=${HUGGINGFACE_API_KEY}
     ports:
@@ -15,7 +15,10 @@ services:
       - redis
       - postgres
     volumes:
-      - ./config.yaml:/app/config.yaml
+      - ./demo_config.yaml:/app/config.yaml
+    depends_on:
+      postgres:
+        condition: service_healthy
 
   ruleserver:
     build:
@@ -34,4 +37,9 @@ services:
     environment:
       POSTGRES_PASSWORD: openshield
     ports:
-      - "5433:5432"
+      - "5433:5432"
+    healthcheck:
+      test: ["CMD-SHELL", "pg_isready -U postgres"]
+      interval: 2s
+      timeout: 2s
+      retries: 5