create sampling response form

This commit is contained in:
Nathan Arseneau
2025-04-01 18:02:29 -04:00
parent 3032a67d4e
commit 71bb89ddf2
6 changed files with 316 additions and 33 deletions

View File

@@ -36,6 +36,7 @@ interface DynamicJsonFormProps {
value: JsonValue; value: JsonValue;
onChange: (value: JsonValue) => void; onChange: (value: JsonValue) => void;
maxDepth?: number; maxDepth?: number;
defaultIsJsonMode?: boolean;
} }
const DynamicJsonForm = ({ const DynamicJsonForm = ({
@@ -43,8 +44,9 @@ const DynamicJsonForm = ({
value, value,
onChange, onChange,
maxDepth = 3, maxDepth = 3,
defaultIsJsonMode = false,
}: DynamicJsonFormProps) => { }: DynamicJsonFormProps) => {
const [isJsonMode, setIsJsonMode] = useState(false); const [isJsonMode, setIsJsonMode] = useState(defaultIsJsonMode);
const [jsonError, setJsonError] = useState<string>(); const [jsonError, setJsonError] = useState<string>();
// Store the raw JSON string to allow immediate feedback during typing // Store the raw JSON string to allow immediate feedback during typing
// while deferring parsing until the user stops typing // while deferring parsing until the user stops typing
@@ -370,11 +372,21 @@ const DynamicJsonForm = ({
<div className="space-y-4"> <div className="space-y-4">
<div className="flex justify-end space-x-2"> <div className="flex justify-end space-x-2">
{isJsonMode && ( {isJsonMode && (
<Button variant="outline" size="sm" onClick={formatJson}> <Button
type="button"
variant="outline"
size="sm"
onClick={formatJson}
>
Format JSON Format JSON
</Button> </Button>
)} )}
<Button variant="outline" size="sm" onClick={handleSwitchToFormMode}> <Button
type="button"
variant="outline"
size="sm"
onClick={handleSwitchToFormMode}
>
{isJsonMode ? "Switch to Form" : "Switch to JSON"} {isJsonMode ? "Switch to Form" : "Switch to JSON"}
</Button> </Button>
</div> </div>

View File

@@ -0,0 +1,162 @@
import { Button } from "@/components/ui/button";
import JsonView from "./JsonView";
import { useMemo, useState } from "react";
import {
CreateMessageResult,
CreateMessageResultSchema,
} from "@modelcontextprotocol/sdk/types.js";
import { PendingRequest } from "./SamplingTab";
import DynamicJsonForm, { JsonSchemaType, JsonValue } from "./DynamicJsonForm";
import { useToast } from "@/hooks/use-toast";
export type SamplingRequestProps = {
request: PendingRequest;
onApprove: (id: number, result: CreateMessageResult) => void;
onReject: (id: number) => void;
};
const SamplingRequest = ({
onApprove,
request,
onReject,
}: SamplingRequestProps) => {
const { toast } = useToast();
const [messageResult, setMessageResult] = useState<JsonValue>({
model: "GPT-4o",
stopReason: "endTurn",
role: "assistant",
content: {
type: "text",
text: "",
},
});
const s = useMemo(() => {
const schema: JsonSchemaType = {
type: "object",
description: "Message result",
properties: {
model: {
type: "string",
default: "GPT-4o",
description: "model name",
},
stopReason: {
type: "string",
default: "endTurn",
description: "Stop reason",
},
role: {
type: "string",
default: "endTurn",
description: "Role of the model",
},
content: {
type: "object",
properties: {
type: {
type: "string",
default: "text",
description: "Type of content",
},
},
},
},
};
const contentType = (messageResult as any)?.content?.type;
if (contentType === "text" && schema.properties) {
schema.properties.content.properties = {
...schema.properties.content.properties,
text: {
type: "string",
default: "",
description: "text content",
},
};
setMessageResult((prev) => ({
...(prev as { [key: string]: JsonValue }),
content: {
type: contentType,
text: "",
},
}));
} else if (contentType === "image" && schema.properties) {
schema.properties.content.properties = {
...schema.properties.content.properties,
data: {
type: "string",
default: "",
description: "Base64 encoded image data",
},
mimeType: {
type: "string",
default: "",
description: "Mime type of the image",
},
};
setMessageResult((prev) => ({
...(prev as { [key: string]: JsonValue }),
content: {
type: contentType,
data: "",
mimeType: "",
},
}));
}
return schema;
}, [(messageResult as any)?.content?.type]);
const handleApprove = (id: number) => {
const validationResult = CreateMessageResultSchema.safeParse(messageResult);
if (!validationResult.success) {
toast({
title: "Error",
description: `There was an error validating the message result: ${validationResult.error.message}`,
variant: "destructive",
});
return;
}
onApprove(id, validationResult.data);
};
return (
<div
data-testid="sampling-request"
className="flex gap-4 p-4 border rounded-lg space-y-4"
>
<div className="flex-1 bg-gray-50 dark:bg-gray-800 dark:text-gray-100 p-2 rounded">
<JsonView data={JSON.stringify(request.request)} />
</div>
<form className="flex-1 space-y-4">
<div className="space-y-2">
<DynamicJsonForm
defaultIsJsonMode={true}
schema={s}
value={messageResult}
onChange={(newValue: JsonValue) => {
setMessageResult(newValue);
}}
/>
</div>
<div className="flex space-x-2 mt-1">
<Button type="button" onClick={() => handleApprove(request.id)}>
Approve
</Button>
<Button
type="button"
variant="outline"
onClick={() => onReject(request.id)}
>
Reject
</Button>
</div>
</form>
</div>
);
};
export default SamplingRequest;

View File

@@ -1,11 +1,10 @@
import { Alert, AlertDescription } from "@/components/ui/alert"; import { Alert, AlertDescription } from "@/components/ui/alert";
import { Button } from "@/components/ui/button";
import { TabsContent } from "@/components/ui/tabs"; import { TabsContent } from "@/components/ui/tabs";
import { import {
CreateMessageRequest, CreateMessageRequest,
CreateMessageResult, CreateMessageResult,
} from "@modelcontextprotocol/sdk/types.js"; } from "@modelcontextprotocol/sdk/types.js";
import JsonView from "./JsonView"; import SamplingRequest from "./SamplingRequest";
export type PendingRequest = { export type PendingRequest = {
id: number; id: number;
@@ -19,21 +18,8 @@ export type Props = {
}; };
const SamplingTab = ({ pendingRequests, onApprove, onReject }: Props) => { const SamplingTab = ({ pendingRequests, onApprove, onReject }: Props) => {
const handleApprove = (id: number) => {
// For now, just return a stub response
onApprove(id, {
model: "stub-model",
stopReason: "endTurn",
role: "assistant",
content: {
type: "text",
text: "This is a stub response.",
},
});
};
return ( return (
<TabsContent value="sampling" className="h-96"> <TabsContent value="sampling" className="mh-96">
<Alert> <Alert>
<AlertDescription> <AlertDescription>
When the server requests LLM sampling, requests will appear here for When the server requests LLM sampling, requests will appear here for
@@ -43,19 +29,12 @@ const SamplingTab = ({ pendingRequests, onApprove, onReject }: Props) => {
<div className="mt-4 space-y-4"> <div className="mt-4 space-y-4">
<h3 className="text-lg font-semibold">Recent Requests</h3> <h3 className="text-lg font-semibold">Recent Requests</h3>
{pendingRequests.map((request) => ( {pendingRequests.map((request) => (
<div key={request.id} className="p-4 border rounded-lg space-y-4"> <SamplingRequest
<JsonView key={request.id}
className="bg-gray-50 dark:bg-gray-800 dark:text-gray-100 rounded" request={request}
data={JSON.stringify(request.request)} onApprove={onApprove}
/> onReject={onReject}
/>
<div className="flex space-x-2">
<Button onClick={() => handleApprove(request.id)}>Approve</Button>
<Button variant="outline" onClick={() => onReject(request.id)}>
Reject
</Button>
</div>
</div>
))} ))}
{pendingRequests.length === 0 && ( {pendingRequests.length === 0 && (
<p className="text-gray-500">No pending requests</p> <p className="text-gray-500">No pending requests</p>

View File

@@ -460,7 +460,9 @@ const Sidebar = ({
</SelectTrigger> </SelectTrigger>
<SelectContent> <SelectContent>
{Object.values(LoggingLevelSchema.enum).map((level) => ( {Object.values(LoggingLevelSchema.enum).map((level) => (
<SelectItem value={level}>{level}</SelectItem> <SelectItem key={level} value={level}>
{level}
</SelectItem>
))} ))}
</SelectContent> </SelectContent>
</Select> </Select>

View File

@@ -0,0 +1,73 @@
import { render, screen, fireEvent } from "@testing-library/react";
import SamplingRequest from "../SamplingRequest";
import { PendingRequest } from "../SamplingTab";
const mockRequest: PendingRequest = {
id: 1,
request: {
method: "sampling/createMessage",
params: {
messages: [
{
role: "user",
content: {
type: "text",
text: "What files are in the current directory?",
},
},
],
systemPrompt: "You are a helpful file system assistant.",
includeContext: "thisServer",
maxTokens: 100,
},
},
};
describe("Form to handle sampling response", () => {
const mockOnApprove = jest.fn();
const mockOnReject = jest.fn();
afterEach(() => {
jest.clearAllMocks();
});
it("should call onApprove with correct text content when Approve button is clicked", () => {
render(
<SamplingRequest
request={mockRequest}
onApprove={mockOnApprove}
onReject={mockOnReject}
/>,
);
// Click the Approve button
fireEvent.click(screen.getByRole("button", { name: /approve/i }));
// Assert that onApprove is called with the correct arguments
expect(mockOnApprove).toHaveBeenCalledWith(mockRequest.id, {
model: "GPT-4o",
stopReason: "endTurn",
role: "assistant",
content: {
type: "text",
text: "",
},
});
});
it("should call onReject with correct request id when Reject button is clicked", () => {
render(
<SamplingRequest
request={mockRequest}
onApprove={mockOnApprove}
onReject={mockOnReject}
/>,
);
// Click the Approve button
fireEvent.click(screen.getByRole("button", { name: /Reject/i }));
// Assert that onApprove is called with the correct arguments
expect(mockOnReject).toHaveBeenCalledWith(mockRequest.id);
});
});

View File

@@ -0,0 +1,55 @@
import { render, screen } from "@testing-library/react";
import { Tabs } from "@/components/ui/tabs";
import SamplingTab, { PendingRequest } from "../SamplingTab";
describe("Sampling tab", () => {
const mockOnApprove = jest.fn();
const mockOnReject = jest.fn();
const renderSamplingTab = (pendingRequests: PendingRequest[]) =>
render(
<Tabs defaultValue="sampling">
<SamplingTab
pendingRequests={pendingRequests}
onApprove={mockOnApprove}
onReject={mockOnReject}
/>
</Tabs>,
);
it("should render 'No pending requests' when there are no pending requests", () => {
renderSamplingTab([]);
expect(
screen.getByText(
"When the server requests LLM sampling, requests will appear here for approval.",
),
).toBeTruthy();
expect(screen.findByText("No pending requests")).toBeTruthy();
});
it("should render the correct number of requests", () => {
renderSamplingTab(
Array.from({ length: 5 }, (_, i) => ({
id: i,
request: {
method: "sampling/createMessage",
params: {
messages: [
{
role: "user",
content: {
type: "text",
text: "What files are in the current directory?",
},
},
],
systemPrompt: "You are a helpful file system assistant.",
includeContext: "thisServer",
maxTokens: 100,
},
},
})),
);
expect(screen.getAllByTestId("sampling-request").length).toBe(5);
});
});