From 8889da73abb0ee34b5a85b941536f9b26dd6b047 Mon Sep 17 00:00:00 2001
From: heimoshuiyu <heimoshuiyu@gmail.com>
Date: Tue, 11 Jul 2023 14:51:19 +0800
Subject: [PATCH] record full request body

---
 main.go   | 10 +++++-----
 record.go | 44 ++++++--------------------------------------
 2 files changed, 11 insertions(+), 43 deletions(-)

diff --git a/main.go b/main.go
index c1df708..0d32e4e 100644
--- a/main.go
+++ b/main.go
@@ -46,7 +46,7 @@ func main() {
 	}
 
 	db.AutoMigrate(&OPENAI_UPSTREAM{})
-	db.AutoMigrate(&UserMessage{})
+	db.AutoMigrate(&RequestRecord{})
 	log.Println("Auto migrate database done")
 
 	if *addMode {
@@ -306,18 +306,18 @@ func main() {
 			"message": "success",
 		})
 	})
-	engine.GET("/admin/user_messages", func(c *gin.Context) {
+	engine.GET("/admin/request_records", func(c *gin.Context) {
 		// check authorization headers
 		if handleAuth(c) != nil {
 			return
 		}
-		userMessages := []UserMessage{}
-		err := db.Order("id desc").Limit(100).Find(&userMessages).Error
+		requestRecords := []RequestRecord{}
+		err := db.Order("id desc").Limit(100).Find(&requestRecords).Error
 		if err != nil {
 			c.AbortWithError(502, err)
 			return
 		}
-		c.JSON(200, userMessages)
+		c.JSON(200, requestRecords)
 	})
 	engine.Run(*listenAddr)
 }
diff --git a/record.go b/record.go
index 08b56e4..0ddbb93 100644
--- a/record.go
+++ b/record.go
@@ -1,51 +1,19 @@
 package main
 
 import (
-	"encoding/json"
-	"log"
-	"strings"
-
 	"github.com/gin-gonic/gin"
 	"gorm.io/gorm"
 )
 
-type UserMessage struct {
+type RequestRecord struct {
 	gorm.Model
-	ModelName string
-	Content   string
-}
-
-// sturcture to parse request
-type ChatRequest struct {
-	Model    string    `json:"model"`
-	Messages []Message `json:"messages"`
-}
-type Message struct {
-	Content string `json:"content"`
+	Body string
 }
 
 func recordUserMessage(c *gin.Context, db *gorm.DB, body []byte) {
-	bodyJson := ChatRequest{}
-	err := json.Unmarshal(body, &bodyJson)
-	if err != nil {
-		c.AbortWithError(502, err)
-		return
+	bodyStr := string(body)
+	requestRecord := RequestRecord{
+		Body: bodyStr,
 	}
-	model := bodyJson.Model
-	if !strings.HasPrefix(model, "gpt-") {
-		return
-	}
-	// get message content
-	if len(bodyJson.Messages) == 0 {
-		return
-	}
-	content := bodyJson.Messages[len(bodyJson.Messages)-1].Content
-
-	log.Println("Record user message", model, content)
-
-	userMessage := UserMessage{
-		ModelName: model,
-		Content:   content,
-	}
-	db.Create(&userMessage)
+	db.Create(&requestRecord)
 }