Add API to download judge votes (#53)

consoull · gordonhart · web-flow · commit e4df7149141f · 2024-09-18T22:11:35.000Z
* Add API to download judge votes

Adds an endpoint to download a judges votes under
`/project/{project_slug}/judge/{judge_id}/download/votes`.
The exported CSV contains the prompt, two responses and the judges
choice.

* Connect to frontend

* Disable download button when no votes are present

* Improve enabled/disabled pill display

* Add test for judge vote download endpoint

* Add second vote to test and catch bug where responses were not ordered

---------

Co-authored-by: Gordon Hart &lt;gordon.hart2@gmail.com&gt;
diff --git a/autoarena/api/router.py b/autoarena/api/router.py
@@ -188,6 +188,18 @@ def delete_judge(project_slug: str, judge_id: int, background_tasks: BackgroundT
         except NotFoundError:
             pass
 
+    @r.get("/project/{project_slug}/judge/{judge_id}/download/votes")
+    async def download_judge_votes_csv(project_slug: str, judge_id: int) -> StreamingResponse:
+        columns = ["prompt", "model_a", "model_b", "response_a", "response_b", "winner"]
+        df_response = JudgeService.get_df_vote(project_slug, judge_id)
+        # TODO: handle case where no votes exist, not a big problem for now as UI buttons are disabled for 0-vote judges
+        judge_name = df_response.iloc[0].judge
+        stream = StringIO()
+        df_response[columns].to_csv(stream, index=False)
+        response = StreamingResponse(iter([stream.getvalue()]), media_type="text/csv")
+        response.headers["Content-Disposition"] = f'attachment; filename="{judge_name}.csv"'
+        return response
+
     @r.put("/project/{project_slug}/elo/reseed-scores")
     def reseed_elo_scores(project_slug: str) -> None:
         EloService.reseed_scores(project_slug)
diff --git a/autoarena/service/judge.py b/autoarena/service/judge.py
@@ -1,4 +1,5 @@
 from loguru import logger
+import pandas as pd
 
 from autoarena.api import api
 from autoarena.judge.factory import verify_judge_type_environment
@@ -44,6 +45,32 @@ def get_all(project_slug: str) -> list[api.Judge]:
         df["judge_type"] = df["judge_type"].apply(lambda j: j if j in judge_types else api.JudgeType.UNRECOGNIZED.value)
         return [api.Judge(**r) for _, r in df.iterrows()]
 
+    @staticmethod
+    def get_df_vote(project_slug: str, judge_id: int) -> pd.DataFrame:
+        with ProjectService.connect(project_slug) as conn:
+            df_vote = conn.execute(
+                """
+                SELECT
+                    j.name as judge,
+                    ra.prompt as prompt,
+                    ma.name as model_a,
+                    mb.name as model_b,
+                    ra.response as response_a,
+                    rb.response as response_b,
+                    h2h.winner as winner
+                FROM judge j
+                JOIN head_to_head h2h ON j.id = h2h.judge_id
+                JOIN response ra ON ra.id = h2h.response_a_id
+                JOIN response rb ON rb.id = h2h.response_b_id
+                JOIN model ma ON ra.model_id = ma.id
+                JOIN model mb ON rb.model_id = mb.id
+                WHERE j.id = $judge_id
+                ORDER BY h2h.id
+                """,
+                dict(judge_id=judge_id),
+            ).df()
+        return df_vote
+
     @staticmethod
     def create(project_slug: str, request: api.CreateJudgeRequest) -> api.Judge:
         with ProjectService.connect(project_slug) as conn:
diff --git a/tests/integration/api/test_judges.py b/tests/integration/api/test_judges.py
@@ -1,3 +1,6 @@
+from io import StringIO
+
+import pandas as pd
 import pytest
 from fastapi.testclient import TestClient
 
@@ -51,6 +54,32 @@ def test__judges__can_access__unrecognized__failed(project_client: TestClient) -
     assert not project_client.get("/judge/unrecognized/can-access").json()
 
 
+def test__judges__download_votes_csv(project_client: TestClient, model_id: int, model_b_id: int) -> None:
+    h2h = project_client.put("/head-to-heads", json=dict(model_a_id=model_id, model_b_id=model_b_id)).json()
+    judge_request = dict(response_a_id=h2h[0]["response_a_id"], response_b_id=h2h[0]["response_b_id"], winner="A")
+    assert project_client.post("/head-to-head/vote", json=judge_request).json() is None
+    judge_request = dict(response_a_id=h2h[1]["response_b_id"], response_b_id=h2h[1]["response_a_id"], winner="-")
+    assert project_client.post("/head-to-head/vote", json=judge_request).json() is None
+    judges = project_client.get("/judges").json()
+    assert len(judges) == 1
+    assert judges[0]["n_votes"] == 2
+    human_judge_id = judges[0]["id"]
+    response = project_client.get(f"/judge/{human_judge_id}/download/votes")
+    assert response.status_code == 200
+    models = project_client.get("/models").json()
+    assert len(models) == 2
+    model_a, model_b = [m for m in models if m["id"] == model_id][0], [m for m in models if m["id"] == model_b_id][0]
+    df_vote_expected = pd.DataFrame(
+        [
+            (h2h[0]["prompt"], model_a["name"], model_b["name"], h2h[0]["response_a"], h2h[0]["response_b"], "A"),
+            (h2h[1]["prompt"], model_b["name"], model_a["name"], h2h[1]["response_b"], h2h[1]["response_a"], "-"),
+        ],
+        columns=["prompt", "model_a", "model_b", "response_a", "response_b", "winner"],
+    )
+    df_vote = pd.read_csv(StringIO(response.text))
+    assert df_vote.equals(df_vote_expected)
+
+
 def test__judges__delete(project_client: TestClient, judge_id: int) -> None:
     assert project_client.delete(f"/judge/{judge_id}").json() is None
     assert len(project_client.get("/judges").json()) == 1  # only default judge is left
diff --git a/ui/src/components/Judges/JudgeAccordionItem.tsx b/ui/src/components/Judges/JudgeAccordionItem.tsx
@@ -1,13 +1,14 @@
 import { useState } from 'react';
-import { Accordion, Button, Checkbox, Collapse, Group, Pill, Stack, Text, Tooltip } from '@mantine/core';
+import { Accordion, Anchor, Button, Checkbox, Collapse, Group, Pill, Stack, Text, Tooltip } from '@mantine/core';
 import { Link } from 'react-router-dom';
 import { useDisclosure } from '@mantine/hooks';
-import { IconGavel } from '@tabler/icons-react';
+import { IconDownload, IconGavel, IconPrompt } from '@tabler/icons-react';
 import { Judge } from '../../hooks/useJudges.ts';
 import { useUrlState } from '../../hooks/useUrlState.ts';
 import { useUpdateJudge } from '../../hooks/useUpdateJudge.ts';
 import { MarkdownContent } from '../MarkdownContent.tsx';
 import { pluralize } from '../../lib/string.ts';
+import { API_ROUTES } from '../../lib/routes.ts';
 import { judgeTypeIconComponent, judgeTypeToHumanReadableName } from './types.ts';
 import { DeleteJudgeButton } from './DeleteJudgeButton.tsx';
 import { CanAccessJudgeStatusIndicator } from './CanAccessJudgeStatusIndicator.tsx';
@@ -30,6 +31,15 @@ export function JudgeAccordionItem({ judge }: Props) {
     setIsEnabled(prev => !prev);
   }
 
+  const canDownload = judge.n_votes > 0;
+  const downloadUrl = canDownload ? API_ROUTES.downloadJudgeVotesCsv(projectSlug, judge.id) : undefined;
+  const DownloadVotesComponent = (
+    <Anchor href={downloadUrl} target="_blank">
+      <Button variant="light" color="teal" size="xs" leftSection={<IconDownload size={20} />} disabled={!canDownload}>
+        Download Votes CSV
+      </Button>
+    </Anchor>
+  );
   const IconComponent = judgeTypeIconComponent(judge_type);
   return (
     <Accordion.Item key={id} value={`${judge_type}-${id}`}>
@@ -53,10 +63,15 @@ export function JudgeAccordionItem({ judge }: Props) {
             </Text>
           </Stack>
           <Group>
-            {isEnabled && (
-              <Pill bg="ice.0" c="gray.8">
+            <Text c="dimmed" size="xs" fs="italic">
+              {pluralize(judge.n_votes, 'vote')}
+            </Text>
+            {isEnabled ? (
+              <Pill bg="ice.0" c="ice.9">
                 Enabled
               </Pill>
+            ) : (
+              <Pill c="gray">Disabled</Pill>
             )}
           </Group>
         </Group>
@@ -73,12 +88,16 @@ export function JudgeAccordionItem({ judge }: Props) {
                   onChange={() => handleToggleEnabled()}
                 />
                 <Group>
-                  <Text c="dimmed" size="xs" fs="italic">
-                    {pluralize(judge.n_votes, 'vote')} submitted
-                  </Text>
-                  <Button variant="light" color="gray" size="xs" onClick={toggleShowSystemPrompt}>
+                  <Button
+                    variant="light"
+                    color="gray"
+                    size="xs"
+                    leftSection={<IconPrompt size={20} />}
+                    onClick={toggleShowSystemPrompt}
+                  >
                     {showSystemPrompt ? 'Hide' : 'Show'} System Prompt
                   </Button>
+                  {DownloadVotesComponent}
                   <Tooltip label="Judge must be enabled" disabled={judge.enabled}>
                     <Button
                       variant="light"
@@ -111,9 +130,7 @@ export function JudgeAccordionItem({ judge }: Props) {
                 </Link>{' '}
                 tab to provide ratings on head-to-head matchups between models.
               </Text>
-              <Text c="dimmed" size="xs" fs="italic">
-                {pluralize(judge.n_votes, 'vote')} submitted
-              </Text>
+              {DownloadVotesComponent}
             </Group>
           )}
           <Collapse in={showSystemPrompt} fz="sm">
diff --git a/ui/src/components/Leaderboard/ExpandedModelDetails.tsx b/ui/src/components/Leaderboard/ExpandedModelDetails.tsx
@@ -64,12 +64,12 @@ export function ExpandedModelDetails({ model }: Props) {
               </Link>
               <Anchor href={API_ROUTES.downloadModelResponsesCsv(projectSlug, model.id)} target="_blank">
                 <Button color="teal" variant="light" size="xs" leftSection={<IconDownload size={20} />}>
-                  Download Response CSV
+                  Download Responses CSV
                 </Button>
               </Anchor>
               <Anchor href={API_ROUTES.downloadModelHeadToHeadsCsv(projectSlug, model.id)} target="_blank">
                 <Button color="teal" variant="light" size="xs" leftSection={<IconDownload size={20} />}>
-                  Download Head-to-Head CSV
+                  Download Head-to-Heads CSV
                 </Button>
               </Anchor>
               <Button
diff --git a/ui/src/lib/routes.ts b/ui/src/lib/routes.ts
@@ -60,6 +60,8 @@ export const API_ROUTES = {
   getDefaultSystemPrompt: (projectSlug: string) => `${getProjectApiUrl(projectSlug)}/judge/default-system-prompt`,
   createJudge: (projectSlug: string) => `${getProjectApiUrl(projectSlug)}/judge`,
   updateJudge: (projectSlug: string, judgeId: number) => `${getProjectApiUrl(projectSlug)}/judge/${judgeId}`,
+  downloadJudgeVotesCsv: (projectSlug: string, judgeId: number) =>
+    `${getProjectApiUrl(projectSlug)}/judge/${judgeId}/download/votes`,
   checkCanAccess: (projectSlug: string, judgeType: JudgeType) =>
     `${getProjectApiUrl(projectSlug)}/judge/${judgeType}/can-access`,
   deleteJudge: (projectSlug: string, judgeId: number) => `${getProjectApiUrl(projectSlug)}/judge/${judgeId}`,
