Skip to content

Commit 1ee2b77

Browse files
vegito22vegito22
authored
Move playground to use ws (#213)
1 parent 2630bcb commit 1ee2b77

File tree

6 files changed

+317
-72
lines changed

6 files changed

+317
-72
lines changed

‎llmstack/base/models.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,19 @@
2020
logger = logging.getLogger(__name__)
2121

2222

23+
def get_vendor_env_platform_defaults():
24+
return {
25+
"azure_openai_api_key": settings.DEFAULT_AZURE_OPENAI_API_KEY,
26+
"openai_api_key": settings.DEFAULT_OPENAI_API_KEY,
27+
"stabilityai_api_key": settings.DEFAULT_DREAMSTUDIO_API_KEY,
28+
"cohere_api_key": settings.DEFAULT_COHERE_API_KEY,
29+
"forefrontai_api_key": settings.DEFAULT_FOREFRONTAI_API_KEY,
30+
"elevenlabs_api_key": settings.DEFAULT_ELEVENLABS_API_KEY,
31+
"google_service_account_json_key": settings.DEFAULT_GOOGLE_SERVICE_ACCOUNT_JSON_KEY,
32+
"aws_access_key_id": settings.DEFAULT_AWS_ACCESS_KEY_ID,
33+
}
34+
35+
2336
class VectorstoreEmbeddingEndpoint(Enum):
2437
OPEN_AI = "openai"
2538
AZURE_OPEN_AI = "azure_openai"

‎llmstack/client/src/pages/playground.jsx‎

Lines changed: 220 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { Box, Button, Grid, Stack } from "@mui/material";
2-
import { useEffect, useState } from "react";
2+
import { useEffect, useState, useRef, useCallback } from "react";
33
import { useRecoilState, useRecoilValue } from "recoil";
44
import ApiBackendSelector from "../components/ApiBackendSelector";
55
import ConfigForm from "../components/ConfigForm";
@@ -13,11 +13,21 @@ import {
1313
isLoggedInState,
1414
templateValueState,
1515
} from "../data/atoms";
16-
import { axios } from "../data/axios";
16+
import {
17+
Messages,
18+
AppErrorMessage,
19+
AppMessage,
20+
} from "../components/apps/renderer/Messages";
21+
import { Ws } from "../data/ws";
22+
import { stitchObjects } from "../data/utils";
1723

1824
export default function PlaygroundPage() {
1925
const isLoggedIn = useRecoilValue(isLoggedInState);
2026
const [input] = useRecoilState(inputValueState);
27+
const appSessionId = useRef(null);
28+
const messagesRef = useRef(new Messages());
29+
const chunkedOutput = useRef({});
30+
const [showLoginDialog, setShowLoginDialog] = useState(false);
2131

2232
const [apiBackendSelected, setApiBackendSelected] = useRecoilState(
2333
apiBackendSelectedState,
@@ -35,72 +45,224 @@ export default function PlaygroundPage() {
3545
const [tokenCount, setTokenCount] = useState(null);
3646
const [processorResult, setProcessorResult] = useState(null);
3747

38-
const run = () => {
39-
setRunError("");
40-
setOutputLoading(true);
41-
axios()
42-
.post(`/api/playground/run`, {
43-
input: input,
44-
config: paramValues,
45-
bypass_cache: true,
46-
api_backend_slug: apiBackendSelected.slug,
47-
api_provider_slug: apiBackendSelected.api_provider.slug,
48-
})
49-
.then((response) => {
50-
if (response?.data?.errors) {
51-
setOutput("");
52-
setRunError(JSON.stringify(response?.data?.errors));
53-
}
48+
const [ws, setWs] = useState(null);
49+
const [appRunData, setAppRunData] = useState({});
50+
51+
const wsUrlPrefix = `${
52+
window.location.protocol === "https:" ? "wss" : "ws"
53+
}://${
54+
process.env.NODE_ENV === "development"
55+
? process.env.REACT_APP_API_SERVER || "localhost:9000"
56+
: window.location.host
57+
}/ws`;
58+
59+
useEffect(() => {
60+
if (!ws) {
61+
setWs(new Ws(`${wsUrlPrefix}/playground`));
62+
}
63+
}, [ws, wsUrlPrefix]);
64+
65+
if (ws) {
66+
ws.setOnMessage((evt) => {
67+
const message = JSON.parse(evt.data);
68+
69+
if (message.session) {
70+
appSessionId.current = message.session.id;
71+
72+
// Add messages from the session to the message list
73+
setAppRunData((prevState) => {
74+
prevState?.playground_messages?.forEach((message) => {
75+
messagesRef.current.add(message);
76+
});
5477

55-
setProcessorResult(response?.data?.output);
56-
if (response?.data?.output?.generations) {
57-
setOutput(response?.data?.output?.generations);
58-
} else if (response?.data?.output?.chat_completions) {
59-
setOutput(response?.data?.output?.chat_completions);
60-
} else {
61-
setOutput([response?.data?.output]);
78+
return {
79+
...prevState,
80+
sessionId: message.session.id,
81+
};
82+
});
83+
}
84+
85+
if (message.event && message.event === "done") {
86+
setAppRunData((prevState) => ({
87+
...prevState,
88+
isRunning: false,
89+
isStreaming: false,
90+
}));
91+
92+
chunkedOutput.current = {};
93+
}
94+
95+
if (message.event && message.event === "ratelimited") {
96+
messagesRef.current.add(
97+
new AppErrorMessage(
98+
null,
99+
message.request_id,
100+
"Rate limit exceeded. Please try after sometime.",
101+
),
102+
);
103+
104+
setAppRunData((prevState) => ({
105+
...prevState,
106+
isRunning: false,
107+
isStreaming: false,
108+
isRateLimited: true,
109+
errors: ["Rate limit exceeded"],
110+
playground_messages: messagesRef.current.get(),
111+
}));
112+
}
113+
114+
if (message.event && message.event === "usagelimited") {
115+
messagesRef.current.add(
116+
new AppErrorMessage(
117+
null,
118+
message.request_id,
119+
isLoggedIn
120+
? "Usage limit exceeded. Please try after adding more credits."
121+
: "Usage limit exceeded. Please login to continue.",
122+
),
123+
);
124+
125+
setAppRunData((prevState) => ({
126+
...prevState,
127+
isRunning: false,
128+
isStreaming: false,
129+
isUsageLimited: true,
130+
errors: ["Usage limit exceeded"],
131+
playground_messages: messagesRef.current.get(),
132+
}));
133+
134+
// If the user is not logged in, show the login dialog
135+
if (!isLoggedIn) {
136+
setShowLoginDialog(true);
62137
}
138+
}
139+
140+
if (message.errors && message.errors.length > 0) {
141+
message.errors.forEach((error) => {
142+
messagesRef.current.add(
143+
new AppErrorMessage(null, message.request_id, error),
144+
);
145+
});
63146

64-
setOutputLoading(false);
65-
66-
// Update token count
67-
if (
68-
response?.data?.output?._response?.api_response?.usage !== undefined
69-
) {
70-
if (
71-
response?.data?.output?._response?.api_response?.usage
72-
.prompt_tokens !== undefined &&
73-
response?.data?.output?._response?.api_response?.usage
74-
.completion_tokens !== undefined
75-
) {
76-
setTokenCount(
77-
`${response?.data?.output?._response?.api_response?.usage.total_tokens} (P${response?.data?.output?._response?.api_response?.usage.prompt_tokens} + C${response?.data?.output?._response?.api_response?.usage.completion_tokens})`,
78-
);
79-
} else if (
80-
response?.data?.output?._response?.api_response?.usage
81-
.total_tokens !== undefined
82-
) {
83-
setTokenCount(
84-
response?.data?.output?._response?.api_response?.usage
85-
.total_tokens,
86-
);
147+
setAppRunData((prevState) => ({
148+
...prevState,
149+
isRunning: false,
150+
isStreaming: false,
151+
errors: message.errors,
152+
playground_messages: messagesRef.current.get(),
153+
}));
154+
chunkedOutput.current = {};
155+
}
156+
157+
// Merge the new output with the existing output
158+
if (message.output) {
159+
let newChunkedOutput = {};
160+
newChunkedOutput = stitchObjects(chunkedOutput.current, message.output);
161+
chunkedOutput.current = newChunkedOutput;
162+
}
163+
164+
if (message.id && message.output) {
165+
const newMessage = message.output;
166+
messagesRef.current.add(
167+
new AppMessage(
168+
message.id,
169+
message.request_id,
170+
message.output,
171+
message.reply_to,
172+
),
173+
);
174+
setAppRunData((prevState) => ({
175+
...prevState,
176+
playground_messages: messagesRef.current.get(),
177+
isStreaming: newMessage.content !== null,
178+
}));
179+
}
180+
});
181+
}
182+
183+
useEffect(() => {
184+
if (appRunData && !appRunData?.isRunning && !appRunData?.isStreaming) {
185+
if (appRunData?.playground_messages) {
186+
const lastMessage =
187+
appRunData?.playground_messages[
188+
appRunData?.playground_messages.length - 1
189+
];
190+
if (lastMessage) {
191+
setOutputLoading(false);
192+
setProcessorResult(lastMessage?.content?.output);
193+
if (lastMessage?.content?.output) {
194+
if (lastMessage?.content?.output?.generations) {
195+
setOutput(lastMessage?.content?.output?.generations);
196+
} else if (lastMessage?.content?.output?.chat_completions) {
197+
setOutput(lastMessage?.content?.output?.chat_completions);
198+
} else {
199+
setOutput([lastMessage?.content?.output]);
200+
}
201+
}
202+
if (lastMessage?.content?.errors) {
203+
setRunError(lastMessage?.content?.errors);
87204
}
88205
}
89-
})
90-
.catch((error) => {
91-
console.error(error);
92-
setRunError(error?.response?.data || "Failed to run given prompt");
93-
setOutputLoading(false);
94-
});
95-
};
206+
}
207+
}
208+
}, [appRunData]);
209+
210+
const runApp = useCallback(
211+
(sessionId, input) => {
212+
setRunError("");
213+
setOutputLoading(true);
214+
215+
chunkedOutput.current = {};
216+
const requestId = Math.random().toString(36).substring(2);
217+
218+
setAppRunData((prevState) => ({
219+
...prevState,
220+
isRunning: true,
221+
isStreaming: false,
222+
errors: null,
223+
playground_messages: messagesRef.current.get(),
224+
input,
225+
}));
226+
227+
ws.send(
228+
JSON.stringify({
229+
event: "run",
230+
input,
231+
id: requestId,
232+
session_id: sessionId,
233+
}),
234+
);
235+
},
236+
[ws, setAppRunData],
237+
);
238+
239+
const cancelAppRun = useCallback(() => {
240+
setAppRunData((prevState) => ({
241+
...prevState,
242+
isRunning: false,
243+
}));
244+
245+
if (ws && ws.ws) {
246+
ws.send(
247+
JSON.stringify({
248+
event: "stop",
249+
}),
250+
);
251+
}
252+
}, [ws, setAppRunData]);
253+
96254
const Run = () => {
97255
return (
98256
<Button
99257
type="primary"
100258
onClick={(e) => {
101-
if (isLoggedIn) {
102-
return run();
103-
}
259+
runApp(appSessionId.current, {
260+
input: input,
261+
config: paramValues,
262+
bypass_cache: true,
263+
api_backend_slug: apiBackendSelected.slug,
264+
api_provider_slug: apiBackendSelected.api_provider.slug,
265+
});
104266
}}
105267
variant="contained"
106268
>
@@ -109,7 +271,6 @@ export default function PlaygroundPage() {
109271
);
110272
};
111273

112-
// Reset endpointSelected, apiBackendSelected, promptValues, paramValues and output on load
113274
useEffect(() => {
114275
setTokenCount(null);
115276
setOutput("");

0 commit comments

Comments
 (0)