backend: create test data for chat completions

go-backend/llms/chat.go

@@ -0,0 +1 @@
+package llms

go-backend/llms/chat_test.go

@@ -0,0 +1,5 @@
+package llms
+
+type ChatCompletionMessage struct {
+
+}

go-backend/models/chat.go

@@ -0,0 +1,22 @@
+package models
+
+import (
+	"time"
+)
+
+type ChatCompletion struct {
+	ID             int       `json:"id"`
+	UserID         int       `json:"user_id"`
+	ConversationID string    `json:"conversation_id"`
+	SequenceNumber int       `json:"sequence_number"`
+	Role           string    `json:"role"`
+	Content        string    `json:"content"`
+	Refusal        *string   `json:"refusal"`
+	Model          string    `json:"model"`
+	Tokens         int       `json:"tokens"`
+	CreatedAt      time.Time `json:"created_at"`
+}
+
+type ChatData struct {
+	ChatCompletions []ChatCompletion `json:"chat_completions"`
+}

go-backend/schema/0029-chat-completions.sql

@@ -0,0 +1,14 @@
+CREATE TABLE IF NOT EXISTS chat_completions (
+       id SERIAL PRIMARY KEY,
+       user_id int,
+       conversation_id UUID,
+       sequence_number INT,  -- position in conversation (1, 2, 3, etc.)
+       role text,
+       content text,
+       refusal text,
+       model text,
+       tokens int,
+       created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+       FOREIGN KEY (user_id) REFERENCES users(id),
+       UNIQUE (conversation_id, sequence_number) -- ensures order numbers are unique per conversation
+)

go-backend/server/database.go

