feat: record model thoughts (#22676)

Depends on https://github.com/coder/aibridge/pull/203
Closes https://github.com/coder/internal/issues/1337

---------

Signed-off-by: Danny Kopping <danny@coder.com>
This commit is contained in:
Danny Kopping
2026-03-17 11:41:10 +00:00
committed by GitHub
parent 5d0eb772da
commit 365de3e367
24 changed files with 723 additions and 208 deletions
+39
View File
@@ -340,6 +340,11 @@ WITH
WHERE started_at < $1::timestamp with time zone
),
-- CTEs are executed in order.
model_thoughts AS (
DELETE FROM aibridge_model_thoughts
WHERE interception_id IN (SELECT id FROM to_delete)
RETURNING 1
),
tool_usages AS (
DELETE FROM aibridge_tool_usages
WHERE interception_id IN (SELECT id FROM to_delete)
@@ -361,6 +366,7 @@ WITH
RETURNING 1
)
SELECT (
(SELECT COUNT(*) FROM model_thoughts) +
(SELECT COUNT(*) FROM tool_usages) +
(SELECT COUNT(*) FROM token_usages) +
(SELECT COUNT(*) FROM user_prompts) +
@@ -661,6 +667,39 @@ func (q *sqlQuerier) InsertAIBridgeInterception(ctx context.Context, arg InsertA
return i, err
}
const insertAIBridgeModelThought = `-- name: InsertAIBridgeModelThought :one
INSERT INTO aibridge_model_thoughts (
interception_id, content, metadata, created_at
) VALUES (
$1, $2, COALESCE($3::jsonb, '{}'::jsonb), $4
)
RETURNING interception_id, content, metadata, created_at
`
type InsertAIBridgeModelThoughtParams struct {
InterceptionID uuid.UUID `db:"interception_id" json:"interception_id"`
Content string `db:"content" json:"content"`
Metadata json.RawMessage `db:"metadata" json:"metadata"`
CreatedAt time.Time `db:"created_at" json:"created_at"`
}
func (q *sqlQuerier) InsertAIBridgeModelThought(ctx context.Context, arg InsertAIBridgeModelThoughtParams) (AIBridgeModelThought, error) {
row := q.db.QueryRowContext(ctx, insertAIBridgeModelThought,
arg.InterceptionID,
arg.Content,
arg.Metadata,
arg.CreatedAt,
)
var i AIBridgeModelThought
err := row.Scan(
&i.InterceptionID,
&i.Content,
&i.Metadata,
&i.CreatedAt,
)
return i, err
}
const insertAIBridgeTokenUsage = `-- name: InsertAIBridgeTokenUsage :one
INSERT INTO aibridge_token_usages (
id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at