77 lines
2.3 KiB
Go
77 lines
2.3 KiB
Go
package postgres
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"strings"
|
|
|
|
"enkly/pkg/pgkit"
|
|
"enkly/services/conversation/internal/domain"
|
|
)
|
|
|
|
// cleanText makes arbitrary tool/model output safe for a Postgres TEXT
|
|
// column, which rejects NUL bytes and invalid UTF-8 outright.
|
|
func cleanText(s string) string {
|
|
if strings.IndexByte(s, 0) >= 0 {
|
|
s = strings.ReplaceAll(s, "\x00", "")
|
|
}
|
|
return strings.ToValidUTF8(s, "\uFFFD")
|
|
}
|
|
|
|
type MessageRepo struct {
|
|
DB *pgkit.DB
|
|
}
|
|
|
|
func NewMessageRepo(db *pgkit.DB) *MessageRepo { return &MessageRepo{DB: db} }
|
|
|
|
func (r *MessageRepo) Add(ctx context.Context, conversationID int64, role domain.Role, content, toolName string) (*domain.Message, error) {
|
|
var tool any
|
|
if toolName != "" {
|
|
tool = toolName
|
|
}
|
|
var m domain.Message
|
|
err := r.DB.QueryRowContext(ctx, `
|
|
INSERT INTO messages (conversation_id, role, content, tool_name)
|
|
VALUES ($1, $2, $3, $4) RETURNING id, conversation_id, role, content, COALESCE(tool_name, ''), created_at`,
|
|
conversationID, string(role), cleanText(content), tool).
|
|
Scan(&m.ID, &m.ConversationID, &m.Role, &m.Content, &m.ToolName, &m.CreatedAt)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("conversation: add message: %w", err)
|
|
}
|
|
return &m, nil
|
|
}
|
|
|
|
func (r *MessageRepo) List(ctx context.Context, conversationID int64) ([]*domain.Message, error) {
|
|
rows, err := r.DB.QueryContext(ctx, `
|
|
SELECT id, conversation_id, role, content, COALESCE(tool_name, ''), created_at
|
|
FROM messages WHERE conversation_id = $1 ORDER BY created_at, id`, conversationID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
|
|
var out []*domain.Message
|
|
for rows.Next() {
|
|
var m domain.Message
|
|
if err := rows.Scan(&m.ID, &m.ConversationID, &m.Role, &m.Content, &m.ToolName, &m.CreatedAt); err != nil {
|
|
return nil, err
|
|
}
|
|
out = append(out, &m)
|
|
}
|
|
return out, rows.Err()
|
|
}
|
|
|
|
func (r *MessageRepo) AddToolCall(ctx context.Context, tc domain.ToolCall) error {
|
|
argsJSON := tc.ArgumentsJSON
|
|
if argsJSON == "" {
|
|
argsJSON = "{}"
|
|
}
|
|
// jsonb rejects the \u0000 escape that encoding/json emits for NUL.
|
|
argsJSON = strings.ReplaceAll(argsJSON, `\u0000`, "")
|
|
_, err := r.DB.ExecContext(ctx, `
|
|
INSERT INTO tool_calls (message_id, tool_name, arguments, result, is_error)
|
|
VALUES ($1, $2, $3::jsonb, $4, $5)`,
|
|
tc.MessageID, tc.ToolName, argsJSON, cleanText(tc.Result), tc.IsError)
|
|
return err
|
|
}
|