diff --git a/main.go b/main.go index eaa8a7a..de76fe3 100644 --- a/main.go +++ b/main.go @@ -97,6 +97,7 @@ func main() { db.Take(&authConfig, "key = ?", "authorization") engine.POST("/v1/*any", func(c *gin.Context) { + begin := time.Now() trackID := uuid.New() // check authorization header if !*noauth { @@ -224,7 +225,7 @@ func main() { if err != nil { log.Println("Failed to read from response tee buffer", err) } - go recordAssistantResponse(contentType, db, trackID, resp) + go recordAssistantResponse(contentType, db, trackID, resp, time.Now().Sub(begin)) }) // --------------------------------- diff --git a/record.go b/record.go index 7cf2ce2..30af0a5 100644 --- a/record.go +++ b/record.go @@ -12,11 +12,12 @@ import ( ) type Record struct { - ID uuid.UUID `gorm:"type:uuid"` - CreatedAt time.Time - IP string - Body string - Response string + ID uuid.UUID `gorm:"type:uuid"` + CreatedAt time.Time + IP string + Body string + Response string + ElapsedTime time.Duration } func recordUserMessage(c *gin.Context, db *gorm.DB, trackID uuid.UUID, body []byte) { @@ -62,7 +63,7 @@ type FetchModeUsage struct { TotalTokens int64 `json:"total_tokens"` } -func recordAssistantResponse(contentType string, db *gorm.DB, trackID uuid.UUID, body []byte) { +func recordAssistantResponse(contentType string, db *gorm.DB, trackID uuid.UUID, body []byte, elapsedTime time.Duration) { result := "" // stream mode if strings.HasPrefix(contentType, "text/event-stream") { @@ -113,6 +114,7 @@ func recordAssistantResponse(contentType string, db *gorm.DB, trackID uuid.UUID, return } record.Response = result + record.ElapsedTime = elapsedTime if db.Save(&record).Error != nil { log.Println("Error to save record:", record) return