mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user