Add db message metadata
This commit is contained in:
+18
-1
@@ -1,6 +1,9 @@
|
||||
package db
|
||||
|
||||
import "time"
|
||||
import (
|
||||
"encoding/json/jsontext"
|
||||
"time"
|
||||
)
|
||||
|
||||
// ExecutionStatus is the state of one tool call, as stored in the executions
|
||||
// table (one row for every tool call in an assistant message).
|
||||
@@ -86,12 +89,26 @@ const (
|
||||
// wire-format JSON for the message — the exact JSON object that
|
||||
// appears in the OpenAI Chat Completions messages array.
|
||||
type Message struct {
|
||||
// Content is the model-visible message JSON, without inference
|
||||
// metadata.
|
||||
Content string `json:"content"`
|
||||
// ConversationId identifies the conversation that owns this message
|
||||
// row.
|
||||
ConversationId string `json:"conversation_id"`
|
||||
// CreatedAt is the database timestamp when this message row was
|
||||
// created, not the provider's completion creation time.
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
// Id identifies this database message row, not a completion or request.
|
||||
Id string `json:"id"`
|
||||
// Metadata holds inference JSON decoded from the stored JSONB value.
|
||||
// Nil means no metadata was saved; content edits preserve metadata.
|
||||
Metadata jsontext.Value `json:"metadata"`
|
||||
// Role is the message role stored separately for row queries.
|
||||
Role MessageRole `json:"role"`
|
||||
// Sequence is this row's ordering position within its conversation.
|
||||
Sequence int `json:"sequence"`
|
||||
// UpdatedAt is the database timestamp of this message row's last
|
||||
// update.
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
|
||||
+23
-4
@@ -3,6 +3,7 @@ package db
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json/jsontext"
|
||||
"fmt"
|
||||
"time"
|
||||
"uuid"
|
||||
@@ -301,33 +302,46 @@ func (q *Queries) DeleteConversationId(ctx context.Context, id string) error {
|
||||
}
|
||||
|
||||
const insertMessage = `
|
||||
INSERT INTO messages (id, content, conversation_id, role, sequence)
|
||||
VALUES (?, ?, ?, ?, ?)
|
||||
RETURNING content, conversation_id, created_at, id, role, sequence, updated_at
|
||||
INSERT INTO messages (id, content, conversation_id, metadata, role, sequence)
|
||||
VALUES (?, ?, ?, jsonb(?), ?, ?)
|
||||
RETURNING content, conversation_id, created_at, id, json(metadata), role, sequence, updated_at
|
||||
`
|
||||
|
||||
// InsertMessageParams holds arguments for InsertMessage.
|
||||
type InsertMessageParams struct {
|
||||
// Content is the model-visible message JSON saved in the new row.
|
||||
Content string
|
||||
// ConversationId identifies the conversation that owns the new row.
|
||||
ConversationId string
|
||||
// Metadata is inference JSON converted to JSONB for storage.
|
||||
// Nil stores SQL NULL; malformed JSON causes insertion to fail.
|
||||
Metadata jsontext.Value
|
||||
// Role is the message's role, saved separately for row queries.
|
||||
Role MessageRole
|
||||
// Sequence is the new row's ordering position within the conversation.
|
||||
Sequence int
|
||||
}
|
||||
|
||||
// InsertMessage creates a new message.
|
||||
func (q *Queries) InsertMessage(ctx context.Context, arg InsertMessageParams) (*Message, error) {
|
||||
var metadataArg any
|
||||
if arg.Metadata != nil {
|
||||
metadataArg = string(arg.Metadata)
|
||||
}
|
||||
row := q.db.QueryRowContext(
|
||||
ctx,
|
||||
insertMessage,
|
||||
uuid.New().String(),
|
||||
arg.Content,
|
||||
arg.ConversationId,
|
||||
metadataArg,
|
||||
arg.Role,
|
||||
arg.Sequence,
|
||||
)
|
||||
|
||||
var (
|
||||
m Message
|
||||
metadata []byte
|
||||
createdAt string
|
||||
updatedAt string
|
||||
err error
|
||||
@@ -337,12 +351,14 @@ func (q *Queries) InsertMessage(ctx context.Context, arg InsertMessageParams) (*
|
||||
&m.ConversationId,
|
||||
&createdAt,
|
||||
&m.Id,
|
||||
&metadata,
|
||||
&m.Role,
|
||||
&m.Sequence,
|
||||
&updatedAt,
|
||||
); err != nil {
|
||||
return nil, fmt.Errorf("InsertMessage: %w", err)
|
||||
}
|
||||
m.Metadata = jsontext.Value(metadata)
|
||||
if m.CreatedAt, err = time.Parse(time.RFC3339, createdAt); err != nil {
|
||||
return nil, fmt.Errorf("InsertMessage: parse created_at: %w", err)
|
||||
}
|
||||
@@ -353,7 +369,7 @@ func (q *Queries) InsertMessage(ctx context.Context, arg InsertMessageParams) (*
|
||||
}
|
||||
|
||||
const selectMessagesConversationId = `
|
||||
SELECT content, conversation_id, created_at, id, role, sequence, updated_at
|
||||
SELECT content, conversation_id, created_at, id, json(metadata), role, sequence, updated_at
|
||||
FROM messages
|
||||
WHERE conversation_id = ?
|
||||
ORDER BY sequence
|
||||
@@ -372,6 +388,7 @@ func (q *Queries) SelectMessagesConversationId(ctx context.Context, conversation
|
||||
for rows.Next() {
|
||||
var (
|
||||
m Message
|
||||
metadata []byte
|
||||
createdAt string
|
||||
updatedAt string
|
||||
)
|
||||
@@ -380,12 +397,14 @@ func (q *Queries) SelectMessagesConversationId(ctx context.Context, conversation
|
||||
&m.ConversationId,
|
||||
&createdAt,
|
||||
&m.Id,
|
||||
&metadata,
|
||||
&m.Role,
|
||||
&m.Sequence,
|
||||
&updatedAt,
|
||||
); err != nil {
|
||||
return nil, fmt.Errorf("SelectMessagesConversationId: %w", err)
|
||||
}
|
||||
m.Metadata = jsontext.Value(metadata)
|
||||
if m.CreatedAt, err = time.Parse(time.RFC3339, createdAt); err != nil {
|
||||
return nil, fmt.Errorf("SelectMessagesConversationId: parse created_at: %w", err)
|
||||
}
|
||||
|
||||
@@ -2,8 +2,11 @@ package db
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json/jsontext"
|
||||
"encoding/json/v2"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"reflect"
|
||||
"testing"
|
||||
|
||||
"github.com/mattn/go-sqlite3"
|
||||
@@ -177,6 +180,7 @@ func testMessaging(t *testing.T, ctx context.Context, db *SQLite) {
|
||||
t.Fatalf("InsertMessage: %v", err)
|
||||
}
|
||||
|
||||
metadata := `{"id":"completion","usage":{"prompt_tokens":0}}`
|
||||
// Insert an assistant message with tool calls.
|
||||
seq, err = qw.NextSequence(ctx, c.Id)
|
||||
if err != nil {
|
||||
@@ -186,6 +190,7 @@ func testMessaging(t *testing.T, ctx context.Context, db *SQLite) {
|
||||
ConversationId: c.Id,
|
||||
Sequence: seq,
|
||||
Role: MessageRoleAssistant,
|
||||
Metadata: jsontext.Value(metadata),
|
||||
Content: `{"role":"assistant","content":null,"tool_calls":[{"id":"call_abc","type":"function","function":{"name":"get_weather","arguments":"{\"location\":\"NYC\"}"}}]}`,
|
||||
})
|
||||
if err != nil {
|
||||
@@ -223,6 +228,22 @@ func testMessaging(t *testing.T, ctx context.Context, db *SQLite) {
|
||||
t.Errorf("expected last message role tool, got %s", messages[3].Role)
|
||||
}
|
||||
|
||||
for i, message := range messages {
|
||||
if i == 2 {
|
||||
if message.Metadata == nil || string(message.Metadata) != metadata {
|
||||
t.Errorf("assistant metadata = %v", message.Metadata)
|
||||
}
|
||||
} else if message.Metadata != nil {
|
||||
t.Errorf("message %d metadata = %v, want nil", i, message.Metadata)
|
||||
}
|
||||
}
|
||||
if err := qw.UpdateMessageContent(ctx, UpdateMessageContentParams{
|
||||
Id: messages[2].Id,
|
||||
Content: `{"role":"assistant","content":"edited"}`,
|
||||
}); err != nil {
|
||||
t.Fatalf("UpdateMessageContent: %v", err)
|
||||
}
|
||||
|
||||
// Test UpdateMessageContent.
|
||||
if err := qw.UpdateMessageContent(ctx, UpdateMessageContentParams{
|
||||
Id: toolMsg.Id,
|
||||
@@ -235,9 +256,116 @@ func testMessaging(t *testing.T, ctx context.Context, db *SQLite) {
|
||||
if err != nil {
|
||||
t.Fatalf("SelectMessagesConversationId: %v", err)
|
||||
}
|
||||
if updated[2].Metadata == nil || string(updated[2].Metadata) != metadata {
|
||||
t.Errorf("edited assistant metadata = %v", updated[2].Metadata)
|
||||
}
|
||||
if updated[3].Content != `{"role":"tool","content":"{\"temperature\":75}","tool_call_id":"call_abc"}` {
|
||||
t.Errorf("updated content mismatch: %s", updated[3].Content)
|
||||
}
|
||||
for _, tt := range []struct {
|
||||
name string
|
||||
metadata jsontext.Value
|
||||
wantErr bool
|
||||
}{
|
||||
{name: "absent"},
|
||||
{name: "JSON null", metadata: jsontext.Value(`null`)},
|
||||
{name: "empty object", metadata: jsontext.Value(`{}`)},
|
||||
{name: "empty array", metadata: jsontext.Value(`[]`)},
|
||||
{name: "malformed", metadata: jsontext.Value(`{`), wantErr: true},
|
||||
{name: "empty", metadata: jsontext.Value{}, wantErr: true},
|
||||
{name: "reported zero", metadata: jsontext.Value(`{"usage":{"total_tokens":0}}`)},
|
||||
{name: "opaque JSON", metadata: jsontext.Value(`{"other":{"value":42}}`)},
|
||||
} {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
seq, err := qw.NextSequence(ctx, c.Id)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
inserted, err := qw.InsertMessage(ctx, InsertMessageParams{
|
||||
Content: `{"role":"assistant","content":"answer"}`,
|
||||
ConversationId: c.Id,
|
||||
Metadata: tt.metadata,
|
||||
Role: MessageRoleAssistant,
|
||||
Sequence: seq,
|
||||
})
|
||||
if tt.wantErr {
|
||||
if err == nil {
|
||||
t.Fatal("invalid metadata was accepted")
|
||||
}
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var storedType string
|
||||
if err := db.Reader.QueryRowContext(ctx,
|
||||
"SELECT typeof(metadata) FROM messages WHERE id = ?",
|
||||
inserted.Id,
|
||||
).Scan(&storedType); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
wantType := "blob"
|
||||
if tt.metadata == nil {
|
||||
wantType = "null"
|
||||
}
|
||||
if storedType != wantType {
|
||||
t.Errorf("storage type = %s, want %s", storedType, wantType)
|
||||
}
|
||||
selected, err := qr.SelectMessagesConversationId(ctx, c.Id)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, message := range []Message{*inserted, selected[len(selected)-1]} {
|
||||
if (message.Metadata == nil) != (tt.metadata == nil) ||
|
||||
(tt.metadata != nil && string(message.Metadata) != string(tt.metadata)) {
|
||||
t.Fatalf("metadata = %v, want %v", message.Metadata, tt.metadata)
|
||||
}
|
||||
encoded, err := json.Marshal(message)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var row map[string]any
|
||||
if err := json.Unmarshal(encoded, &row); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
value, exists := row["metadata"]
|
||||
var wantValue any
|
||||
if tt.metadata != nil {
|
||||
if err := json.Unmarshal(tt.metadata, &wantValue); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if !exists || !reflect.DeepEqual(value, wantValue) {
|
||||
t.Errorf("row JSON = %s", encoded)
|
||||
}
|
||||
}
|
||||
if err := qw.DeleteMessageId(ctx, inserted.Id); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
})
|
||||
}
|
||||
tx, err := db.Writer.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, err = New(tx).InsertMessage(ctx, InsertMessageParams{
|
||||
Content: `{"role":"assistant","content":"rolled back"}`,
|
||||
ConversationId: c.Id,
|
||||
Metadata: jsontext.Value(metadata),
|
||||
Role: MessageRoleAssistant,
|
||||
Sequence: 100,
|
||||
})
|
||||
if err != nil {
|
||||
tx.Rollback()
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := tx.Rollback(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
selected, err := qr.SelectMessagesConversationId(ctx, c.Id)
|
||||
if err != nil || len(selected) != 4 {
|
||||
t.Fatalf("rollback messages = %d, error = %v", len(selected), err)
|
||||
}
|
||||
}
|
||||
|
||||
func testConversationTimestamp(t *testing.T, ctx context.Context, database *SQLite) {
|
||||
|
||||
@@ -55,7 +55,8 @@ CREATE INDEX IF NOT EXISTS idx_conv_user ON conversations(user_id);
|
||||
-- Messages are ordered entries within a conversation. The content column
|
||||
-- stores the wire-format JSON for the message — the exact JSON object
|
||||
-- that appears in the OpenAI Chat Completions messages array. The role
|
||||
-- column is extracted for querying without JSON parsing. The sequence
|
||||
-- column is extracted for querying without JSON parsing. Nullable metadata
|
||||
-- stores inference JSONB separately from model-visible content. The sequence
|
||||
-- column defines ordering within the conversation. Deleting a message
|
||||
-- cascades to nothing further (all message data is in the row).
|
||||
CREATE TABLE IF NOT EXISTS messages (
|
||||
@@ -63,6 +64,7 @@ CREATE TABLE IF NOT EXISTS messages (
|
||||
content TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
|
||||
id TEXT PRIMARY KEY,
|
||||
metadata BLOB CHECK (metadata IS NULL OR json_valid(metadata, 8)),
|
||||
role TEXT NOT NULL,
|
||||
sequence INTEGER NOT NULL,
|
||||
updated_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
|
||||
@@ -149,6 +151,8 @@ CREATE TABLE IF NOT EXISTS executions (
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_executions_message ON executions(message_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_executions_pending ON executions(status) WHERE status = 'pending';
|
||||
-- Prevent duplicate execution rows for the same tool call in a message.
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS idx_executions_call ON executions(message_id, tool_call_id);
|
||||
|
||||
-- Default user. Temporary workaround until authentication is built.
|
||||
INSERT INTO users (id, name, password, role)
|
||||
|
||||
Reference in New Issue
Block a user