You've already forked agentic-coding-workflow
1154 lines
37 KiB
Go
1154 lines
37 KiB
Go
// Package chat — repository.go 是 sqlc 生成层之上的薄包装:
|
|
// - pgtype.* ↔ 领域类型互转
|
|
// - pgx.ErrNoRows → errs.CodeChat*NotFound
|
|
// - PG 唯一约束冲突 → 可识别的 AppError
|
|
// - WithTx 提供事务作用域
|
|
package chat
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"math/big"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/jackc/pgerrcode"
|
|
"github.com/jackc/pgx/v5"
|
|
"github.com/jackc/pgx/v5/pgconn"
|
|
"github.com/jackc/pgx/v5/pgtype"
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
|
|
"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
|
|
chatsqlc "github.com/yan1h/agent-coding-workflow/internal/chat/sqlc"
|
|
)
|
|
|
|
// Postgres epoch: 2000-01-01 00:00:00 UTC
|
|
var pgEpoch = time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC)
|
|
|
|
// Repository 是 chat 模块对 PG 的全部依赖。
|
|
// Service 单元测试通过手写 fakeRepo 满足该接口。
|
|
type Repository interface {
|
|
// ===== Conversation =====
|
|
InsertConversation(ctx context.Context, p InsertConversationParams) (*Conversation, error)
|
|
GetConversation(ctx context.Context, id uuid.UUID) (*Conversation, error)
|
|
ListConversationsByUser(ctx context.Context, userID uuid.UUID, f ConversationFilter) ([]Conversation, error)
|
|
SetConversationTitleManual(ctx context.Context, id uuid.UUID, title *string) error
|
|
UpdateConversationTitle(ctx context.Context, id uuid.UUID, title *string) error
|
|
SoftDeleteConversation(ctx context.Context, id uuid.UUID) error
|
|
|
|
// ===== Message =====
|
|
InsertMessage(ctx context.Context, conversationID uuid.UUID, role MessageRole, content string, status MessageStatus) (*Message, error)
|
|
GetMessage(ctx context.Context, id int64) (*Message, error)
|
|
GetFirstUserMessage(ctx context.Context, conversationID uuid.UUID) (*Message, error)
|
|
ListMessagesByConversation(ctx context.Context, conversationID uuid.UUID, beforeID *int64, limit int) ([]Message, error)
|
|
ListMessagesAscending(ctx context.Context, conversationID uuid.UUID) ([]Message, error)
|
|
UpdateMessageOK(ctx context.Context, p UpdateMessageOKParams) error
|
|
UpdateMessageError(ctx context.Context, id int64, errMsg *string) error
|
|
UpdateMessageCancelled(ctx context.Context, id int64) error
|
|
MarkPendingAsErrorOnStartup(ctx context.Context) error
|
|
DeleteMessage(ctx context.Context, id int64) error
|
|
|
|
// ===== Attachment =====
|
|
InsertAttachment(ctx context.Context, p InsertAttachmentParams) (*Attachment, error)
|
|
GetAttachment(ctx context.Context, id uuid.UUID) (*Attachment, error)
|
|
AttachToMessage(ctx context.Context, ids []uuid.UUID, messageID int64) error
|
|
ListAttachmentsForMessage(ctx context.Context, messageID int64) ([]Attachment, error)
|
|
ListOrphanAttachmentsForUser(ctx context.Context, userID uuid.UUID, ids []uuid.UUID) ([]Attachment, error)
|
|
ListExpiredOrphans(ctx context.Context, olderThan time.Duration) ([]AttachmentRef, error)
|
|
ListAttachmentsForSoftDeletedConversations(ctx context.Context, olderThan time.Duration) ([]AttachmentRef, error)
|
|
DeleteAttachment(ctx context.Context, id uuid.UUID) error
|
|
|
|
// ===== PromptTemplate =====
|
|
InsertPromptTemplate(ctx context.Context, p InsertTemplateParams) (*PromptTemplate, error)
|
|
GetPromptTemplate(ctx context.Context, id uuid.UUID) (*PromptTemplate, error)
|
|
ListPromptTemplatesForUser(ctx context.Context, userID uuid.UUID) ([]PromptTemplate, error)
|
|
UpdatePromptTemplate(ctx context.Context, id uuid.UUID, name, content string) error
|
|
DeletePromptTemplate(ctx context.Context, id uuid.UUID) error
|
|
|
|
// ===== LLMEndpoint =====
|
|
InsertLLMEndpoint(ctx context.Context, p InsertEndpointParams) (*LLMEndpoint, error)
|
|
GetLLMEndpoint(ctx context.Context, id uuid.UUID) (*LLMEndpoint, error)
|
|
ListLLMEndpoints(ctx context.Context) ([]LLMEndpoint, error)
|
|
UpdateLLMEndpoint(ctx context.Context, id uuid.UUID, in UpdateEndpointInput) error
|
|
DeleteLLMEndpoint(ctx context.Context, id uuid.UUID) error
|
|
|
|
// ===== LLMModel =====
|
|
InsertLLMModel(ctx context.Context, p InsertModelParams) (*LLMModel, error)
|
|
GetLLMModel(ctx context.Context, id uuid.UUID) (*LLMModel, error)
|
|
GetTitleGeneratorModel(ctx context.Context) (*LLMModel, error)
|
|
ListLLMModels(ctx context.Context, onlyEnabled bool) ([]LLMModel, error)
|
|
UpdateLLMModel(ctx context.Context, id uuid.UUID, in UpdateModelInput) error
|
|
DeleteLLMModel(ctx context.Context, id uuid.UUID) error
|
|
|
|
// ===== Usage =====
|
|
InsertUsage(ctx context.Context, r UsageRecord) error
|
|
SummarizeUsageByUser(ctx context.Context, from, to time.Time) ([]UserUsageRow, error)
|
|
SummarizeUsageByModel(ctx context.Context, from, to time.Time) ([]ModelUsageRow, error)
|
|
SummarizeUsageByDay(ctx context.Context, from, to time.Time) ([]DayUsageRow, error)
|
|
SummarizeUserDaily(ctx context.Context, userID uuid.UUID, from, to time.Time) ([]DayUsageRow, error)
|
|
|
|
// ===== Tx =====
|
|
WithTx(ctx context.Context) (Repository, Tx, error)
|
|
}
|
|
|
|
// Tx 是暴露给调用方的事务句柄;调用方必须在 defer 或 cleanup 中调用 Rollback,
|
|
// 成功后调用 Commit。
|
|
type Tx interface {
|
|
Commit(ctx context.Context) error
|
|
Rollback(ctx context.Context) error
|
|
}
|
|
|
|
// ===== Repo-layer input structs =====
|
|
|
|
// InsertConversationParams 是 InsertConversation 的参数。
|
|
type InsertConversationParams struct {
|
|
ID uuid.UUID
|
|
UserID uuid.UUID
|
|
ProjectID *uuid.UUID
|
|
WorkspaceID *uuid.UUID
|
|
IssueID *uuid.UUID
|
|
ModelID uuid.UUID
|
|
SystemPrompt string
|
|
Title *string
|
|
}
|
|
|
|
// UpdateMessageOKParams 是 UpdateMessageOK 的参数。
|
|
type UpdateMessageOKParams struct {
|
|
ID int64
|
|
Content string
|
|
Thinking *string
|
|
PromptTokens *int32
|
|
CompletionTokens *int32
|
|
ThinkingTokens *int32
|
|
}
|
|
|
|
// InsertAttachmentParams 是 InsertAttachment 的参数。
|
|
type InsertAttachmentParams struct {
|
|
ID uuid.UUID
|
|
UserID uuid.UUID
|
|
Filename string
|
|
MimeType string
|
|
SizeBytes int64
|
|
SHA256 string
|
|
StoragePath string
|
|
}
|
|
|
|
// InsertTemplateParams 是 InsertPromptTemplate 的参数。
|
|
type InsertTemplateParams struct {
|
|
ID uuid.UUID
|
|
Name string
|
|
Content string
|
|
Scope TemplateScope
|
|
OwnerID *uuid.UUID
|
|
}
|
|
|
|
// InsertEndpointParams 是 InsertLLMEndpoint 的参数。
|
|
type InsertEndpointParams struct {
|
|
ID uuid.UUID
|
|
Provider LLMProvider
|
|
DisplayName string
|
|
BaseURL string
|
|
APIKeyEncrypted []byte
|
|
CreatedBy uuid.UUID
|
|
}
|
|
|
|
// InsertModelParams 是 InsertLLMModel 的参数。
|
|
type InsertModelParams struct {
|
|
ID uuid.UUID
|
|
EndpointID uuid.UUID
|
|
ModelID string
|
|
DisplayName string
|
|
Capabilities ModelCapabilities
|
|
ContextWindow int
|
|
MaxOutputTokens int
|
|
PromptPrice float64
|
|
CompletionPrice float64
|
|
ThinkingPrice float64
|
|
IsTitleGenerator bool
|
|
Enabled bool
|
|
SortOrder int
|
|
}
|
|
|
|
// AttachmentRef は 孤立 attachment 清理时使用的轻量引用。
|
|
type AttachmentRef struct {
|
|
ID uuid.UUID
|
|
StoragePath string
|
|
}
|
|
|
|
// ===== compile-time assertion =====
|
|
var _ Repository = (*pgRepo)(nil)
|
|
|
|
// pgRepo 是 Repository 的 Postgres 实现。
|
|
type pgRepo struct {
|
|
pool *pgxpool.Pool
|
|
tx pgx.Tx // 非 nil 时处于事务中
|
|
q *chatsqlc.Queries
|
|
}
|
|
|
|
// NewRepository 用 pgxpool 构造 Repository。
|
|
func NewRepository(pool *pgxpool.Pool) Repository {
|
|
return &pgRepo{pool: pool, q: chatsqlc.New(pool)}
|
|
}
|
|
|
|
// WithTx 开启一个 pgx 事务,返回绑定在该事务的新 Repository 和 Tx 句柄。
|
|
func (r *pgRepo) WithTx(ctx context.Context) (Repository, Tx, error) {
|
|
tx, err := r.pool.Begin(ctx)
|
|
if err != nil {
|
|
return nil, nil, errs.Wrap(err, errs.CodeInternal, "begin tx")
|
|
}
|
|
return &pgRepo{pool: r.pool, tx: tx, q: chatsqlc.New(tx)}, tx, nil
|
|
}
|
|
|
|
// ===== pgx.Tx 直接实现 Tx 接口 =====
|
|
// pgx.Tx 已有 Commit/Rollback(ctx) error,满足接口要求。
|
|
|
|
// ===== 类型转换 helpers =====
|
|
|
|
func toPgUUID(u uuid.UUID) pgtype.UUID { return pgtype.UUID{Bytes: u, Valid: true} }
|
|
|
|
func fromPgUUID(p pgtype.UUID) uuid.UUID {
|
|
if !p.Valid {
|
|
return uuid.Nil
|
|
}
|
|
return uuid.UUID(p.Bytes)
|
|
}
|
|
|
|
func toPgUUIDPtr(u *uuid.UUID) pgtype.UUID {
|
|
if u == nil {
|
|
return pgtype.UUID{Valid: false}
|
|
}
|
|
return pgtype.UUID{Bytes: *u, Valid: true}
|
|
}
|
|
|
|
func fromPgUUIDPtr(p pgtype.UUID) *uuid.UUID {
|
|
if !p.Valid {
|
|
return nil
|
|
}
|
|
u := uuid.UUID(p.Bytes)
|
|
return &u
|
|
}
|
|
|
|
func fromPgTime(t pgtype.Timestamptz) time.Time {
|
|
if !t.Valid {
|
|
return time.Time{}
|
|
}
|
|
return t.Time
|
|
}
|
|
|
|
func fromPgTimePtr(t pgtype.Timestamptz) *time.Time {
|
|
if !t.Valid {
|
|
return nil
|
|
}
|
|
v := t.Time
|
|
return &v
|
|
}
|
|
|
|
func toPgTimestamptz(t time.Time) pgtype.Timestamptz {
|
|
return pgtype.Timestamptz{Time: t, Valid: true}
|
|
}
|
|
|
|
// numericToFloat64 は pgtype.Numeric → float64 変換。
|
|
// Float64Value() は (Float8, error) を返す; Float8.Float64 を使う。
|
|
func numericToFloat64(n pgtype.Numeric) float64 {
|
|
if !n.Valid {
|
|
return 0
|
|
}
|
|
f8, err := n.Float64Value()
|
|
if err != nil {
|
|
return 0
|
|
}
|
|
return f8.Float64
|
|
}
|
|
|
|
// float64ToNumeric は float64 → pgtype.Numeric 変換。
|
|
// pgtype.Numeric は Int (big.Int) + Exp (int32) で表現するため、
|
|
// fmt.Sprintf で文字列変換してから ScanScientific で読む。
|
|
func float64ToNumeric(f float64) pgtype.Numeric {
|
|
var n pgtype.Numeric
|
|
if err := n.ScanScientific(fmt.Sprintf("%.10e", f)); err != nil {
|
|
// fallback: 整数部のみ
|
|
bi := new(big.Int)
|
|
bi.SetInt64(int64(f))
|
|
return pgtype.Numeric{Int: bi, Exp: 0, Valid: true}
|
|
}
|
|
return n
|
|
}
|
|
|
|
// intervalToTime は date_trunc('day',...) が pgtype.Interval として返ってくる
|
|
// sqlc 型推論ドリフトへの対処。Postgres は timestamptz を binary で
|
|
// 「2000-01-01 エポックからのマイクロ秒」として送るが、pgx が pgtype.Interval
|
|
// へスキャンした場合 Microseconds フィールドに入る。
|
|
// Drift 注記: sqlc v1.31.1 は date_trunc の返り型を誤って interval と推論する。
|
|
func intervalToTime(iv pgtype.Interval) time.Time {
|
|
if !iv.Valid {
|
|
return time.Time{}
|
|
}
|
|
// Microseconds はここでは「Postgres epoch からのマイクロ秒数」として解釈する。
|
|
return pgEpoch.Add(time.Duration(iv.Microseconds) * time.Microsecond)
|
|
}
|
|
|
|
// ===== PG エラー変換 =====
|
|
|
|
// translatePgError は PG 固有のエラーをドメイン AppError に変換する。
|
|
// 対象外のエラーはそのまま返す。
|
|
func translatePgError(err error) error {
|
|
var pgErr *pgconn.PgError
|
|
if !errors.As(err, &pgErr) {
|
|
return err
|
|
}
|
|
if pgErr.Code == pgerrcode.UniqueViolation {
|
|
switch pgErr.ConstraintName {
|
|
case "messages_one_pending_per_conv":
|
|
return errs.Wrap(err, errs.CodeChatConcurrentMessage, "another message is already streaming")
|
|
case "llm_models_one_title_generator":
|
|
return errs.Wrap(err, errs.CodeConflict, "another model is already the title generator")
|
|
}
|
|
}
|
|
return err
|
|
}
|
|
|
|
// ===== Row → Domain converters =====
|
|
|
|
func convFromRow(r chatsqlc.Conversation) *Conversation {
|
|
return &Conversation{
|
|
ID: fromPgUUID(r.ID),
|
|
UserID: fromPgUUID(r.UserID),
|
|
ProjectID: fromPgUUIDPtr(r.ProjectID),
|
|
WorkspaceID: fromPgUUIDPtr(r.WorkspaceID),
|
|
IssueID: fromPgUUIDPtr(r.IssueID),
|
|
ModelID: fromPgUUID(r.ModelID),
|
|
SystemPrompt: r.SystemPrompt,
|
|
Title: r.Title,
|
|
TitleGeneratedAt: fromPgTimePtr(r.TitleGeneratedAt),
|
|
CreatedAt: fromPgTime(r.CreatedAt),
|
|
UpdatedAt: fromPgTime(r.UpdatedAt),
|
|
DeletedAt: fromPgTimePtr(r.DeletedAt),
|
|
}
|
|
}
|
|
|
|
func messageFromRow(r chatsqlc.Message) *Message {
|
|
var ptrs *int
|
|
var ctrs *int
|
|
var ttrs *int
|
|
if r.PromptTokens != nil {
|
|
v := int(*r.PromptTokens)
|
|
ptrs = &v
|
|
}
|
|
if r.CompletionTokens != nil {
|
|
v := int(*r.CompletionTokens)
|
|
ctrs = &v
|
|
}
|
|
if r.ThinkingTokens != nil {
|
|
v := int(*r.ThinkingTokens)
|
|
ttrs = &v
|
|
}
|
|
return &Message{
|
|
ID: r.ID,
|
|
ConversationID: fromPgUUID(r.ConversationID),
|
|
Role: MessageRole(r.Role),
|
|
Content: r.Content,
|
|
Thinking: r.Thinking,
|
|
ToolCalls: r.ToolCalls,
|
|
Status: MessageStatus(r.Status),
|
|
ErrorMessage: r.ErrorMessage,
|
|
PromptTokens: ptrs,
|
|
CompletionTokens: ctrs,
|
|
ThinkingTokens: ttrs,
|
|
CreatedAt: fromPgTime(r.CreatedAt),
|
|
UpdatedAt: fromPgTime(r.UpdatedAt),
|
|
}
|
|
}
|
|
|
|
func attachmentFromRow(r chatsqlc.MessageAttachment) *Attachment {
|
|
return &Attachment{
|
|
ID: fromPgUUID(r.ID),
|
|
UserID: fromPgUUID(r.UserID),
|
|
MessageID: r.MessageID,
|
|
Filename: r.Filename,
|
|
MimeType: r.MimeType,
|
|
SizeBytes: r.SizeBytes,
|
|
SHA256: r.Sha256,
|
|
StoragePath: r.StoragePath,
|
|
CreatedAt: fromPgTime(r.CreatedAt),
|
|
}
|
|
}
|
|
|
|
func templateFromRow(r chatsqlc.PromptTemplate) *PromptTemplate {
|
|
return &PromptTemplate{
|
|
ID: fromPgUUID(r.ID),
|
|
Name: r.Name,
|
|
Content: r.Content,
|
|
Scope: TemplateScope(r.Scope),
|
|
OwnerID: fromPgUUIDPtr(r.OwnerID),
|
|
CreatedAt: fromPgTime(r.CreatedAt),
|
|
UpdatedAt: fromPgTime(r.UpdatedAt),
|
|
}
|
|
}
|
|
|
|
func endpointFromRow(r chatsqlc.LlmEndpoint) *LLMEndpoint {
|
|
return &LLMEndpoint{
|
|
ID: fromPgUUID(r.ID),
|
|
Provider: LLMProvider(r.Provider),
|
|
DisplayName: r.DisplayName,
|
|
BaseURL: r.BaseUrl,
|
|
APIKeyEncrypted: r.ApiKeyEncrypted,
|
|
CreatedBy: fromPgUUID(r.CreatedBy),
|
|
CreatedAt: fromPgTime(r.CreatedAt),
|
|
UpdatedAt: fromPgTime(r.UpdatedAt),
|
|
}
|
|
}
|
|
|
|
func modelFromRow(r chatsqlc.LlmModel) *LLMModel {
|
|
var caps ModelCapabilities
|
|
if len(r.Capabilities) > 0 {
|
|
_ = json.Unmarshal(r.Capabilities, &caps)
|
|
}
|
|
return &LLMModel{
|
|
ID: fromPgUUID(r.ID),
|
|
EndpointID: fromPgUUID(r.EndpointID),
|
|
ModelID: r.ModelID,
|
|
DisplayName: r.DisplayName,
|
|
Capabilities: caps,
|
|
ContextWindow: int(r.ContextWindow),
|
|
MaxOutputTokens: int(r.MaxOutputTokens),
|
|
PromptPricePerMillionUSD: numericToFloat64(r.PromptPricePerMillionUsd),
|
|
CompletionPricePerMillionUSD: numericToFloat64(r.CompletionPricePerMillionUsd),
|
|
ThinkingPricePerMillionUSD: numericToFloat64(r.ThinkingPricePerMillionUsd),
|
|
IsTitleGenerator: r.IsTitleGenerator,
|
|
Enabled: r.Enabled,
|
|
SortOrder: int(r.SortOrder),
|
|
CreatedAt: fromPgTime(r.CreatedAt),
|
|
UpdatedAt: fromPgTime(r.UpdatedAt),
|
|
}
|
|
}
|
|
|
|
func usageFromRow(r chatsqlc.LlmUsage) *UsageRecord {
|
|
return &UsageRecord{
|
|
ID: r.ID,
|
|
ConversationID: fromPgUUID(r.ConversationID),
|
|
MessageID: r.MessageID,
|
|
UserID: fromPgUUID(r.UserID),
|
|
EndpointID: fromPgUUID(r.EndpointID),
|
|
ModelID: fromPgUUID(r.ModelID),
|
|
PromptTokens: int(r.PromptTokens),
|
|
CompletionTokens: int(r.CompletionTokens),
|
|
ThinkingTokens: int(r.ThinkingTokens),
|
|
CostUSD: numericToFloat64(r.CostUsd),
|
|
CreatedAt: fromPgTime(r.CreatedAt),
|
|
}
|
|
}
|
|
|
|
func userUsageRowFromRow(r chatsqlc.SummarizeUsageByUserRow) UserUsageRow {
|
|
return UserUsageRow{
|
|
UserID: fromPgUUID(r.UserID),
|
|
PromptTokens: int(r.PromptTokens),
|
|
CompletionTokens: int(r.CompletionTokens),
|
|
ThinkingTokens: int(r.ThinkingTokens),
|
|
CostUSD: numericToFloat64(r.CostUsd),
|
|
}
|
|
}
|
|
|
|
func modelUsageRowFromRow(r chatsqlc.SummarizeUsageByModelRow) ModelUsageRow {
|
|
return ModelUsageRow{
|
|
ModelID: fromPgUUID(r.ModelID),
|
|
PromptTokens: int(r.PromptTokens),
|
|
CompletionTokens: int(r.CompletionTokens),
|
|
ThinkingTokens: int(r.ThinkingTokens),
|
|
CostUSD: numericToFloat64(r.CostUsd),
|
|
}
|
|
}
|
|
|
|
func dayUsageRowFromRow(r chatsqlc.SummarizeUsageByDayRow) DayUsageRow {
|
|
return DayUsageRow{
|
|
// Drift: sqlc 把 date_trunc('day', timestamptz) 推断为 pgtype.Interval。
|
|
// pgx 在 binary 模式下把 timestamptz 值的 "Postgres epoch + microseconds"
|
|
// 送入 Interval.Microseconds 字段。用 pgEpoch + Duration 还原。
|
|
Day: intervalToTime(r.Day),
|
|
PromptTokens: int(r.PromptTokens),
|
|
CompletionTokens: int(r.CompletionTokens),
|
|
ThinkingTokens: int(r.ThinkingTokens),
|
|
CostUSD: numericToFloat64(r.CostUsd),
|
|
}
|
|
}
|
|
|
|
// ===== Conversation =====
|
|
|
|
func (r *pgRepo) InsertConversation(ctx context.Context, p InsertConversationParams) (*Conversation, error) {
|
|
row, err := r.q.InsertConversation(ctx, chatsqlc.InsertConversationParams{
|
|
ID: toPgUUID(p.ID),
|
|
UserID: toPgUUID(p.UserID),
|
|
ProjectID: toPgUUIDPtr(p.ProjectID),
|
|
WorkspaceID: toPgUUIDPtr(p.WorkspaceID),
|
|
IssueID: toPgUUIDPtr(p.IssueID),
|
|
ModelID: toPgUUID(p.ModelID),
|
|
SystemPrompt: p.SystemPrompt,
|
|
Title: p.Title,
|
|
})
|
|
if err != nil {
|
|
return nil, errs.Wrap(translatePgError(err), errs.CodeInternal, "insert conversation")
|
|
}
|
|
return convFromRow(row), nil
|
|
}
|
|
|
|
func (r *pgRepo) GetConversation(ctx context.Context, id uuid.UUID) (*Conversation, error) {
|
|
row, err := r.q.GetConversation(ctx, toPgUUID(id))
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return nil, errs.New(errs.CodeChatConversationNotFound, "conversation not found")
|
|
}
|
|
if err != nil {
|
|
return nil, errs.Wrap(err, errs.CodeInternal, "get conversation")
|
|
}
|
|
return convFromRow(row), nil
|
|
}
|
|
|
|
func (r *pgRepo) ListConversationsByUser(ctx context.Context, userID uuid.UUID, f ConversationFilter) ([]Conversation, error) {
|
|
limit := f.Limit
|
|
if limit <= 0 {
|
|
limit = 20
|
|
}
|
|
rows, err := r.q.ListConversationsByUser(ctx, chatsqlc.ListConversationsByUserParams{
|
|
UserID: toPgUUID(userID),
|
|
Column2: toPgUUIDPtr(f.ProjectID),
|
|
Column3: toPgUUIDPtr(f.WorkspaceID),
|
|
Column4: toPgUUIDPtr(f.IssueID),
|
|
Column5: f.TitleQuery,
|
|
Limit: int32(limit),
|
|
})
|
|
if err != nil {
|
|
return nil, errs.Wrap(err, errs.CodeInternal, "list conversations")
|
|
}
|
|
out := make([]Conversation, 0, len(rows))
|
|
for _, row := range rows {
|
|
out = append(out, *convFromRow(row))
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
func (r *pgRepo) SetConversationTitleManual(ctx context.Context, id uuid.UUID, title *string) error {
|
|
if err := r.q.SetConversationTitleManual(ctx, chatsqlc.SetConversationTitleManualParams{
|
|
ID: toPgUUID(id),
|
|
Title: title,
|
|
}); err != nil {
|
|
return errs.Wrap(err, errs.CodeInternal, "set conversation title manual")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *pgRepo) UpdateConversationTitle(ctx context.Context, id uuid.UUID, title *string) error {
|
|
if err := r.q.UpdateConversationTitle(ctx, chatsqlc.UpdateConversationTitleParams{
|
|
ID: toPgUUID(id),
|
|
Title: title,
|
|
}); err != nil {
|
|
return errs.Wrap(err, errs.CodeInternal, "update conversation title")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *pgRepo) SoftDeleteConversation(ctx context.Context, id uuid.UUID) error {
|
|
if err := r.q.SoftDeleteConversation(ctx, toPgUUID(id)); err != nil {
|
|
return errs.Wrap(err, errs.CodeInternal, "soft delete conversation")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// ===== Message =====
|
|
|
|
func (r *pgRepo) InsertMessage(ctx context.Context, conversationID uuid.UUID, role MessageRole, content string, status MessageStatus) (*Message, error) {
|
|
row, err := r.q.InsertMessage(ctx, chatsqlc.InsertMessageParams{
|
|
ConversationID: toPgUUID(conversationID),
|
|
Role: string(role),
|
|
Content: content,
|
|
Status: string(status),
|
|
})
|
|
if err != nil {
|
|
return nil, translatePgError(err)
|
|
}
|
|
return messageFromRow(row), nil
|
|
}
|
|
|
|
func (r *pgRepo) GetMessage(ctx context.Context, id int64) (*Message, error) {
|
|
row, err := r.q.GetMessage(ctx, id)
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return nil, errs.New(errs.CodeChatMessageNotFound, "message not found")
|
|
}
|
|
if err != nil {
|
|
return nil, errs.Wrap(err, errs.CodeInternal, "get message")
|
|
}
|
|
return messageFromRow(row), nil
|
|
}
|
|
|
|
func (r *pgRepo) GetFirstUserMessage(ctx context.Context, conversationID uuid.UUID) (*Message, error) {
|
|
row, err := r.q.GetFirstUserMessage(ctx, toPgUUID(conversationID))
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return nil, errs.New(errs.CodeChatMessageNotFound, "no user message found")
|
|
}
|
|
if err != nil {
|
|
return nil, errs.Wrap(err, errs.CodeInternal, "get first user message")
|
|
}
|
|
return messageFromRow(row), nil
|
|
}
|
|
|
|
func (r *pgRepo) ListMessagesByConversation(ctx context.Context, conversationID uuid.UUID, beforeID *int64, limit int) ([]Message, error) {
|
|
var before int64
|
|
if beforeID != nil {
|
|
before = *beforeID
|
|
}
|
|
if limit <= 0 {
|
|
limit = 50
|
|
}
|
|
rows, err := r.q.ListMessagesByConversation(ctx, chatsqlc.ListMessagesByConversationParams{
|
|
ConversationID: toPgUUID(conversationID),
|
|
Column2: before,
|
|
Limit: int32(limit),
|
|
})
|
|
if err != nil {
|
|
return nil, errs.Wrap(err, errs.CodeInternal, "list messages")
|
|
}
|
|
out := make([]Message, 0, len(rows))
|
|
for _, row := range rows {
|
|
out = append(out, *messageFromRow(row))
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
func (r *pgRepo) ListMessagesAscending(ctx context.Context, conversationID uuid.UUID) ([]Message, error) {
|
|
rows, err := r.q.ListMessagesAscending(ctx, toPgUUID(conversationID))
|
|
if err != nil {
|
|
return nil, errs.Wrap(err, errs.CodeInternal, "list messages ascending")
|
|
}
|
|
out := make([]Message, 0, len(rows))
|
|
for _, row := range rows {
|
|
out = append(out, *messageFromRow(row))
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
func (r *pgRepo) UpdateMessageOK(ctx context.Context, p UpdateMessageOKParams) error {
|
|
if err := r.q.UpdateMessageOK(ctx, chatsqlc.UpdateMessageOKParams{
|
|
ID: p.ID,
|
|
Content: p.Content,
|
|
Thinking: p.Thinking,
|
|
PromptTokens: p.PromptTokens,
|
|
CompletionTokens: p.CompletionTokens,
|
|
ThinkingTokens: p.ThinkingTokens,
|
|
}); err != nil {
|
|
return errs.Wrap(err, errs.CodeInternal, "update message ok")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *pgRepo) UpdateMessageError(ctx context.Context, id int64, errMsg *string) error {
|
|
if err := r.q.UpdateMessageError(ctx, chatsqlc.UpdateMessageErrorParams{
|
|
ID: id,
|
|
ErrorMessage: errMsg,
|
|
}); err != nil {
|
|
return errs.Wrap(err, errs.CodeInternal, "update message error")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *pgRepo) UpdateMessageCancelled(ctx context.Context, id int64) error {
|
|
if err := r.q.UpdateMessageCancelled(ctx, id); err != nil {
|
|
return errs.Wrap(err, errs.CodeInternal, "update message cancelled")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *pgRepo) MarkPendingAsErrorOnStartup(ctx context.Context) error {
|
|
if err := r.q.MarkPendingAsErrorOnStartup(ctx); err != nil {
|
|
return errs.Wrap(err, errs.CodeInternal, "mark pending as error on startup")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *pgRepo) DeleteMessage(ctx context.Context, id int64) error {
|
|
if err := r.q.DeleteMessage(ctx, id); err != nil {
|
|
return errs.Wrap(err, errs.CodeInternal, "delete message")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// ===== Attachment =====
|
|
|
|
func (r *pgRepo) InsertAttachment(ctx context.Context, p InsertAttachmentParams) (*Attachment, error) {
|
|
row, err := r.q.InsertAttachment(ctx, chatsqlc.InsertAttachmentParams{
|
|
ID: toPgUUID(p.ID),
|
|
UserID: toPgUUID(p.UserID),
|
|
Filename: p.Filename,
|
|
MimeType: p.MimeType,
|
|
SizeBytes: p.SizeBytes,
|
|
Sha256: p.SHA256,
|
|
StoragePath: p.StoragePath,
|
|
})
|
|
if err != nil {
|
|
return nil, errs.Wrap(err, errs.CodeInternal, "insert attachment")
|
|
}
|
|
return attachmentFromRow(row), nil
|
|
}
|
|
|
|
func (r *pgRepo) GetAttachment(ctx context.Context, id uuid.UUID) (*Attachment, error) {
|
|
row, err := r.q.GetAttachment(ctx, toPgUUID(id))
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return nil, errs.New(errs.CodeNotFound, "attachment not found")
|
|
}
|
|
if err != nil {
|
|
return nil, errs.Wrap(err, errs.CodeInternal, "get attachment")
|
|
}
|
|
return attachmentFromRow(row), nil
|
|
}
|
|
|
|
func (r *pgRepo) AttachToMessage(ctx context.Context, ids []uuid.UUID, messageID int64) error {
|
|
pgIDs := make([]pgtype.UUID, len(ids))
|
|
for i, id := range ids {
|
|
pgIDs[i] = toPgUUID(id)
|
|
}
|
|
if err := r.q.AttachToMessage(ctx, chatsqlc.AttachToMessageParams{
|
|
Column1: pgIDs,
|
|
MessageID: &messageID,
|
|
}); err != nil {
|
|
return errs.Wrap(err, errs.CodeInternal, "attach to message")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *pgRepo) ListAttachmentsForMessage(ctx context.Context, messageID int64) ([]Attachment, error) {
|
|
rows, err := r.q.ListAttachmentsForMessage(ctx, &messageID)
|
|
if err != nil {
|
|
return nil, errs.Wrap(err, errs.CodeInternal, "list attachments for message")
|
|
}
|
|
out := make([]Attachment, 0, len(rows))
|
|
for _, row := range rows {
|
|
out = append(out, *attachmentFromRow(row))
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
func (r *pgRepo) ListOrphanAttachmentsForUser(ctx context.Context, userID uuid.UUID, ids []uuid.UUID) ([]Attachment, error) {
|
|
pgIDs := make([]pgtype.UUID, len(ids))
|
|
for i, id := range ids {
|
|
pgIDs[i] = toPgUUID(id)
|
|
}
|
|
rows, err := r.q.ListOrphanAttachmentsForUser(ctx, chatsqlc.ListOrphanAttachmentsForUserParams{
|
|
UserID: toPgUUID(userID),
|
|
Column2: pgIDs,
|
|
})
|
|
if err != nil {
|
|
return nil, errs.Wrap(err, errs.CodeInternal, "list orphan attachments")
|
|
}
|
|
out := make([]Attachment, 0, len(rows))
|
|
for _, row := range rows {
|
|
out = append(out, *attachmentFromRow(row))
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
func (r *pgRepo) ListExpiredOrphans(ctx context.Context, olderThan time.Duration) ([]AttachmentRef, error) {
|
|
rows, err := r.q.ListExpiredOrphans(ctx, pgtype.Interval{
|
|
Microseconds: int64(olderThan / time.Microsecond),
|
|
Valid: true,
|
|
})
|
|
if err != nil {
|
|
return nil, errs.Wrap(err, errs.CodeInternal, "list expired orphans")
|
|
}
|
|
out := make([]AttachmentRef, 0, len(rows))
|
|
for _, row := range rows {
|
|
out = append(out, AttachmentRef{
|
|
ID: fromPgUUID(row.ID),
|
|
StoragePath: row.StoragePath,
|
|
})
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
func (r *pgRepo) ListAttachmentsForSoftDeletedConversations(ctx context.Context, olderThan time.Duration) ([]AttachmentRef, error) {
|
|
rows, err := r.q.ListAttachmentsForSoftDeletedConversations(ctx, pgtype.Interval{
|
|
Microseconds: int64(olderThan / time.Microsecond),
|
|
Valid: true,
|
|
})
|
|
if err != nil {
|
|
return nil, errs.Wrap(err, errs.CodeInternal, "list attachments for soft deleted")
|
|
}
|
|
out := make([]AttachmentRef, 0, len(rows))
|
|
for _, row := range rows {
|
|
out = append(out, AttachmentRef{
|
|
ID: fromPgUUID(row.ID),
|
|
StoragePath: row.StoragePath,
|
|
})
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
func (r *pgRepo) DeleteAttachment(ctx context.Context, id uuid.UUID) error {
|
|
if err := r.q.DeleteAttachment(ctx, toPgUUID(id)); err != nil {
|
|
return errs.Wrap(err, errs.CodeInternal, "delete attachment")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// ===== PromptTemplate =====
|
|
|
|
func (r *pgRepo) InsertPromptTemplate(ctx context.Context, p InsertTemplateParams) (*PromptTemplate, error) {
|
|
row, err := r.q.InsertPromptTemplate(ctx, chatsqlc.InsertPromptTemplateParams{
|
|
ID: toPgUUID(p.ID),
|
|
Name: p.Name,
|
|
Content: p.Content,
|
|
Scope: string(p.Scope),
|
|
OwnerID: toPgUUIDPtr(p.OwnerID),
|
|
})
|
|
if err != nil {
|
|
return nil, errs.Wrap(err, errs.CodeInternal, "insert template")
|
|
}
|
|
return templateFromRow(row), nil
|
|
}
|
|
|
|
func (r *pgRepo) GetPromptTemplate(ctx context.Context, id uuid.UUID) (*PromptTemplate, error) {
|
|
row, err := r.q.GetPromptTemplate(ctx, toPgUUID(id))
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return nil, errs.New(errs.CodeChatTemplateNotFound, "template not found")
|
|
}
|
|
if err != nil {
|
|
return nil, errs.Wrap(err, errs.CodeInternal, "get template")
|
|
}
|
|
return templateFromRow(row), nil
|
|
}
|
|
|
|
func (r *pgRepo) ListPromptTemplatesForUser(ctx context.Context, userID uuid.UUID) ([]PromptTemplate, error) {
|
|
rows, err := r.q.ListPromptTemplatesForUser(ctx, toPgUUID(userID))
|
|
if err != nil {
|
|
return nil, errs.Wrap(err, errs.CodeInternal, "list templates")
|
|
}
|
|
out := make([]PromptTemplate, 0, len(rows))
|
|
for _, row := range rows {
|
|
out = append(out, *templateFromRow(row))
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
func (r *pgRepo) UpdatePromptTemplate(ctx context.Context, id uuid.UUID, name, content string) error {
|
|
if err := r.q.UpdatePromptTemplate(ctx, chatsqlc.UpdatePromptTemplateParams{
|
|
ID: toPgUUID(id),
|
|
Name: name,
|
|
Content: content,
|
|
}); err != nil {
|
|
return errs.Wrap(err, errs.CodeInternal, "update template")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *pgRepo) DeletePromptTemplate(ctx context.Context, id uuid.UUID) error {
|
|
if err := r.q.DeletePromptTemplate(ctx, toPgUUID(id)); err != nil {
|
|
return errs.Wrap(err, errs.CodeInternal, "delete template")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// ===== LLMEndpoint =====
|
|
|
|
func (r *pgRepo) InsertLLMEndpoint(ctx context.Context, p InsertEndpointParams) (*LLMEndpoint, error) {
|
|
row, err := r.q.InsertLLMEndpoint(ctx, chatsqlc.InsertLLMEndpointParams{
|
|
ID: toPgUUID(p.ID),
|
|
Provider: string(p.Provider),
|
|
DisplayName: p.DisplayName,
|
|
BaseUrl: p.BaseURL,
|
|
ApiKeyEncrypted: p.APIKeyEncrypted,
|
|
CreatedBy: toPgUUID(p.CreatedBy),
|
|
})
|
|
if err != nil {
|
|
return nil, errs.Wrap(err, errs.CodeInternal, "insert endpoint")
|
|
}
|
|
return endpointFromRow(row), nil
|
|
}
|
|
|
|
func (r *pgRepo) GetLLMEndpoint(ctx context.Context, id uuid.UUID) (*LLMEndpoint, error) {
|
|
row, err := r.q.GetLLMEndpoint(ctx, toPgUUID(id))
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return nil, errs.New(errs.CodeChatEndpointNotFound, "endpoint not found")
|
|
}
|
|
if err != nil {
|
|
return nil, errs.Wrap(err, errs.CodeInternal, "get endpoint")
|
|
}
|
|
return endpointFromRow(row), nil
|
|
}
|
|
|
|
func (r *pgRepo) ListLLMEndpoints(ctx context.Context) ([]LLMEndpoint, error) {
|
|
rows, err := r.q.ListLLMEndpoints(ctx)
|
|
if err != nil {
|
|
return nil, errs.Wrap(err, errs.CodeInternal, "list endpoints")
|
|
}
|
|
out := make([]LLMEndpoint, 0, len(rows))
|
|
for _, row := range rows {
|
|
out = append(out, *endpointFromRow(row))
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
// UpdateLLMEndpoint は UpdateEndpointInput に基づいて既存エンドポイントを更新する。
|
|
// nil フィールドは変更なし(先に Get して現在値でマージする)。
|
|
func (r *pgRepo) UpdateLLMEndpoint(ctx context.Context, id uuid.UUID, in UpdateEndpointInput) error {
|
|
existing, err := r.GetLLMEndpoint(ctx, id)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
displayName := existing.DisplayName
|
|
if in.DisplayName != nil {
|
|
displayName = *in.DisplayName
|
|
}
|
|
baseURL := existing.BaseURL
|
|
if in.BaseURL != nil {
|
|
baseURL = *in.BaseURL
|
|
}
|
|
apiKey := existing.APIKeyEncrypted
|
|
if in.APIKey != nil {
|
|
// in.APIKey は明文;service 層で暗号化後にここに渡す想定だが、
|
|
// repository 層では []byte として受け取る。UpdateEndpointInput.APIKey は
|
|
// *string(明文)なので、service 層で変換済みのものが渡る。
|
|
// ここでは文字列をバイト列に変換するだけ。
|
|
apiKey = []byte(*in.APIKey)
|
|
}
|
|
if err := r.q.UpdateLLMEndpoint(ctx, chatsqlc.UpdateLLMEndpointParams{
|
|
ID: toPgUUID(id),
|
|
DisplayName: displayName,
|
|
BaseUrl: baseURL,
|
|
ApiKeyEncrypted: apiKey,
|
|
}); err != nil {
|
|
return errs.Wrap(err, errs.CodeInternal, "update endpoint")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *pgRepo) DeleteLLMEndpoint(ctx context.Context, id uuid.UUID) error {
|
|
if err := r.q.DeleteLLMEndpoint(ctx, toPgUUID(id)); err != nil {
|
|
return errs.Wrap(err, errs.CodeInternal, "delete endpoint")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// ===== LLMModel =====
|
|
|
|
func (r *pgRepo) InsertLLMModel(ctx context.Context, p InsertModelParams) (*LLMModel, error) {
|
|
capsBytes, err := json.Marshal(p.Capabilities)
|
|
if err != nil {
|
|
return nil, errs.Wrap(err, errs.CodeInternal, "marshal capabilities")
|
|
}
|
|
row, err := r.q.InsertLLMModel(ctx, chatsqlc.InsertLLMModelParams{
|
|
ID: toPgUUID(p.ID),
|
|
EndpointID: toPgUUID(p.EndpointID),
|
|
ModelID: p.ModelID,
|
|
DisplayName: p.DisplayName,
|
|
Capabilities: capsBytes,
|
|
ContextWindow: int32(p.ContextWindow),
|
|
MaxOutputTokens: int32(p.MaxOutputTokens),
|
|
PromptPricePerMillionUsd: float64ToNumeric(p.PromptPrice),
|
|
CompletionPricePerMillionUsd: float64ToNumeric(p.CompletionPrice),
|
|
ThinkingPricePerMillionUsd: float64ToNumeric(p.ThinkingPrice),
|
|
IsTitleGenerator: p.IsTitleGenerator,
|
|
Enabled: p.Enabled,
|
|
SortOrder: int32(p.SortOrder),
|
|
})
|
|
if err != nil {
|
|
return nil, translatePgError(err)
|
|
}
|
|
return modelFromRow(row), nil
|
|
}
|
|
|
|
func (r *pgRepo) GetLLMModel(ctx context.Context, id uuid.UUID) (*LLMModel, error) {
|
|
row, err := r.q.GetLLMModel(ctx, toPgUUID(id))
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return nil, errs.New(errs.CodeChatModelNotFound, "model not found")
|
|
}
|
|
if err != nil {
|
|
return nil, errs.Wrap(err, errs.CodeInternal, "get model")
|
|
}
|
|
return modelFromRow(row), nil
|
|
}
|
|
|
|
func (r *pgRepo) GetTitleGeneratorModel(ctx context.Context) (*LLMModel, error) {
|
|
row, err := r.q.GetTitleGeneratorModel(ctx)
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return nil, errs.New(errs.CodeChatModelNotFound, "no title generator model found")
|
|
}
|
|
if err != nil {
|
|
return nil, errs.Wrap(err, errs.CodeInternal, "get title generator model")
|
|
}
|
|
return modelFromRow(row), nil
|
|
}
|
|
|
|
func (r *pgRepo) ListLLMModels(ctx context.Context, onlyEnabled bool) ([]LLMModel, error) {
|
|
rows, err := r.q.ListLLMModels(ctx, onlyEnabled)
|
|
if err != nil {
|
|
return nil, errs.Wrap(err, errs.CodeInternal, "list models")
|
|
}
|
|
out := make([]LLMModel, 0, len(rows))
|
|
for _, row := range rows {
|
|
out = append(out, *modelFromRow(row))
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
// UpdateLLMModel は UpdateModelInput に基づいて既存モデルを更新する。
|
|
// nil フィールドは変更なし(先に Get して現在値でマージする)。
|
|
func (r *pgRepo) UpdateLLMModel(ctx context.Context, id uuid.UUID, in UpdateModelInput) error {
|
|
existing, err := r.GetLLMModel(ctx, id)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
displayName := existing.DisplayName
|
|
if in.DisplayName != nil {
|
|
displayName = *in.DisplayName
|
|
}
|
|
caps := existing.Capabilities
|
|
if in.Capabilities != nil {
|
|
caps = *in.Capabilities
|
|
}
|
|
contextWindow := existing.ContextWindow
|
|
if in.ContextWindow != nil {
|
|
contextWindow = *in.ContextWindow
|
|
}
|
|
maxOutputTokens := existing.MaxOutputTokens
|
|
if in.MaxOutputTokens != nil {
|
|
maxOutputTokens = *in.MaxOutputTokens
|
|
}
|
|
promptPrice := existing.PromptPricePerMillionUSD
|
|
if in.PromptPrice != nil {
|
|
promptPrice = *in.PromptPrice
|
|
}
|
|
completionPrice := existing.CompletionPricePerMillionUSD
|
|
if in.CompletionPrice != nil {
|
|
completionPrice = *in.CompletionPrice
|
|
}
|
|
thinkingPrice := existing.ThinkingPricePerMillionUSD
|
|
if in.ThinkingPrice != nil {
|
|
thinkingPrice = *in.ThinkingPrice
|
|
}
|
|
isTitleGenerator := existing.IsTitleGenerator
|
|
if in.IsTitleGenerator != nil {
|
|
isTitleGenerator = *in.IsTitleGenerator
|
|
}
|
|
enabled := existing.Enabled
|
|
if in.Enabled != nil {
|
|
enabled = *in.Enabled
|
|
}
|
|
sortOrder := existing.SortOrder
|
|
if in.SortOrder != nil {
|
|
sortOrder = *in.SortOrder
|
|
}
|
|
|
|
capsBytes, err := json.Marshal(caps)
|
|
if err != nil {
|
|
return errs.Wrap(err, errs.CodeInternal, "marshal capabilities")
|
|
}
|
|
|
|
if err := r.q.UpdateLLMModel(ctx, chatsqlc.UpdateLLMModelParams{
|
|
ID: toPgUUID(id),
|
|
DisplayName: displayName,
|
|
Capabilities: capsBytes,
|
|
ContextWindow: int32(contextWindow),
|
|
MaxOutputTokens: int32(maxOutputTokens),
|
|
PromptPricePerMillionUsd: float64ToNumeric(promptPrice),
|
|
CompletionPricePerMillionUsd: float64ToNumeric(completionPrice),
|
|
ThinkingPricePerMillionUsd: float64ToNumeric(thinkingPrice),
|
|
IsTitleGenerator: isTitleGenerator,
|
|
Enabled: enabled,
|
|
SortOrder: int32(sortOrder),
|
|
}); err != nil {
|
|
return translatePgError(err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *pgRepo) DeleteLLMModel(ctx context.Context, id uuid.UUID) error {
|
|
if err := r.q.DeleteLLMModel(ctx, toPgUUID(id)); err != nil {
|
|
return errs.Wrap(err, errs.CodeInternal, "delete model")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// ===== Usage =====
|
|
|
|
func (r *pgRepo) InsertUsage(ctx context.Context, rec UsageRecord) error {
|
|
if err := r.q.InsertUsage(ctx, chatsqlc.InsertUsageParams{
|
|
ConversationID: toPgUUID(rec.ConversationID),
|
|
MessageID: rec.MessageID,
|
|
UserID: toPgUUID(rec.UserID),
|
|
EndpointID: toPgUUID(rec.EndpointID),
|
|
ModelID: toPgUUID(rec.ModelID),
|
|
PromptTokens: int32(rec.PromptTokens),
|
|
CompletionTokens: int32(rec.CompletionTokens),
|
|
ThinkingTokens: int32(rec.ThinkingTokens),
|
|
CostUsd: float64ToNumeric(rec.CostUSD),
|
|
}); err != nil {
|
|
return errs.Wrap(err, errs.CodeInternal, "insert usage")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *pgRepo) SummarizeUsageByUser(ctx context.Context, from, to time.Time) ([]UserUsageRow, error) {
|
|
rows, err := r.q.SummarizeUsageByUser(ctx, chatsqlc.SummarizeUsageByUserParams{
|
|
CreatedAt: toPgTimestamptz(from),
|
|
CreatedAt_2: toPgTimestamptz(to),
|
|
})
|
|
if err != nil {
|
|
return nil, errs.Wrap(err, errs.CodeInternal, "summarize usage by user")
|
|
}
|
|
out := make([]UserUsageRow, 0, len(rows))
|
|
for _, row := range rows {
|
|
out = append(out, userUsageRowFromRow(row))
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
func (r *pgRepo) SummarizeUsageByModel(ctx context.Context, from, to time.Time) ([]ModelUsageRow, error) {
|
|
rows, err := r.q.SummarizeUsageByModel(ctx, chatsqlc.SummarizeUsageByModelParams{
|
|
CreatedAt: toPgTimestamptz(from),
|
|
CreatedAt_2: toPgTimestamptz(to),
|
|
})
|
|
if err != nil {
|
|
return nil, errs.Wrap(err, errs.CodeInternal, "summarize usage by model")
|
|
}
|
|
out := make([]ModelUsageRow, 0, len(rows))
|
|
for _, row := range rows {
|
|
out = append(out, modelUsageRowFromRow(row))
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
func (r *pgRepo) SummarizeUsageByDay(ctx context.Context, from, to time.Time) ([]DayUsageRow, error) {
|
|
rows, err := r.q.SummarizeUsageByDay(ctx, chatsqlc.SummarizeUsageByDayParams{
|
|
CreatedAt: toPgTimestamptz(from),
|
|
CreatedAt_2: toPgTimestamptz(to),
|
|
})
|
|
if err != nil {
|
|
return nil, errs.Wrap(err, errs.CodeInternal, "summarize usage by day")
|
|
}
|
|
out := make([]DayUsageRow, 0, len(rows))
|
|
for _, row := range rows {
|
|
out = append(out, dayUsageRowFromRow(row))
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
func (r *pgRepo) SummarizeUserDaily(ctx context.Context, userID uuid.UUID, from, to time.Time) ([]DayUsageRow, error) {
|
|
rows, err := r.q.SummarizeUserDaily(ctx, chatsqlc.SummarizeUserDailyParams{
|
|
UserID: toPgUUID(userID),
|
|
CreatedAt: toPgTimestamptz(from),
|
|
CreatedAt_2: toPgTimestamptz(to),
|
|
})
|
|
if err != nil {
|
|
return nil, errs.Wrap(err, errs.CodeInternal, "summarize user daily")
|
|
}
|
|
out := make([]DayUsageRow, 0, len(rows))
|
|
for _, row := range rows {
|
|
out = append(out, DayUsageRow{
|
|
Day: intervalToTime(row.Day),
|
|
PromptTokens: int(row.PromptTokens),
|
|
CompletionTokens: int(row.CompletionTokens),
|
|
ThinkingTokens: int(row.ThinkingTokens),
|
|
CostUSD: numericToFloat64(row.CostUsd),
|
|
})
|
|
}
|
|
return out, nil
|
|
}
|