Fixing a couple of broken commands and added more details

README.md

@@ -1,29 +1,33 @@
 # OpenShield - Firewall for AI models
 
-
 >💡 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.
 
 ## 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
@@ -32,72 +36,90 @@ 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
+cp ./example_config ./config.yaml/config
 ```
 
-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 "DEMO API KEYS" 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)
+🔑 ======== DEMO API KEYS ============ 🔑
+🚀 Behold! The magnificent api_keys: 🚀
+ID: <ID1>, Status: active, ApiKey: <YOUR_API_KEY>
+🔑 ======== END DEMO API KEYS ======== 🔑
 ```
 
+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>

cmd/cli.go

@@ -2,13 +2,15 @@
 
 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/models"
+	"github.com/openshieldai/openshield/server"
+	"github.com/spf13/cobra"
 )
 
 var rootCmd = &cobra.Command{
@@ -27,6 +29,7 @@
 	rootCmd.AddCommand(stopServerCmd)
 	dbCmd.AddCommand(createTablesCmd)
 	dbCmd.AddCommand(createMockDataCmd)
+	dbCmd.AddCommand(queryApiKeysCmd)
 	configCmd.AddCommand(editConfigCmd)
 	configCmd.AddCommand(addRuleCmd)
 	configCmd.AddCommand(removeRuleCmd)
@@ -54,6 +57,29 @@
 	},
 }
 
+var queryApiKeysCmd = &cobra.Command{
+	Use:   "query-api-keys",
+	Short: "Query and display data from the api_keys table",
+	Run: func(cmd *cobra.Command, args []string) {
+		queryApiKeys()
+	},
+}
+
+func queryApiKeys() {
+	db := lib.DB()
+	var apiKeys []models.ApiKeys
+	result := db.Limit(5).Find(&apiKeys)
+	if result.Error != nil {
+		fmt.Printf("Error querying api_keys: %v\n", result.Error)
+		return
+	}
+
+	fmt.Println("Sample data from api_keys table:")
+	for _, key := range apiKeys {
+		fmt.Printf("ID: %s, Status: %s, ApiKey: %s\n", key.Id, key.Status, key.ApiKey)
+	}
+}
+
 var configCmd = &cobra.Command{
 	Use:   "config",
 	Short: "Configuration related commands",

demo/.env.example

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

demo/example_config

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

docker/docker-entrypoint.sh

@@ -11,6 +11,12 @@ if [ -n "$DEMO_MODE" ]; then
   echo "Running in DEMO_MODE"
   ./openshield db create-tables
   ./openshield db create-mock-data
+
+  # Query and display data from the api_keys table
+  echo "🔑 ======== DEMO API KEYS ============ 🔑"
+  echo "🚀 Behold! The magnificent api_keys: 🚀"
+  ./openshield db query-api-keys
+  echo "🔑 ======== END DEMO API KEYS ======== 🔑"
 fi
 
 # Run the command passed to the script using dumb-init

lib/config.go

@@ -117,7 +117,7 @@ func init() {
 
 	viperCfg.SetConfigName("config")
 	viperCfg.SetConfigType("yaml")
-	viperCfg.AddConfigPath(".")
+	viperCfg.AddConfigPath("./config.yaml")
 	viperCfg.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
 
 	err := viperCfg.ReadInConfig()