8000 feat: record model thoughts by dannykopping · Pull Request #22676 · coder/coder · GitHub
[go: up one dir, main page]

Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions coderd/database/dbauthz/dbauthz.go
Original file line number Diff line number Diff line change
Expand Up @@ -4441,6 +4441,13 @@ func (q *querier) InsertAIBridgeInterception(ctx context.Context, arg database.I
return insert(q.log, q.auth, rbac.ResourceAibridgeInterception.WithOwner(arg.InitiatorID.String()), q.db.InsertAIBridgeInterception)(ctx, arg)
}

func (q *querier) InsertAIBridgeModelThought(ctx context.Context, arg database.InsertAIBridgeModelThoughtParams) (database.AIBridgeModelThought, error) {
if err := q.authorizeAIBridgeInterceptionAction(ctx, policy.ActionUpdate, arg.InterceptionID); err != nil {
return database.AIBridgeModelThought{}, err
}
return q.db.InsertAIBridgeModelThought(ctx, arg)
}

func (q *querier) InsertAIBridgeTokenUsage(ctx context.Context, arg database.InsertAIBridgeTokenUsageParams) (database.AIBridgeTokenUsage, error) {
// All aibridge_token_usages records belong to the initiator of their associated interception.
if err := q.authorizeAIBridgeInterceptionAction(ctx, policy.ActionUpdate, arg.InterceptionID); err != nil {
Expand Down
11 changes: 11 additions & 0 deletions coderd/database/dbauthz/dbauthz_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -5027,6 +5027,17 @@ func (s *MethodTestSuite) TestAIBridge() {
check.Args(params).Asserts(intc, policy.ActionCreate)
}))

s.Run("InsertAIBridgeModelThought", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
intID := uuid.UUID{2}
intc := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: intID})
db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), intID).Return(intc, nil).AnyTimes() // Validation.

params := database.InsertAIBridgeModelThoughtParams{InterceptionID: intc.ID}
expected := testutil.Fake(s.T(), faker, database.AIBridgeModelThought{InterceptionID: intc.ID})
db.EXPECT().InsertAIBridgeModelThought(gomock.Any(), params).Return(expected, nil).AnyTimes()
check.Args(params).Asserts(intc, policy.ActionUpdate)
}))

s.Run("InsertAIBridgeTokenUsage", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
intID := uuid.UUID{2}
intc := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: intID})
Expand Down
8 changes: 8 additions & 0 deletions coderd/database/dbmetrics/querymetrics.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

15 changes: 15 additions & 0 deletions coderd/database/dbmock/dbmock.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

15 changes: 15 additions & 0 deletions coderd/database/dump.sql

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
DROP INDEX idx_aibridge_model_thoughts_interception_id;

DROP TABLE aibridge_model_thoughts;
11 changes: 11 additions & 0 deletions coderd/database/migrations/000440_aibridge_model_thoughts.up.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
CREATE TABLE aibridge_model_thoughts (
id UUID NOT NULL PRIMARY KEY,
interception_id UUID NOT NULL,
content TEXT NOT NULL,
metadata jsonb,
created_at TIMESTAMPTZ NOT NULL
);

COMMENT ON TABLE aibridge_model_thoughts IS 'Audit log of model thinking in intercepted requests in AI Bridge';

CREATE INDEX idx_aibridge_model_thoughts_interception_id ON aibridge_model_thoughts(interception_id);
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
INSERT INTO
aibridge_model_thoughts (
id,
interception_id,
content,
metadata,
created_at
)
VALUES (
'a1b2c3d4-e5f6-7890-abcd-ef1234567890',
'be003e1e-b38f-43bf-847d-928074dd0aa8', -- from 000370_aibridge.up.sql
'The user is asking about their workspaces. I should use the coder_list_workspaces tool to retrieve this information.',
'{"source": "commentary"}',
'2025-09-15 12:45:19.123456+00'
);
9 changes: 9 additions & 0 deletions coderd/database/models.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions coderd/database/querier.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

42 changes: 42 additions & 0 deletions coderd/database/queries.sql.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

14 changes: 14 additions & 0 deletions coderd/database/queries/aibridge.sql
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,14 @@ INSERT INTO aibridge_tool_usages (
)
RETURNING *;

-- name: InsertAIBridgeModelThought :one
INSERT INTO aibridge_model_thoughts (
id, interception_id, content, metadata, created_at
) VALUES (
@id, @interception_id, @content, COALESCE(@metadata::jsonb, '{}'::jsonb), @created_at
)
RETURNING *;

-- name: GetAIBridgeInterceptionByID :one
SELECT
*
Expand Down Expand Up @@ -362,6 +370,11 @@ WITH
WHERE started_at < @before_time::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)
Expand All @@ -384,6 +397,7 @@ WITH
)
-- Cumulative count.
SELECT (
(SELECT COUNT(*) FROM model_thoughts) +
(SELECT COUNT(*) FROM tool_usages) +
(SELECT COUNT(*) FROM token_usages) +
(SELECT COUNT(*) FROM user_prompts) +
Expand Down
1 change: 1 addition & 0 deletions coderd/database/sqlc.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -235,6 +235,7 @@ sql:
aibridge_tool_usage: AIBridgeToolUsage
aibridge_token_usage: AIBridgeTokenUsage
aibridge_user_prompt: AIBridgeUserPrompt
aibridge_model_thought: AIBridgeModelThought
rules:
- name: do-not-use-public-schema-in-queries
message: "do not use public schema in queries"
Expand Down
1 change: 1 addition & 0 deletions coderd/database/unique_constraint.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion docs/ai-coder/ai-bridge/monitoring.md
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
# Monitoring

AI Bridge records the last `user` prompt, token usage, and every tool invocation for each intercepted request. Each capture is tied to a single "interception" that maps back to the authenticated Coder identity, making it easy to attribute spend and behaviour.
AI Bridge records the last `user` prompt, token usage, model reasoning, and every tool invocation for each intercepted request. Each capture is tied to a single "interception" that maps back to the authenticated Coder identity, making it easy to attribute spend and behaviour.

![User Prompt logging](../../images/aibridge/grafana_user_prompts_logging.png)

Expand Down
2 changes: 1 addition & 1 deletion docs/ai-coder/ai-bridge/setup.md
Original file line number Diff line number Diff line change
Expand Up @@ -103,7 +103,7 @@ AI Bridge can relay traffic to other OpenAI- or Anthropic-compatible services or

## Data Retention

AI Bridge records prompts, token usage, and tool invocations for auditing and
AI Bridge records prompts, token usage, tool invocations, and model reasoning for auditing and
monitoring purposes. By default, this data is retained for **60 days**.

Configure retention using `--aibridge-retention` or `CODER_AIBRIDGE_RETENTION`:
Expand Down
15 changes: 15 additions & 0 deletions enterprise/aibridged/aibridgedmock/clientmock.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading
Loading
0