Update environment configuration, enhance token management, and add UUID support for user models

Author: farhoud.m7

.env.example

@@ -7,6 +7,7 @@ DB_PATH=app.db
 AZURE_OPENAI_ENDPOINT=https://your-resource-name.openai.azure.com
 AZURE_OPENAI_KEY=your-api-key-here
 AZURE_OPENAI_DEPLOYMENT=your-deployment-name
+AZURE_OPENAI_DEPLOYMENT_VERSION=2024-12-01-preview
 
 # Description of each variable:
 # PORT: The port number the server will listen on (default: 8080)

.gitignore

@@ -8,3 +8,4 @@ docs/swagger.json
 docs/swagger.yaml
 docs/docs.go
 app.db
+.env

auth.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"fmt"
 	"net/http"
 	"time"
 
@@ -39,6 +40,7 @@ func generateToken(username, tokenType string, expiresIn time.Duration) (string,
 func authMiddleware() gin.HandlerFunc {
 	return func(c *gin.Context) {
 		tokenString := c.GetHeader("Authorization")
+		fmt.Println(tokenString)
 		if tokenString == "" {
 			c.JSON(http.StatusUnauthorized, models.NewErrorResponse("Authorization header is required"))
 			c.Abort()
@@ -68,7 +70,17 @@ func authMiddleware() gin.HandlerFunc {
 			return
 		}
 
+		// Get user from database to get the ID
+		var user database.DBUser
+		if err := database.GetDB().Where("username = ?", claims.Username).First(&user).Error; err != nil {
+			c.JSON(http.StatusUnauthorized, models.NewErrorResponse("User not found"))
+			c.Abort()
+			return
+		}
+
+		// Set both username and user_id in context
 		c.Set("username", claims.Username)
+		c.Set("user_id", user.ID)
 		c.Next()
 	}
 }

azure/openai.go

@@ -8,6 +8,7 @@ import (
 	"time"
 
 	"github.com/gin-gonic/gin"
+	"github.com/google/uuid"
 	"github.com/vhybZApp/api/config"
 	"github.com/vhybZApp/api/database"
 	"github.com/vhybZApp/api/models"
@@ -79,6 +80,12 @@ func ChatCompletion(c *gin.Context) {
 	// Initialize token quota service
 	tokenQuotaService := services.NewTokenQuotaService(database.GetDB())
 
+	// Check if user has enough quota for a reasonable estimate (e.g., 1000 tokens)
+	if err := tokenQuotaService.UpdateUsage(userID.(uuid.UUID), 1000); err != nil {
+		c.JSON(http.StatusTooManyRequests, models.NewErrorResponse(err.Error()))
+		return
+	}
+
 	// Validate Azure OpenAI configuration
 	if config.AppConfig.AzureOpenAIEndpoint == "" || config.AppConfig.AzureOpenAIKey == "" || config.AppConfig.AzureOpenAIDeployment == "" {
 		c.JSON(http.StatusInternalServerError, models.NewErrorResponse("Azure OpenAI configuration is incomplete"))
@@ -94,7 +101,7 @@ func ChatCompletion(c *gin.Context) {
 
 	// Create HTTP client
 	client := &http.Client{
-		Timeout: 30 * time.Second,
+		Timeout: 300 * time.Second,
 	}
 
 	// Marshal request body
@@ -105,7 +112,9 @@ func ChatCompletion(c *gin.Context) {
 	}
 
 	// Create request
-	url := config.AppConfig.AzureOpenAIEndpoint + "/openai/deployments/" + config.AppConfig.AzureOpenAIDeployment + "/chat/completions?api-version=2023-05-15"
+	url := config.AppConfig.AzureOpenAIEndpoint + "/openai/deployments/" +
+		config.AppConfig.AzureOpenAIDeployment + "/chat/completions?api-version=" + config.AppConfig.AzureOpenAIDeploymentVersion
+
 	httpReq, err := http.NewRequest("POST", url, bytes.NewBuffer(reqBody))
 	if err != nil {
 		c.JSON(http.StatusInternalServerError, models.NewErrorResponse("Error creating request"))
@@ -141,14 +150,7 @@ func ChatCompletion(c *gin.Context) {
 	var chatResp ChatCompletionResponse
 	if err := json.Unmarshal(body, &chatResp); err != nil {
 		c.JSON(http.StatusInternalServerError, models.NewErrorResponse("Error parsing response"))
-		return
-	}
-
-	// Update token usage
-	if err := tokenQuotaService.UpdateUsage(userID.(string), chatResp.Usage.TotalTokens); err != nil {
-		c.JSON(http.StatusTooManyRequests, models.NewErrorResponse(err.Error()))
-		return
-	}
+		returnasdu
 
 	c.JSON(http.StatusOK, chatResp)
 }

config/config.go

@@ -12,9 +12,10 @@ type Config struct {
 	JWTSecret string
 	DBPath    string
 	// Azure OpenAI Configuration
-	AzureOpenAIEndpoint   string
-	AzureOpenAIKey        string
-	AzureOpenAIDeployment string
+	AzureOpenAIEndpoint          string
+	AzureOpenAIKey               string
+	AzureOpenAIDeployment        string
+	AzureOpenAIDeploymentVersion string
 }
 
 var AppConfig Config
@@ -27,12 +28,13 @@ func LoadConfig() {
 
 	// Set default values
 	AppConfig = Config{
-		Port:                  getEnv("PORT", "8080"),
-		JWTSecret:             getEnv("JWT_SECRET", "your-default-secret-key"),
-		DBPath:                getEnv("DB_PATH", "app.db"),
-		AzureOpenAIEndpoint:   getEnv("AZURE_OPENAI_ENDPOINT", ""),
-		AzureOpenAIKey:        getEnv("AZURE_OPENAI_KEY", ""),
-		AzureOpenAIDeployment: getEnv("AZURE_OPENAI_DEPLOYMENT", ""),
+		Port:                         getEnv("PORT", "8080"),
+		JWTSecret:                    getEnv("JWT_SECRET", "your-default-secret-key"),
+		DBPath:                       getEnv("DB_PATH", "app.db"),
+		AzureOpenAIEndpoint:          getEnv("AZURE_OPENAI_ENDPOINT", ""),
+		AzureOpenAIKey:               getEnv("AZURE_OPENAI_KEY", ""),
+		AzureOpenAIDeployment:        getEnv("AZURE_OPENAI_DEPLOYMENT", "gpt-4o"),
+		AzureOpenAIDeploymentVersion: getEnv("AZURE_OPENAI_DEPLOYMENT_VERSION", "gpt-4o"),
 	}
 
 	// Validate required configurations

database/models.go

@@ -3,32 +3,44 @@ package database
 import (
 	"time"
 
+	"github.com/google/uuid"
 	"golang.org/x/crypto/bcrypt"
 	"gorm.io/gorm"
 )
 
 // DBUser represents the user data as stored in the database
 type DBUser struct {
-	gorm.Model
-	Username     string `gorm:"uniqueIndex;not null"`
-	Password     string `gorm:"not null"`
-	Email        string `gorm:"uniqueIndex;not null"`
-	RefreshToken string `gorm:"-"` // Not stored in DB, only used temporarily
+	ID           uuid.UUID `gorm:"type:uuid;primary_key"`
+	CreatedAt    time.Time
+	UpdatedAt    time.Time
+	DeletedAt    gorm.DeletedAt `gorm:"index"`
+	Username     string         `gorm:"uniqueIndex;not null"`
+	Password     string         `gorm:"not null"`
+	Email        string         `gorm:"uniqueIndex;not null"`
+	RefreshToken string         `gorm:"-"` // Not stored in DB, only used temporarily
+}
+
+// BeforeCreate will set a UUID rather than numeric ID
+func (u *DBUser) BeforeCreate(tx *gorm.DB) error {
+	u.ID = uuid.New()
+	return nil
 }
 
 // DBTokenUsage represents the daily token usage for a user in the database
 type DBTokenUsage struct {
 	gorm.Model
-	UserID string    `gorm:"index"`
+	UserID uuid.UUID `gorm:"type:uuid;index;foreignKey:ID;references:ID;onDelete:CASCADE"`
+	User   DBUser    `gorm:"foreignKey:UserID"`
 	Date   time.Time `gorm:"index"`
 	Tokens int       `gorm:"default:0"`
 }
 
 // DBTokenQuota represents the daily token quota for a user in the database
 type DBTokenQuota struct {
 	gorm.Model
-	UserID     string `gorm:"uniqueIndex"`
-	DailyQuota int    `gorm:"default:100000"` // Default 100k tokens per day
+	UserID     uuid.UUID `gorm:"type:uuid;uniqueIndex;foreignKey:ID;references:ID;onDelete:CASCADE"`
+	User       DBUser    `gorm:"foreignKey:UserID"`
+	DailyQuota int       `gorm:"default:100000"` // Default 100k tokens per day
 }
 
 // HashPassword hashes the password using bcrypt

go.mod

@@ -5,6 +5,7 @@ go 1.21
 require (
 	github.com/gin-gonic/gin v1.9.1
 	github.com/golang-jwt/jwt/v5 v5.0.0
+	github.com/google/uuid v1.6.0
 	github.com/joho/godotenv v1.5.1
 	github.com/swaggo/files v1.0.1
 	github.com/swaggo/gin-swagger v1.6.0

go.sum

@@ -48,6 +48,8 @@ github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaS
 github.com/google/go-cmp v0.5.5 h1:Khx7svrCpmxxtHBq5j2mp/xVjsi8hQMfNLvJFAlrGgU=
 github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 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=
 github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E=
 github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc=
 github.com/jinzhu/now v1.1.5 h1:/o9tlHleP7gOFmsnYNz3RGnqzefHA47wQpKrrdTIwXQ=

main.go

@@ -3,13 +3,14 @@ package main
 import (
 	"log"
 	"net/http"
-	_ "github.com/vhybZApp/api/docs"
+
 	"github.com/gin-gonic/gin"
 	swaggerFiles "github.com/swaggo/files"
 	ginSwagger "github.com/swaggo/gin-swagger"
 	"github.com/vhybZApp/api/azure"
 	"github.com/vhybZApp/api/config"
 	"github.com/vhybZApp/api/database"
+	_ "github.com/vhybZApp/api/docs"
 )
 
 // @title           Vhybz API

services/token_quota.go

@@ -4,6 +4,7 @@ import (
 	"errors"
 	"time"
 
+	"github.com/google/uuid"
 	"github.com/vhybZApp/api/database"
 	"gorm.io/gorm"
 )
@@ -17,7 +18,7 @@ func NewTokenQuotaService(db *gorm.DB) *TokenQuotaService {
 }
 
 // GetUserQuota returns the daily token quota for a user
-func (s *TokenQuotaService) GetUserQuota(userID string) (*database.DBTokenQuota, error) {
+func (s *TokenQuotaService) GetUserQuota(userID uuid.UUID) (*database.DBTokenQuota, error) {
 	var quota database.DBTokenQuota
 	result := s.db.Where("user_id = ?", userID).First(&quota)
 	if result.Error != nil {
@@ -38,7 +39,7 @@ func (s *TokenQuotaService) GetUserQuota(userID string) (*database.DBTokenQuota,
 }
 
 // GetDailyUsage returns the token usage for a user on a specific date
-func (s *TokenQuotaService) GetDailyUsage(userID string, date time.Time) (*database.DBTokenUsage, error) {
+func (s *TokenQuotaService) GetDailyUsage(userID uuid.UUID, date time.Time) (*database.DBTokenUsage, error) {
 	var usage database.DBTokenUsage
 	startOfDay := time.Date(date.Year(), date.Month(), date.Day(), 0, 0, 0, 0, date.Location())
 	result := s.db.Where("user_id = ? AND date = ?", userID, startOfDay).First(&usage)
@@ -61,7 +62,7 @@ func (s *TokenQuotaService) GetDailyUsage(userID string, date time.Time) (*datab
 }
 
 // UpdateUsage updates the token usage for a user
-func (s *TokenQuotaService) UpdateUsage(userID string, tokens int) error {
+func (s *TokenQuotaService) UpdateUsage(userID uuid.UUID, tokens int) error {
 	now := time.Now()
 
 	// Get or create daily usage