121 lines
2.9 KiB
Go
121 lines
2.9 KiB
Go
package store
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
|
|
"github.com/jackc/pgx/v5"
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
|
|
"github.com/hhs/camtalk/internal/models"
|
|
)
|
|
|
|
// PgMessageRepository 基于 PostgreSQL 的 MessageRepository 实现。
|
|
type PgMessageRepository struct {
|
|
pool *pgxpool.Pool
|
|
}
|
|
|
|
// NewPgMessageRepository 创建 PgMessageRepository。
|
|
func NewPgMessageRepository(pool *pgxpool.Pool) *PgMessageRepository {
|
|
return &PgMessageRepository{pool: pool}
|
|
}
|
|
|
|
func (r *PgMessageRepository) SaveMessage(ctx context.Context, sessionID string, msg models.Message, tokensUsed int) error {
|
|
_, err := r.pool.Exec(ctx,
|
|
`INSERT INTO messages (session_id, role, content, tokens_used) VALUES ($1, $2, $3, $4)`,
|
|
sessionID, msg.Role, msg.Content, tokensUsed,
|
|
)
|
|
return err
|
|
}
|
|
|
|
func (r *PgMessageRepository) GetMessages(ctx context.Context, sessionID string, limit int, beforeID int64) ([]StoredMessage, error) {
|
|
if limit <= 0 {
|
|
limit = 50
|
|
}
|
|
|
|
var rows []StoredMessage
|
|
var err error
|
|
|
|
if beforeID > 0 {
|
|
rows, err = r.queryMessages(ctx,
|
|
`SELECT id, session_id, role, content, tokens_used, created_at
|
|
FROM messages
|
|
WHERE session_id = $1 AND id < $2
|
|
ORDER BY id DESC
|
|
LIMIT $3`,
|
|
sessionID, beforeID, limit,
|
|
)
|
|
} else {
|
|
rows, err = r.queryMessages(ctx,
|
|
`SELECT id, session_id, role, content, tokens_used, created_at
|
|
FROM messages
|
|
WHERE session_id = $1
|
|
ORDER BY id DESC
|
|
LIMIT $2`,
|
|
sessionID, limit,
|
|
)
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// 反转为升序
|
|
for i, j := 0, len(rows)-1; i < j; i, j = i+1, j-1 {
|
|
rows[i], rows[j] = rows[j], rows[i]
|
|
}
|
|
|
|
return rows, nil
|
|
}
|
|
|
|
func (r *PgMessageRepository) queryMessages(ctx context.Context, query string, args ...any) ([]StoredMessage, error) {
|
|
pgxRows, err := r.pool.Query(ctx, query, args...)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer pgxRows.Close()
|
|
|
|
var messages []StoredMessage
|
|
for pgxRows.Next() {
|
|
var m StoredMessage
|
|
if err := pgxRows.Scan(&m.ID, &m.SessionID, &m.Role, &m.Content, &m.TokensUsed, &m.CreatedAt); err != nil {
|
|
return nil, err
|
|
}
|
|
messages = append(messages, m)
|
|
}
|
|
if err := pgxRows.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
return messages, nil
|
|
}
|
|
|
|
func (r *PgMessageRepository) GetLastMessage(ctx context.Context, sessionID string) (*StoredMessage, error) {
|
|
var m StoredMessage
|
|
err := r.pool.QueryRow(ctx,
|
|
`SELECT id, session_id, role, content, tokens_used, created_at
|
|
FROM messages
|
|
WHERE session_id = $1
|
|
ORDER BY id DESC
|
|
LIMIT 1`,
|
|
sessionID,
|
|
).Scan(&m.ID, &m.SessionID, &m.Role, &m.Content, &m.TokensUsed, &m.CreatedAt)
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return nil, ErrMessageNotFound
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &m, nil
|
|
}
|
|
|
|
func (r *PgMessageRepository) GetMessageCount(ctx context.Context, sessionID string) (int, error) {
|
|
var count int
|
|
err := r.pool.QueryRow(ctx,
|
|
`SELECT COUNT(*) FROM messages WHERE session_id = $1`,
|
|
sessionID,
|
|
).Scan(&count)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
return count, nil
|
|
}
|