Allow mapping PRs to episode-level metrics (#29)

* Add clickhouse methods for PR - episode mapping

* TEMP: feedback

* Configure episode-level feedback metric

* Remove duplicate episode mapping function

* chore: automated fix for PR #29 (#30)

Co-authored-by: TensorZero-Experimental-CI-Bot[bot] <hello@tensorzero.com>

* Lint

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: TensorZero-Experimental-CI-Bot[bot] <hello@tensorzero.com>

Author: 3 people

.github/workflows/ci-failure-diagnosis.yml

@@ -67,6 +67,7 @@ jobs:
         with:
           token: ${{ secrets.GITHUB_TOKEN }}
           tensorzero-base-url: http://localhost:3000
+          # TODO: Remove when agent creates PRs
           tensorzero-diff-patched-successfully-metric-name: tensorzero_github_ci_bot_diff_patched_successfully
           output-artifacts-dir: debug-logs
           clickhouse-url: ${{ secrets.CI_BOT_CLICKHOUSE_URL }}

.github/workflows/provide-pull-request-feedback.yml

@@ -43,6 +43,7 @@ jobs:
         uses: tensorzero/experimental-ci-bot/create-pr-feedback@main
         with:
           tensorzero-base-url: http://localhost:3000
+          # TODO: Switch to tensorzero_github_ci_bot_agent_pr_merged for episode-level feedback when agent creates PRs
           tensorzero-pr-merged-metric-name: tensorzero_github_ci_bot_pr_merged
           clickhouse-url: ${{ secrets.CI_BOT_CLICKHOUSE_URL }}
           clickhouse-table: GitHubBotPullRequestToInferenceMap

AGENTS.md

@@ -52,5 +52,5 @@ alongside source when behavior changes.
 
 ## IMPORTANT RULES
 
-Before submitting a PR, run `npm run bundle` to update the javascript bundles.
+Before submitting a PR, run `npm run bundle` to update the JavaScript bundles.
 Otherwise, CI will fail.

README.md

@@ -38,6 +38,7 @@ CREATE TABLE GitHubBotPullRequestToInferenceMap
 (
    pull_request_id UInt128,
    inference_id String,
+   episode_id String,
    created_at DateTime DEFAULT now(),
    original_pull_request_url String
 )

dist/create-pr-feedback/index.js

@@ -4,7 +4,8 @@ import {
   type ClickHouseClientLike,
   type ClickHouseConfig,
   createPullRequestToInferenceRecord,
-  getPullRequestToInferenceRecords
+  getPullRequestToInferenceRecords,
+  getPullRequestToEpisodeRecords
 } from './clickhouseClient.js'
 
 const defaultConfig: ClickHouseConfig = {
@@ -31,6 +32,7 @@ describe('clickhouseClient', () => {
     await createPullRequestToInferenceRecord(
       {
         inferenceId: 'abc-123',
+        episodeId: 'episode-123',
         pullRequestId: 42,
         originalPullRequestUrl: 'https://github.com/org/repo/pull/42'
       },
@@ -42,6 +44,7 @@ describe('clickhouseClient', () => {
       table: 'tensorzero.inference_records',
       values: [
         {
+          episode_id: 'episode-123',
           pull_request_id: 42,
           inference_id: 'abc-123',
           original_pull_request_url: 'https://github.com/org/repo/pull/42'
@@ -86,6 +89,7 @@ describe('clickhouseClient', () => {
       createPullRequestToInferenceRecord(
         {
           inferenceId: 'abc',
+          episodeId: 'episode-123',
           pullRequestId: 1,
           originalPullRequestUrl: 'https://example.com/pr/1'
         },
@@ -99,6 +103,7 @@ describe('clickhouseClient', () => {
       createPullRequestToInferenceRecord(
         {
           inferenceId: 'abc',
+          episodeId: 'episode-123',
           pullRequestId: 1,
           originalPullRequestUrl: 'https://example.com/pr/1'
         },
@@ -114,6 +119,7 @@ describe('clickhouseClient', () => {
       createPullRequestToInferenceRecord(
         {
           inferenceId: 'abc',
+          episodeId: 'episode-123',
           pullRequestId: 1,
           originalPullRequestUrl: 'https://example.com/pr/1'
         },
@@ -124,4 +130,43 @@ describe('clickhouseClient', () => {
 
     expect(client.close).not.toHaveBeenCalled()
   })
+
+  it('queries episode records with parameter binding', async () => {
+    const expectedRecords = [
+      {
+        episode_id: 'episode-789',
+        pull_request_id: 55,
+        created_at: '2024-02-01T00:00:00Z',
+        original_pull_request_url: 'https://github.com/org/repo/pull/55'
+      }
+    ]
+    const jsonMock = jest
+      .fn<() => Promise<unknown>>()
+      .mockResolvedValue(expectedRecords)
+    // @ts-expect-error(Mock type is inaccurate)
+    client.query.mockResolvedValueOnce({ json: jsonMock })
+
+    const records = await getPullRequestToEpisodeRecords(55, defaultConfig, {
+      client
+    })
+
+    expect(client.query).toHaveBeenCalledWith({
+      query:
+        'SELECT episode_id, pull_request_id, created_at, original_pull_request_url FROM tensorzero.inference_records WHERE pull_request_id = {pullRequestId:UInt64}',
+      query_params: { pullRequestId: 55 },
+      format: 'JSONEachRow'
+    })
+    expect(jsonMock).toHaveBeenCalledTimes(1)
+    expect(records).toEqual(expectedRecords)
+  })
+
+  it('propagates query failures for episode records without closing the client', async () => {
+    client.query.mockRejectedValueOnce(new Error('query failed'))
+
+    await expect(
+      getPullRequestToEpisodeRecords(123, defaultConfig, { client })
+    ).rejects.toThrow('query failed')
+
+    expect(client.close).not.toHaveBeenCalled()
+  })
 })

dist/create-pr-feedback/index.js.map

@@ -16,6 +16,7 @@ export interface ClickHouseConfig {
 
 export interface CreatePullRequestToInferenceRequest {
   inferenceId: string
+  episodeId: string
   pullRequestId: number
   originalPullRequestUrl: string
 }
@@ -27,6 +28,13 @@ export interface PullRequestToInferenceRecord {
   original_pull_request_url: string
 }
 
+export interface PullRequestToEpisodeRecord {
+  episode_id: string
+  pull_request_id: number
+  created_at: string
+  original_pull_request_url: string
+}
+
 const CLICKHOUSE_TABLE_NAME_REGEX = /^[a-zA-Z0-9_.]+$/
 
 function assertValidTableName(table: string): void {
@@ -103,6 +111,7 @@ export async function createPullRequestToInferenceRecord(
         {
           pull_request_id: request.pullRequestId,
           inference_id: request.inferenceId,
+          episode_id: request.episodeId,
           original_pull_request_url: request.originalPullRequestUrl
         }
       ],
@@ -140,3 +149,29 @@ export async function getPullRequestToInferenceRecords(
   }
   return records
 }
+
+// Returns all episode records for a given pull request.
+export async function getPullRequestToEpisodeRecords(
+  pullRequestId: number,
+  config: ClickHouseConfig,
+  dependencies?: ClickHouseDependencies
+): Promise<PullRequestToEpisodeRecord[]> {
+  const { client, table, shouldClose } = createTensorZeroClickHouseClient(
+    config,
+    dependencies
+  )
+  let records: PullRequestToEpisodeRecord[] = []
+  try {
+    const response = await client.query({
+      query: `SELECT episode_id, pull_request_id, created_at, original_pull_request_url FROM ${table} WHERE pull_request_id = {pullRequestId:UInt64}`,
+      query_params: { pullRequestId },
+      format: 'JSONEachRow'
+    })
+    records = await response.json()
+  } finally {
+    if (shouldClose) {
+      await client.close()
+    }
+  }
+  return records
+}