Merge pull request #4 from NeuralTrust/prometheus-metrics

Prometheus metrics

CONTRIBUTING.md

@@ -1,4 +1,4 @@
-# Contributing to AI Gateway CE
+# Contributing to TrustGate
 
 We love your input! We want to make contributing to this project as easy and transparent as possible, whether it's:
 

Dockerfile

@@ -46,6 +46,7 @@ ENV GIN_MODE=release
 COPY scripts/docker-entrypoint.sh /app/
 RUN chmod +x /app/docker-entrypoint.sh
 
-EXPOSE 8080 8081
+# Expose API and metrics ports
+EXPOSE 8080 8081 9090
 
 ENTRYPOINT ["/app/docker-entrypoint.sh"] 

README.md

@@ -66,7 +66,7 @@ docker compose up -d redis postgres
 
 ## 🏗️ Architecture
 
-AI Gateway CE consists of two main components:
+TrustGate consists of two main components:
 
 1. **Admin API** (Port 8080)
    - Gateway management
@@ -79,6 +79,37 @@ AI Gateway CE consists of two main components:
    - Load balancing
    - Plugin execution
 
+## 📊 Monitoring
+
+TrustGate provides comprehensive monitoring through Prometheus metrics:
+
+```yaml
+# config.yaml
+metrics:
+  enabled: true                # Enable metrics collection
+  enable_latency: true        # Basic latency metrics
+  enable_upstream: true       # Upstream latency tracking
+  enable_connections: true    # Connection tracking
+  enable_per_route: true     # Per-route metrics
+```
+
+### Key Metrics
+- `trustgate_requests_total` - Request counts by gateway, method, and status
+- `trustgate_latency_ms` - Overall request processing time
+- `trustgate_detailed_latency_ms` - Granular latency by service/route
+- `trustgate_upstream_latency_ms` - Upstream service latency
+- `trustgate_connections` - Active connection tracking
+
+### Prometheus Configuration
+```yaml
+scrape_configs:
+  - job_name: 'trustgate'
+    static_configs:
+      - targets: ['localhost:9090']
+```
+
+For detailed metrics documentation, dashboards, and advanced queries, see our [Monitoring Guide](https://docs.neuraltrust.ai/monitoring).
+
 ## 🔌 Plugins
 
 Extend functionality with plugins:
@@ -96,8 +127,10 @@ type Plugin interface {
 - Rate Limiter
 - Token Rate Limiter
 - External API Call
+- Data Masking (Pre-defined entities, custom keywords, regex)
+- Prompt Moderation (keywords, regex)
 
-### 🔜 Coming Soon Plugins
+### 🔜 Work in progress (Coming Soon)
 
 #### Security
 - **Jailbreak Protection**
@@ -110,24 +143,13 @@ type Plugin interface {
   - Legacy toxicity detection methods
 
 - **Prompt Moderation**
-  - Keywords & REGEX filtering
   - Topic detection (accepted/denied)
 
-- **Data Masking**
-  - Pre-defined entity masking
-  - Custom data masking (keywords, regex patterns)
-
 - **Network Security**
   - CORS protection
   - SQL Injection prevention
   - Cross-site injection protection
 
-- **Load Balancing**
-  - Weight-based routing
-  - Round-robin distribution
-  - Prompt templates support
-
-
 ## 🤝 Contributing
 
 We love contributions! Please see our [Contributing Guide](CONTRIBUTING.md) for details.

cmd/gateway/main.go

@@ -64,6 +64,25 @@ func (h *ConsoleHook) Levels() []logrus.Level {
 	return logrus.AllLevels
 }
 
+// Get server type safely
+func getServerType() string {
+	if len(os.Args) > 1 {
+		return os.Args[1]
+	}
+	return "proxy" // default to proxy server
+}
+
+func initializeServer(cfg *config.Config, cache *cache.Cache, repo *database.Repository, logger *logrus.Logger) server.Server {
+	serverType := getServerType()
+
+	switch serverType {
+	case "admin":
+		return server.NewAdminServer(cfg, cache, repo, logger)
+	default:
+		return server.NewProxyServer(cfg, cache, repo, logger, false)
+	}
+}
+
 func main() {
 	// Initialize logger
 	logger := logrus.New()
@@ -80,11 +99,8 @@ func main() {
 		logger.SetLevel(logrus.DebugLevel)
 	}
 
-	// Determine server type from command line
-	serverType := "proxy"
-	if len(os.Args) > 1 {
-		serverType = os.Args[1]
-	}
+	// Get server type once at the start
+	serverType := getServerType()
 
 	// Set up logging to file
 	var logFile string
@@ -165,21 +181,8 @@ func main() {
 	// Initialize repository
 	repo := database.NewRepository(db.DB, logger, cacheInstance)
 
-	var srv server.Server
-	switch serverType {
-	case "admin":
-		srv = server.NewAdminServer(cfg, cacheInstance, repo, logger)
-	case "proxy":
-		srv = server.NewProxyServer(
-			cfg,
-			cacheInstance,
-			repo,
-			logger,
-			false, // debug mode
-		)
-	default:
-		logger.Fatalf("Unknown server type: %s", serverType)
-	}
+	// Create and initialize the server
+	srv := initializeServer(cfg, cacheInstance, repo, logger)
 
 	if err := srv.Run(); err != nil {
 		logger.Fatalf("Server failed: %v", err)

config/config.yaml

@@ -1,18 +1,27 @@
 server:
   admin_port: 8080
+  metrics_port: 9090
   proxy_port: 8081
   base_domain: example.com
 
-redis:
-  host: localhost
-  port: 6379
-  password: ""
-  db: 0
+metrics:
+  enabled: true
+  enable_latency: true        # Basic latency metrics
+  enable_upstream: true      # Upstream latency (high cardinality)
+  enable_connections: true   # Connection tracking
+  enable_per_route: true    # Per-route metrics (high cardinality)
+  enable_detailed_status: true # Detailed status codes
 
 database:
   host: localhost
   port: 5432
   user: postgres
   password: postgres
   dbname: ai_gateway
-  ssl_mode: disable 
+  ssl_mode: disable 
+
+redis:
+  host: localhost
+  port: 6379
+  password: ""
+  db: 0

examples/test_rate_limiter.sh

@@ -159,14 +159,23 @@ echo -e "\n${GREEN}6. Testing different rate limit types...${NC}"
 
 # Test IP-based rate limit
 echo -e "\n${GREEN}6.1 Testing IP-based rate limit (limit: 5/min)...${NC}"
-for i in {1..6}; do
-    RESPONSE=$(curl -s -w "\n%{http_code}" "$PROXY_URL/path2" \
+for i in {1..1}; do
+    echo "Making request $i..."
+    start_time=$(date +%s.%N)
+
+    RESPONSE=$(curl -s -w "\n%{http_code}\n%{time_total}" "$PROXY_URL/path2" \
         -H "Host: ${SUBDOMAIN}.${BASE_DOMAIN}" \
         -H "X-API-Key: ${API_KEY}" \
         -H "X-Real-IP: 192.168.1.1")
 
-    HTTP_CODE=$(echo "$RESPONSE" | tail -n1)
+    duration=$(echo "$RESPONSE" | tail -n1)
+    HTTP_CODE=$(echo "$RESPONSE" | tail -n2 | head -n1)
     BODY=$(echo "$RESPONSE" | head -n1)
+    end_time=$(date +%s.%N)
+    total_time=$(echo "$end_time - $start_time" | bc)
+
+    echo -e "Request duration (curl): ${duration}s"
+    echo -e "Total script duration: ${total_time}s"
 
     if [ "$HTTP_CODE" == "200" ]; then
         echo -e "${GREEN}IP-based Request $i: Success${NC}"

go.mod

@@ -9,6 +9,7 @@ require (
 	github.com/google/uuid v1.6.0
 	github.com/lib/pq v1.10.9
 	github.com/mitchellh/mapstructure v1.5.0
+	github.com/prometheus/client_golang v1.20.5
 	github.com/sirupsen/logrus v1.9.3
 	github.com/spf13/viper v1.19.0
 	github.com/stretchr/testify v1.9.0
@@ -21,8 +22,9 @@ require (
 
 require (
 	github.com/andybalholm/brotli v1.1.1 // indirect
+	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/bytedance/sonic/loader v0.2.0 // indirect
-	github.com/cespare/xxhash/v2 v2.1.2 // indirect
+	github.com/cespare/xxhash/v2 v2.3.0 // indirect
 	github.com/cloudwego/base64x v0.1.4 // indirect
 	github.com/cloudwego/iasm v0.2.0 // indirect
 	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
@@ -49,8 +51,12 @@ require (
 	github.com/mattn/go-isatty v0.0.20 // indirect
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.2 // indirect
+	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
 	github.com/pelletier/go-toml/v2 v2.2.2 // indirect
 	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
+	github.com/prometheus/client_model v0.6.1 // indirect
+	github.com/prometheus/common v0.55.0 // indirect
+	github.com/prometheus/procfs v0.15.1 // indirect
 	github.com/rogpeppe/go-internal v1.13.1 // indirect
 	github.com/sagikazarmark/locafero v0.4.0 // indirect
 	github.com/sagikazarmark/slog-shim v0.1.0 // indirect
@@ -70,6 +76,6 @@ require (
 	golang.org/x/sync v0.10.0 // indirect
 	golang.org/x/sys v0.28.0 // indirect
 	golang.org/x/text v0.21.0 // indirect
-	google.golang.org/protobuf v1.34.1 // indirect
+	google.golang.org/protobuf v1.34.2 // indirect
 	gopkg.in/ini.v1 v1.67.0 // indirect
 )

go.sum

@@ -1,12 +1,14 @@
 github.com/andybalholm/brotli v1.1.1 h1:PR2pgnyFznKEugtsUo0xLdDop5SKXd5Qf5ysW+7XdTA=
 github.com/andybalholm/brotli v1.1.1/go.mod h1:05ib4cKhjx3OQYUY22hTVd34Bc8upXjOLL2rKwwZBoA=
+github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
+github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=
 github.com/bytedance/sonic v1.12.5 h1:hoZxY8uW+mT+OpkcUWw4k0fDINtOcVavEsGfzwzFU/w=
 github.com/bytedance/sonic v1.12.5/go.mod h1:B8Gt/XvtZ3Fqj+iSKMypzymZxw/FVwgIGKzMzT9r/rk=
 github.com/bytedance/sonic/loader v0.1.1/go.mod h1:ncP89zfokxS5LZrJxl5z0UJcsk4M4yY2JpfqGeCtNLU=
 github.com/bytedance/sonic/loader v0.2.0 h1:zNprn+lsIP06C/IqCHs3gPQIvnvpKbbxyXQP1iU4kWM=
 github.com/bytedance/sonic/loader v0.2.0/go.mod h1:ncP89zfokxS5LZrJxl5z0UJcsk4M4yY2JpfqGeCtNLU=
-github.com/cespare/xxhash/v2 v2.1.2 h1:YRXhKfTDauu4ajMg1TPgFO5jnlC2HCbmLXMcTG5cbYE=
-github.com/cespare/xxhash/v2 v2.1.2/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
+github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
+github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
 github.com/cloudwego/base64x v0.1.4 h1:jwCgWpFanWmN8xoIUHa2rtzmkd5J2plF/dnLS6Xd/0Y=
 github.com/cloudwego/base64x v0.1.4/go.mod h1:0zlkT4Wn5C6NdauXdJRhSKRlJvmclQ1hhJgA0rcu/8w=
 github.com/cloudwego/iasm v0.2.0 h1:1KNIy1I1H9hNNFEEH3DVnI4UujN+1zjpuk6gwHLTssg=
@@ -39,8 +41,8 @@ github.com/go-redis/redis/v8 v8.11.5 h1:AcZZR7igkdvfVmQTPnu9WE37LRrO/YrBH5zWyjDC
 github.com/go-redis/redis/v8 v8.11.5/go.mod h1:gREzHqY1hg6oD9ngVRbLStwAWKhA0FEgq8Jd4h5lpwo=
 github.com/goccy/go-json v0.10.2 h1:CrxCmQqYDkv1z7lO7Wbh2HN93uovUHgrECaO5ZrCXAU=
 github.com/goccy/go-json v0.10.2/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I=
-github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38=
-github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
+github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
+github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
 github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
@@ -70,6 +72,8 @@ github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
 github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
 github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
 github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
+github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc=
+github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw=
 github.com/leodido/go-urn v1.4.0 h1:WT9HwE9SGECu3lg4d/dIA+jxlljEa1/ffXKmRjqdmIQ=
 github.com/leodido/go-urn v1.4.0/go.mod h1:bvxc+MVxLKB4z00jd1z+Dvzr47oO32F/QSNjSBOlFxI=
 github.com/lib/pq v1.10.9 h1:YXG7RB+JIjhP29X+OtkiDnYaXQwpS4JEWq7dtCCRUEw=
@@ -85,6 +89,8 @@ github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w
 github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
 github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M=
 github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk=
+github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA=
+github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ=
 github.com/nxadm/tail v1.4.8 h1:nPr65rt6Y5JFSKQO7qToXr7pePgD6Gwiw05lkbyAQTE=
 github.com/nxadm/tail v1.4.8/go.mod h1:+ncqLTQzXmGhMZNUePPaPqPvBxHAIsmXswZKocGu+AU=
 github.com/onsi/ginkgo v1.16.5 h1:8xi0RTUf59SOSfEtZMvwTvXYMzG4gV23XVHOZiXNtnE=
@@ -96,6 +102,14 @@ github.com/pelletier/go-toml/v2 v2.2.2/go.mod h1:1t835xjRzz80PqgE6HHgN2JOsmgYu/h
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U=
 github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/prometheus/client_golang v1.20.5 h1:cxppBPuYhUnsO6yo/aoRol4L7q7UFfdm+bR9r+8l63Y=
+github.com/prometheus/client_golang v1.20.5/go.mod h1:PIEt8X02hGcP8JWbeHyeZ53Y/jReSnHgO035n//V5WE=
+github.com/prometheus/client_model v0.6.1 h1:ZKSh/rekM+n3CeS952MLRAdFwIKqeY8b62p8ais2e9E=
+github.com/prometheus/client_model v0.6.1/go.mod h1:OrxVMOVHjw3lKMa8+x6HeMGkHMQyHDk9E3jmP2AmGiY=
+github.com/prometheus/common v0.55.0 h1:KEi6DK7lXW/m7Ig5i47x0vRzuBsHuvJdi5ee6Y3G1dc=
+github.com/prometheus/common v0.55.0/go.mod h1:2SECS4xJG1kd8XF9IcM1gMX6510RAEL65zxzNImwdc8=
+github.com/prometheus/procfs v0.15.1 h1:YagwOFzUgYfKKHX6Dr+sHT7km/hxC76UB0learggepc=
+github.com/prometheus/procfs v0.15.1/go.mod h1:fB45yRUv8NstnjriLhBQLuOUt+WW4BsoGhij/e3PBqk=
 github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR38lUII=
 github.com/rogpeppe/go-internal v1.13.1/go.mod h1:uMEvuHeurkdAXX61udpOXGD/AzZDWNMNyH2VO9fmH0o=
 github.com/sagikazarmark/locafero v0.4.0 h1:HApY1R9zGo4DBgr7dqsTH/JJxLTTsOt7u6keLGt6kNQ=
@@ -159,8 +173,8 @@ golang.org/x/sys v0.28.0 h1:Fksou7UEQUWlKvIdsqzJmUmCX3cZuD2+P3XyyzwMhlA=
 golang.org/x/sys v0.28.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/text v0.21.0 h1:zyQAAkrwaneQ066sspRyJaG9VNi/YJ1NfzcGB3hZ/qo=
 golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ=
-google.golang.org/protobuf v1.34.1 h1:9ddQBjfCyZPOHPUiPxpYESBLc+T8P3E+Vo4IbKZgFWg=
-google.golang.org/protobuf v1.34.1/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos=
+google.golang.org/protobuf v1.34.2 h1:6xV6lTsCfpGD21XK49h7MhtcApnLqkfYgPcdHftf6hg=
+google.golang.org/protobuf v1.34.2/go.mod h1:qYOHts0dSfpeUzUFpOMr/WGzszTmLH+DiWniOlNbLDw=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
 gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=