@@ -44,6 +44,7 @@ func ResetDatabase(S *Server) error {
                         DROP TABLE IF EXISTS card_embeddings CASCADE;
                         DROP TABLE IF EXISTS card_chunks CASCADE;
                         DROP TABLE IF EXISTS mailing_list CASCADE;
+                        DROP TABLE IF EXISTS chat_completions CASCADE;
 			CREATE TABLE IF NOT EXISTS migrations (
 				id SERIAL PRIMARY KEY,
 				migration_name VARCHAR(255) NOT NULL,

go-backend/testdata/chat.json

@@ -0,0 +1,124 @@
+{
+  "chat_completions": [
+    {
+      "id": 1,
+      "user_id": 1,
+      "conversation_id": "550e8400-e29b-41d4-a716-446655440000",
+      "sequence_number": 1,
+      "role": "user",
+      "content": "What are the best apples for baking a pie?",
+      "refusal": null,
+      "model": "gpt-4",
+      "tokens": 10,
+      "created_at": "2023-09-20T14:28:00Z"
+    },
+    {
+      "id": 2,
+      "user_id": 1,
+      "conversation_id": "550e8400-e29b-41d4-a716-446655440000",
+      "sequence_number": 2,
+      "role": "assistant",
+      "content": "Granny Smith apples are considered excellent for pies due to their firm texture and tart flavor. Other good options include Honeycrisp, Braeburn, and Northern Spy. These varieties hold their shape well when baked and offer a good balance of sweet and tart flavors.",
+      "refusal": null,
+      "model": "gpt-4",
+      "tokens": 52,
+      "created_at": "2023-09-20T14:28:05Z"
+    },
+    {
+      "id": 3,
+      "user_id": 1,
+      "conversation_id": "550e8400-e29b-41d4-a716-446655440000",
+      "sequence_number": 3,
+      "role": "user",
+      "content": "How many apples do I need for one pie?",
+      "refusal": null,
+      "model": "gpt-4",
+      "tokens": 9,
+      "created_at": "2023-09-20T14:28:30Z"
+    },
+    {
+      "id": 4,
+      "user_id": 1,
+      "conversation_id": "550e8400-e29b-41d4-a716-446655440000",
+      "sequence_number": 4,
+      "role": "assistant",
+      "content": "For a standard 9-inch pie, you'll need about 6-8 medium-sized apples, which equals roughly 3 pounds. It's always good to have an extra apple or two just in case some have bruised spots you need to cut away.",
+      "refusal": null,
+      "model": "gpt-4",
+      "tokens": 45,
+      "created_at": "2023-09-20T14:28:35Z"
+    },
+    {
+      "id": 5,
+      "user_id": 1,
+      "conversation_id": "550e8400-e29b-41d4-a716-446655440000",
+      "sequence_number": 5,
+      "role": "user",
+      "content": "Should I peel the apples?",
+      "refusal": null,
+      "model": "gpt-4",
+      "tokens": 6,
+      "created_at": "2023-09-20T14:29:00Z"
+    },
+    {
+      "id": 6,
+      "user_id": 1,
+      "conversation_id": "550e8400-e29b-41d4-a716-446655440000",
+      "sequence_number": 6,
+      "role": "assistant",
+      "content": "Yes, it's generally recommended to peel apples for pie. The skin can become tough and chewy when baked, affecting the smooth texture of your filling. Make sure to core them and slice them evenly as well for the best results.",
+      "refusal": null,
+      "model": "gpt-4",
+      "tokens": 1,
+      "created_at": "2023-09-20T14:29:05Z"
+    },
+    {
+      "id": 7,
+      "user_id": 1,
+      "conversation_id": "660e8400-e29b-41d4-a716-446655440001",
+      "sequence_number": 1,
+      "role": "user",
+      "content": "Are apples good for weight loss?",
+      "refusal": null,
+      "model": "gpt-4",
+      "tokens": 7,
+      "created_at": "2023-09-20T15:30:00Z"
+    },
+    {
+      "id": 8,
+      "user_id": 1,
+      "conversation_id": "660e8400-e29b-41d4-a716-446655440001",
+      "sequence_number": 2,
+      "role": "assistant",
+      "content": "Yes, apples can be great for weight loss! They're high in fiber, which helps you feel full longer, and relatively low in calories. The average apple contains about 95 calories and provides various beneficial nutrients.",
+      "refusal": null,
+      "model": "gpt-4",
+      "tokens": 39,
+      "created_at": "2023-09-20T15:30:05Z"
+    },
+    {
+      "id": 9,
+      "user_id": 2,
+      "conversation_id": "770e8400-e29b-41d4-a716-446655440002",
+      "sequence_number": 1,
+      "role": "user",
+      "content": "What's the history of apple cultivation?",
+      "refusal": null,
+      "model": "gpt-4",
+      "tokens": 8,
+      "created_at": "2023-09-20T16:15:00Z"
+    },
+    {
+      "id": 10,
+      "user_id": 2,
+      "conversation_id": "770e8400-e29b-41d4-a716-446655440002",
+      "sequence_number": 2,
+      "role": "assistant",
+      "content": "Apple cultivation began in Central Asia thousands of years ago, with Kazakhstan believed to be the origin of the modern apple. The Silk Road helped spread apple cultivation throughout Europe and eventually to other continents. Ancient Romans played a crucial role in developing grafting techniques still used today.",
+      "refusal": null,
+      "model": "gpt-4",
+      "tokens": 55,
+      "created_at": "2023-09-20T16:15:05Z"
+    }
+  ]
+}

go-backend/tests/conftest.go

@@ -1,6 +1,7 @@
 package tests
 
 import (
+	"database/sql"
 	"encoding/json"
 	"fmt"
 	"go-backend/mail"
@@ -45,7 +46,10 @@ func Setup() *server.Server {
 	S.TestInspector = &server.TestInspector{}
 
 	server.RunMigrations(S)
-	importTestData(S)
+	err = importTestData(S)
+	if err != nil {
+		log.Fatal(err)
+	}
 	return S
 }
 
@@ -199,6 +203,10 @@ func importTestData(s *server.Server) error {
 			return err
 		}
 	}
+	log.Printf("starting chat")
+	if err := loadChatData(tx); err != nil {
+		return err
+	}
 
 	tx.Commit()
 	return nil
@@ -480,3 +488,39 @@ func GenerateTestJWT(userID int) (string, error) {
 	}
 	return tokenString, nil
 }
+func loadChatData(tx *sql.Tx) error {
+	// Read the JSON file
+	jsonData, err := os.ReadFile("../testdata/chat.json")
+	if err != nil {
+		return fmt.Errorf("failed to read chat JSON file: %w", err)
+	}
+
+	var chatData models.ChatData
+	if err := json.Unmarshal(jsonData, &chatData); err != nil {
+		return fmt.Errorf("failed to unmarshal chat data: %w", err)
+	}
+
+	// Insert each chat completion
+	for _, chat := range chatData.ChatCompletions {
+		_, err := tx.Exec(`
+            INSERT INTO chat_completions 
+            (user_id, conversation_id, sequence_number, role, content, refusal, model, tokens, created_at)
+            VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)`,
+			chat.UserID,
+			chat.ConversationID,
+			chat.SequenceNumber,
+			chat.Role,
+			chat.Content,
+			chat.Refusal,
+			chat.Model,
+			chat.Tokens,
+			chat.CreatedAt,
+		)
+		log.Printf("chat %v", chat)
+		if err != nil {
+			return fmt.Errorf("failed to insert chat completion: %w", err)
+		}
+	}
+
+	return nil
+}