mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: route chatd provider traffic through aibridge (#25629)
## Summary Routes chatd model calls backed by concrete AI Provider rows through the in-process aibridge transport by default, with deployment options to use direct provider routing when AI Gateway is disabled or chat AI Gateway routing is disabled. - Splits model routing into common, direct provider, and AI Gateway paths behind a single deployment-mode entry point. - Builds chatd models through explicit request, route, and options data. Active API key attribution is passed explicitly instead of being hidden inside generic model construction. - For AI Gateway BYOK routes, resolves the user's provider key in chatd, forwards it through provider-specific auth headers, and sets `X-Coder-AI-Governance-Token` to the `delegated` marker so aibridge preserves those headers while still stripping Coder-specific metadata. - Keeps central provider credentials and deployment fallback credentials out of forwarded provider auth headers, so AI Gateway central policy remains authoritative. - Redacts delegated provider auth from default string formatting to avoid accidental plaintext logging of user BYOK credentials. - Covers selected chat models, advisor overrides, title and quickgen paths, subagent overrides, computer use model selection, and an integration-style chat turn through the aibridge transport path. - Persists initiating API key IDs on chat and queued user messages, including subagent child messages, and fails closed for AI Gateway-routed model builds without an active key. - Removes unused `api_key_id` indexes while keeping the persistence columns and foreign keys. - Keeps the deployment option available through config and env parsing, but hides it from CLI help and generated docs. - Stabilizes the subagent poll fallback test so background CreateChat processing cannot win the state transition under slower CI environments. ## Tests - `go test ./coderd/x/chatd -run 'TestAIGatewayProviderAuthForUser|TestAIGatewayProviderAuthRedactsFormatting|TestResolveModelRouteForConfigAIGatewayProviderAuth|TestAIGatewayModelForwardsProviderAuth|TestProcessChat_AIGatewayRoutingUsesDelegatedAPIKey|TestAwaitSubagentCompletion' -count=1` - `go test ./coderd/aibridged -run 'TestServeHTTP_DelegatedAPIKey|TestServeHTTP_StripCoderToken' -count=1` - `git diff --check HEAD~1..HEAD` - `make lint` > Mux working on behalf of Mike.
This commit is contained in:
Michael Suchacz
@@ -119,6 +119,7 @@ func ChatMessage(t testing.TB, db database.Store, seed database.ChatMessage) dat
|
||||
msgs, err := db.InsertChatMessages(genCtx, database.InsertChatMessagesParams{
|
||||
ChatID: seed.ChatID,
|
||||
CreatedBy: []uuid.UUID{seed.CreatedBy.UUID},
|
||||
APIKeyID: []string{seed.APIKeyID.String},
|
||||
ModelConfigID: []uuid.UUID{seed.ModelConfigID.UUID},
|
||||
Role: []database.ChatMessageRole{takeFirst(seed.Role, database.ChatMessageRoleUser)},
|
||||
Content: []string{content},
|
||||
|
||||
Generated
+10
-2
@@ -1554,7 +1554,8 @@ CREATE TABLE chat_messages (
|
||||
total_cost_micros bigint,
|
||||
runtime_ms bigint,
|
||||
deleted boolean DEFAULT false NOT NULL,
|
||||
provider_response_id text
|
||||
provider_response_id text,
|
||||
api_key_id text
|
||||
);
|
||||
|
||||
CREATE SEQUENCE chat_messages_id_seq
|
||||
@@ -1593,7 +1594,8 @@ CREATE TABLE chat_queued_messages (
|
||||
chat_id uuid NOT NULL,
|
||||
content jsonb NOT NULL,
|
||||
created_at timestamp with time zone DEFAULT now() NOT NULL,
|
||||
model_config_id uuid
|
||||
model_config_id uuid,
|
||||
api_key_id text
|
||||
);
|
||||
|
||||
CREATE SEQUENCE chat_queued_messages_id_seq
|
||||
@@ -4453,6 +4455,9 @@ ALTER TABLE ONLY chat_files
|
||||
ALTER TABLE ONLY chat_files
|
||||
ADD CONSTRAINT chat_files_owner_id_fkey FOREIGN KEY (owner_id) REFERENCES users(id) ON DELETE CASCADE;
|
||||
|
||||
ALTER TABLE ONLY chat_messages
|
||||
ADD CONSTRAINT chat_messages_api_key_id_fkey FOREIGN KEY (api_key_id) REFERENCES api_keys(id) ON DELETE SET NULL;
|
||||
|
||||
ALTER TABLE ONLY chat_messages
|
||||
ADD CONSTRAINT chat_messages_chat_id_fkey FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE;
|
||||
|
||||
@@ -4468,6 +4473,9 @@ ALTER TABLE ONLY chat_model_configs
|
||||
ALTER TABLE ONLY chat_model_configs
|
||||
ADD CONSTRAINT chat_model_configs_updated_by_fkey FOREIGN KEY (updated_by) REFERENCES users(id);
|
||||
|
||||
ALTER TABLE ONLY chat_queued_messages
|
||||
ADD CONSTRAINT chat_queued_messages_api_key_id_fkey FOREIGN KEY (api_key_id) REFERENCES api_keys(id) ON DELETE SET NULL;
|
||||
|
||||
ALTER TABLE ONLY chat_queued_messages
|
||||
ADD CONSTRAINT chat_queued_messages_chat_id_fkey FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE;
|
||||
|
||||
|
||||
@@ -21,11 +21,13 @@ const (
|
||||
ForeignKeyChatFileLinksFileID ForeignKeyConstraint = "chat_file_links_file_id_fkey" // ALTER TABLE ONLY chat_file_links ADD CONSTRAINT chat_file_links_file_id_fkey FOREIGN KEY (file_id) REFERENCES chat_files(id) ON DELETE CASCADE;
|
||||
ForeignKeyChatFilesOrganizationID ForeignKeyConstraint = "chat_files_organization_id_fkey" // ALTER TABLE ONLY chat_files ADD CONSTRAINT chat_files_organization_id_fkey FOREIGN KEY (organization_id) REFERENCES organizations(id) ON DELETE CASCADE;
|
||||
ForeignKeyChatFilesOwnerID ForeignKeyConstraint = "chat_files_owner_id_fkey" // ALTER TABLE ONLY chat_files ADD CONSTRAINT chat_files_owner_id_fkey FOREIGN KEY (owner_id) REFERENCES users(id) ON DELETE CASCADE;
|
||||
ForeignKeyChatMessagesAPIKeyID ForeignKeyConstraint = "chat_messages_api_key_id_fkey" // ALTER TABLE ONLY chat_messages ADD CONSTRAINT chat_messages_api_key_id_fkey FOREIGN KEY (api_key_id) REFERENCES api_keys(id) ON DELETE SET NULL;
|
||||
ForeignKeyChatMessagesChatID ForeignKeyConstraint = "chat_messages_chat_id_fkey" // ALTER TABLE ONLY chat_messages ADD CONSTRAINT chat_messages_chat_id_fkey FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE;
|
||||
ForeignKeyChatMessagesModelConfigID ForeignKeyConstraint = "chat_messages_model_config_id_fkey" // ALTER TABLE ONLY chat_messages ADD CONSTRAINT chat_messages_model_config_id_fkey FOREIGN KEY (model_config_id) REFERENCES chat_model_configs(id);
|
||||
ForeignKeyChatModelConfigsAiProviderID ForeignKeyConstraint = "chat_model_configs_ai_provider_id_fkey" // ALTER TABLE ONLY chat_model_configs ADD CONSTRAINT chat_model_configs_ai_provider_id_fkey FOREIGN KEY (ai_provider_id) REFERENCES ai_providers(id);
|
||||
ForeignKeyChatModelConfigsCreatedBy ForeignKeyConstraint = "chat_model_configs_created_by_fkey" // ALTER TABLE ONLY chat_model_configs ADD CONSTRAINT chat_model_configs_created_by_fkey FOREIGN KEY (created_by) REFERENCES users(id);
|
||||
ForeignKeyChatModelConfigsUpdatedBy ForeignKeyConstraint = "chat_model_configs_updated_by_fkey" // ALTER TABLE ONLY chat_model_configs ADD CONSTRAINT chat_model_configs_updated_by_fkey FOREIGN KEY (updated_by) REFERENCES users(id);
|
||||
ForeignKeyChatQueuedMessagesAPIKeyID ForeignKeyConstraint = "chat_queued_messages_api_key_id_fkey" // ALTER TABLE ONLY chat_queued_messages ADD CONSTRAINT chat_queued_messages_api_key_id_fkey FOREIGN KEY (api_key_id) REFERENCES api_keys(id) ON DELETE SET NULL;
|
||||
ForeignKeyChatQueuedMessagesChatID ForeignKeyConstraint = "chat_queued_messages_chat_id_fkey" // ALTER TABLE ONLY chat_queued_messages ADD CONSTRAINT chat_queued_messages_chat_id_fkey FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE;
|
||||
ForeignKeyChatsAgentID ForeignKeyConstraint = "chats_agent_id_fkey" // ALTER TABLE ONLY chats ADD CONSTRAINT chats_agent_id_fkey FOREIGN KEY (agent_id) REFERENCES workspace_agents(id) ON DELETE SET NULL;
|
||||
ForeignKeyChatsBuildID ForeignKeyConstraint = "chats_build_id_fkey" // ALTER TABLE ONLY chats ADD CONSTRAINT chats_build_id_fkey FOREIGN KEY (build_id) REFERENCES workspace_builds(id) ON DELETE SET NULL;
|
||||
|
||||
@@ -0,0 +1,5 @@
|
||||
ALTER TABLE chat_queued_messages
|
||||
DROP COLUMN api_key_id;
|
||||
|
||||
ALTER TABLE chat_messages
|
||||
DROP COLUMN api_key_id;
|
||||
@@ -0,0 +1,8 @@
|
||||
-- Preserve chat history when API keys are deleted. Pending work whose latest
|
||||
-- user turn loses this attribution will fail closed under AI Gateway routing;
|
||||
-- operators can retry the turn or temporarily use direct routing.
|
||||
ALTER TABLE chat_messages
|
||||
ADD COLUMN api_key_id text REFERENCES api_keys(id) ON DELETE SET NULL;
|
||||
|
||||
ALTER TABLE chat_queued_messages
|
||||
ADD COLUMN api_key_id text REFERENCES api_keys(id) ON DELETE SET NULL;
|
||||
@@ -4699,6 +4699,7 @@ type ChatMessage struct {
|
||||
RuntimeMs sql.NullInt64 `db:"runtime_ms" json:"runtime_ms"`
|
||||
Deleted bool `db:"deleted" json:"deleted"`
|
||||
ProviderResponseID sql.NullString `db:"provider_response_id" json:"provider_response_id"`
|
||||
APIKeyID sql.NullString `db:"api_key_id" json:"api_key_id"`
|
||||
}
|
||||
|
||||
type ChatModelConfig struct {
|
||||
@@ -4726,6 +4727,7 @@ type ChatQueuedMessage struct {
|
||||
Content json.RawMessage `db:"content" json:"content"`
|
||||
CreatedAt time.Time `db:"created_at" json:"created_at"`
|
||||
ModelConfigID uuid.NullUUID `db:"model_config_id" json:"model_config_id"`
|
||||
APIKeyID sql.NullString `db:"api_key_id" json:"api_key_id"`
|
||||
}
|
||||
|
||||
type ChatTable struct {
|
||||
|
||||
@@ -7157,7 +7157,7 @@ func (q *sqlQuerier) GetChatDiffStatusesByChatIDs(ctx context.Context, chatIds [
|
||||
|
||||
const getChatMessageByID = `-- name: GetChatMessageByID :one
|
||||
SELECT
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id
|
||||
FROM
|
||||
chat_messages
|
||||
WHERE
|
||||
@@ -7190,6 +7190,7 @@ func (q *sqlQuerier) GetChatMessageByID(ctx context.Context, id int64) (ChatMess
|
||||
&i.RuntimeMs,
|
||||
&i.Deleted,
|
||||
&i.ProviderResponseID,
|
||||
&i.APIKeyID,
|
||||
)
|
||||
return i, err
|
||||
}
|
||||
@@ -7279,7 +7280,7 @@ func (q *sqlQuerier) GetChatMessageSummariesPerChat(ctx context.Context, created
|
||||
|
||||
const getChatMessagesByChatID = `-- name: GetChatMessagesByChatID :many
|
||||
SELECT
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id
|
||||
FROM
|
||||
chat_messages
|
||||
WHERE
|
||||
@@ -7327,6 +7328,7 @@ func (q *sqlQuerier) GetChatMessagesByChatID(ctx context.Context, arg GetChatMes
|
||||
&i.RuntimeMs,
|
||||
&i.Deleted,
|
||||
&i.ProviderResponseID,
|
||||
&i.APIKeyID,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -7343,7 +7345,7 @@ func (q *sqlQuerier) GetChatMessagesByChatID(ctx context.Context, arg GetChatMes
|
||||
|
||||
const getChatMessagesByChatIDAscPaginated = `-- name: GetChatMessagesByChatIDAscPaginated :many
|
||||
SELECT
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id
|
||||
FROM
|
||||
chat_messages
|
||||
WHERE
|
||||
@@ -7394,6 +7396,7 @@ func (q *sqlQuerier) GetChatMessagesByChatIDAscPaginated(ctx context.Context, ar
|
||||
&i.RuntimeMs,
|
||||
&i.Deleted,
|
||||
&i.ProviderResponseID,
|
||||
&i.APIKeyID,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -7410,7 +7413,7 @@ func (q *sqlQuerier) GetChatMessagesByChatIDAscPaginated(ctx context.Context, ar
|
||||
|
||||
const getChatMessagesByChatIDDescPaginated = `-- name: GetChatMessagesByChatIDDescPaginated :many
|
||||
SELECT
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id
|
||||
FROM
|
||||
chat_messages
|
||||
WHERE
|
||||
@@ -7474,6 +7477,7 @@ func (q *sqlQuerier) GetChatMessagesByChatIDDescPaginated(ctx context.Context, a
|
||||
&i.RuntimeMs,
|
||||
&i.Deleted,
|
||||
&i.ProviderResponseID,
|
||||
&i.APIKeyID,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -7506,7 +7510,7 @@ WITH latest_compressed_summary AS (
|
||||
1
|
||||
)
|
||||
SELECT
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id
|
||||
FROM
|
||||
chat_messages
|
||||
WHERE
|
||||
@@ -7578,6 +7582,7 @@ func (q *sqlQuerier) GetChatMessagesForPromptByChatID(ctx context.Context, chatI
|
||||
&i.RuntimeMs,
|
||||
&i.Deleted,
|
||||
&i.ProviderResponseID,
|
||||
&i.APIKeyID,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -7639,7 +7644,7 @@ func (q *sqlQuerier) GetChatModelConfigsForTelemetry(ctx context.Context) ([]Get
|
||||
}
|
||||
|
||||
const getChatQueuedMessages = `-- name: GetChatQueuedMessages :many
|
||||
SELECT id, chat_id, content, created_at, model_config_id FROM chat_queued_messages
|
||||
SELECT id, chat_id, content, created_at, model_config_id, api_key_id FROM chat_queued_messages
|
||||
WHERE chat_id = $1
|
||||
ORDER BY created_at ASC, id ASC
|
||||
`
|
||||
@@ -7659,6 +7664,7 @@ func (q *sqlQuerier) GetChatQueuedMessages(ctx context.Context, chatID uuid.UUID
|
||||
&i.Content,
|
||||
&i.CreatedAt,
|
||||
&i.ModelConfigID,
|
||||
&i.APIKeyID,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -8354,7 +8360,7 @@ func (q *sqlQuerier) GetChildChatsByParentIDs(ctx context.Context, arg GetChildC
|
||||
|
||||
const getLastChatMessageByRole = `-- name: GetLastChatMessageByRole :one
|
||||
SELECT
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id
|
||||
FROM
|
||||
chat_messages
|
||||
WHERE
|
||||
@@ -8397,6 +8403,7 @@ func (q *sqlQuerier) GetLastChatMessageByRole(ctx context.Context, arg GetLastCh
|
||||
&i.RuntimeMs,
|
||||
&i.Deleted,
|
||||
&i.ProviderResponseID,
|
||||
&i.APIKeyID,
|
||||
)
|
||||
return i, err
|
||||
}
|
||||
@@ -8709,7 +8716,7 @@ WITH updated_chat AS (
|
||||
SET
|
||||
last_model_config_id = (
|
||||
SELECT val
|
||||
FROM UNNEST($3::uuid[])
|
||||
FROM UNNEST($4::uuid[])
|
||||
WITH ORDINALITY AS t(val, ord)
|
||||
WHERE val != '00000000-0000-0000-0000-000000000000'::uuid
|
||||
ORDER BY ord DESC
|
||||
@@ -8719,12 +8726,12 @@ WITH updated_chat AS (
|
||||
id = $1::uuid
|
||||
AND EXISTS (
|
||||
SELECT 1
|
||||
FROM UNNEST($3::uuid[])
|
||||
FROM UNNEST($4::uuid[])
|
||||
WHERE unnest != '00000000-0000-0000-0000-000000000000'::uuid
|
||||
)
|
||||
AND chats.last_model_config_id IS DISTINCT FROM (
|
||||
SELECT val
|
||||
FROM UNNEST($3::uuid[])
|
||||
FROM UNNEST($4::uuid[])
|
||||
WITH ORDINALITY AS t(val, ord)
|
||||
WHERE val != '00000000-0000-0000-0000-000000000000'::uuid
|
||||
ORDER BY ord DESC
|
||||
@@ -8734,6 +8741,7 @@ WITH updated_chat AS (
|
||||
INSERT INTO chat_messages (
|
||||
chat_id,
|
||||
created_by,
|
||||
api_key_id,
|
||||
model_config_id,
|
||||
role,
|
||||
content,
|
||||
@@ -8754,29 +8762,31 @@ INSERT INTO chat_messages (
|
||||
SELECT
|
||||
$1::uuid,
|
||||
NULLIF(UNNEST($2::uuid[]), '00000000-0000-0000-0000-000000000000'::uuid),
|
||||
NULLIF(UNNEST($3::uuid[]), '00000000-0000-0000-0000-000000000000'::uuid),
|
||||
UNNEST($4::chat_message_role[]),
|
||||
UNNEST($5::text[])::jsonb,
|
||||
UNNEST($6::smallint[]),
|
||||
UNNEST($7::chat_message_visibility[]),
|
||||
NULLIF(UNNEST($8::bigint[]), 0),
|
||||
NULLIF(UNNEST($3::text[]), ''),
|
||||
NULLIF(UNNEST($4::uuid[]), '00000000-0000-0000-0000-000000000000'::uuid),
|
||||
UNNEST($5::chat_message_role[]),
|
||||
UNNEST($6::text[])::jsonb,
|
||||
UNNEST($7::smallint[]),
|
||||
UNNEST($8::chat_message_visibility[]),
|
||||
NULLIF(UNNEST($9::bigint[]), 0),
|
||||
NULLIF(UNNEST($10::bigint[]), 0),
|
||||
NULLIF(UNNEST($11::bigint[]), 0),
|
||||
NULLIF(UNNEST($12::bigint[]), 0),
|
||||
NULLIF(UNNEST($13::bigint[]), 0),
|
||||
NULLIF(UNNEST($14::bigint[]), 0),
|
||||
UNNEST($15::boolean[]),
|
||||
NULLIF(UNNEST($16::bigint[]), 0),
|
||||
NULLIF(UNNEST($15::bigint[]), 0),
|
||||
UNNEST($16::boolean[]),
|
||||
NULLIF(UNNEST($17::bigint[]), 0),
|
||||
NULLIF(UNNEST($18::text[]), '')
|
||||
NULLIF(UNNEST($18::bigint[]), 0),
|
||||
NULLIF(UNNEST($19::text[]), '')
|
||||
RETURNING
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id
|
||||
`
|
||||
|
||||
type InsertChatMessagesParams struct {
|
||||
ChatID uuid.UUID `db:"chat_id" json:"chat_id"`
|
||||
CreatedBy []uuid.UUID `db:"created_by" json:"created_by"`
|
||||
APIKeyID []string `db:"api_key_id" json:"api_key_id"`
|
||||
ModelConfigID []uuid.UUID `db:"model_config_id" json:"model_config_id"`
|
||||
Role []ChatMessageRole `db:"role" json:"role"`
|
||||
Content []string `db:"content" json:"content"`
|
||||
@@ -8799,6 +8809,7 @@ func (q *sqlQuerier) InsertChatMessages(ctx context.Context, arg InsertChatMessa
|
||||
rows, err := q.db.QueryContext(ctx, insertChatMessages,
|
||||
arg.ChatID,
|
||||
pq.Array(arg.CreatedBy),
|
||||
pq.Array(arg.APIKeyID),
|
||||
pq.Array(arg.ModelConfigID),
|
||||
pq.Array(arg.Role),
|
||||
pq.Array(arg.Content),
|
||||
@@ -8845,6 +8856,7 @@ func (q *sqlQuerier) InsertChatMessages(ctx context.Context, arg InsertChatMessa
|
||||
&i.RuntimeMs,
|
||||
&i.Deleted,
|
||||
&i.ProviderResponseID,
|
||||
&i.APIKeyID,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -8860,23 +8872,30 @@ func (q *sqlQuerier) InsertChatMessages(ctx context.Context, arg InsertChatMessa
|
||||
}
|
||||
|
||||
const insertChatQueuedMessage = `-- name: InsertChatQueuedMessage :one
|
||||
INSERT INTO chat_queued_messages (chat_id, content, model_config_id)
|
||||
INSERT INTO chat_queued_messages (chat_id, content, model_config_id, api_key_id)
|
||||
VALUES (
|
||||
$1,
|
||||
$2,
|
||||
$3::uuid
|
||||
$3::uuid,
|
||||
$4::text
|
||||
)
|
||||
RETURNING id, chat_id, content, created_at, model_config_id
|
||||
RETURNING id, chat_id, content, created_at, model_config_id, api_key_id
|
||||
`
|
||||
|
||||
type InsertChatQueuedMessageParams struct {
|
||||
ChatID uuid.UUID `db:"chat_id" json:"chat_id"`
|
||||
Content json.RawMessage `db:"content" json:"content"`
|
||||
ModelConfigID uuid.NullUUID `db:"model_config_id" json:"model_config_id"`
|
||||
APIKeyID sql.NullString `db:"api_key_id" json:"api_key_id"`
|
||||
}
|
||||
|
||||
func (q *sqlQuerier) InsertChatQueuedMessage(ctx context.Context, arg InsertChatQueuedMessageParams) (ChatQueuedMessage, error) {
|
||||
row := q.db.QueryRowContext(ctx, insertChatQueuedMessage, arg.ChatID, arg.Content, arg.ModelConfigID)
|
||||
row := q.db.QueryRowContext(ctx, insertChatQueuedMessage,
|
||||
arg.ChatID,
|
||||
arg.Content,
|
||||
arg.ModelConfigID,
|
||||
arg.APIKeyID,
|
||||
)
|
||||
var i ChatQueuedMessage
|
||||
err := row.Scan(
|
||||
&i.ID,
|
||||
@@ -8884,6 +8903,7 @@ func (q *sqlQuerier) InsertChatQueuedMessage(ctx context.Context, arg InsertChat
|
||||
&i.Content,
|
||||
&i.CreatedAt,
|
||||
&i.ModelConfigID,
|
||||
&i.APIKeyID,
|
||||
)
|
||||
return i, err
|
||||
}
|
||||
@@ -9108,7 +9128,7 @@ WHERE id = (
|
||||
ORDER BY cqm.created_at ASC, cqm.id ASC
|
||||
LIMIT 1
|
||||
)
|
||||
RETURNING id, chat_id, content, created_at, model_config_id
|
||||
RETURNING id, chat_id, content, created_at, model_config_id, api_key_id
|
||||
`
|
||||
|
||||
func (q *sqlQuerier) PopNextQueuedMessage(ctx context.Context, chatID uuid.UUID) (ChatQueuedMessage, error) {
|
||||
@@ -9120,6 +9140,7 @@ func (q *sqlQuerier) PopNextQueuedMessage(ctx context.Context, chatID uuid.UUID)
|
||||
&i.Content,
|
||||
&i.CreatedAt,
|
||||
&i.ModelConfigID,
|
||||
&i.APIKeyID,
|
||||
)
|
||||
return i, err
|
||||
}
|
||||
@@ -10145,7 +10166,7 @@ SET
|
||||
WHERE
|
||||
id = $3::bigint
|
||||
RETURNING
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id
|
||||
`
|
||||
|
||||
type UpdateChatMessageByIDParams struct {
|
||||
@@ -10179,6 +10200,7 @@ func (q *sqlQuerier) UpdateChatMessageByID(ctx context.Context, arg UpdateChatMe
|
||||
&i.RuntimeMs,
|
||||
&i.Deleted,
|
||||
&i.ProviderResponseID,
|
||||
&i.APIKeyID,
|
||||
)
|
||||
return i, err
|
||||
}
|
||||
|
||||
@@ -763,6 +763,7 @@ WITH updated_chat AS (
|
||||
INSERT INTO chat_messages (
|
||||
chat_id,
|
||||
created_by,
|
||||
api_key_id,
|
||||
model_config_id,
|
||||
role,
|
||||
content,
|
||||
@@ -783,6 +784,7 @@ INSERT INTO chat_messages (
|
||||
SELECT
|
||||
@chat_id::uuid,
|
||||
NULLIF(UNNEST(@created_by::uuid[]), '00000000-0000-0000-0000-000000000000'::uuid),
|
||||
NULLIF(UNNEST(@api_key_id::text[]), ''),
|
||||
NULLIF(UNNEST(@model_config_id::uuid[]), '00000000-0000-0000-0000-000000000000'::uuid),
|
||||
UNNEST(@role::chat_message_role[]),
|
||||
UNNEST(@content::text[])::jsonb,
|
||||
@@ -1688,11 +1690,12 @@ RETURNING
|
||||
*;
|
||||
|
||||
-- name: InsertChatQueuedMessage :one
|
||||
INSERT INTO chat_queued_messages (chat_id, content, model_config_id)
|
||||
INSERT INTO chat_queued_messages (chat_id, content, model_config_id, api_key_id)
|
||||
VALUES (
|
||||
@chat_id,
|
||||
@content,
|
||||
sqlc.narg('model_config_id')::uuid
|
||||
sqlc.narg('model_config_id')::uuid,
|
||||
sqlc.narg('api_key_id')::text
|
||||
)
|
||||
RETURNING *;
|
||||
|
||||
|
||||
Reference in New Issue
Block a user