Enable multiple llms
This commit is contained in:
Niels Geens
@@ -1,127 +1,187 @@
|
||||
import React, { useState, useEffect } from 'react';
|
||||
import { Header } from '@/components/Header';
|
||||
import { Header, Tabs } from '@/components/Header';
|
||||
import { Sidebar } from '@/components/Sidebar';
|
||||
import { FloatingButton } from '@/components/FloatingButton';
|
||||
import { ContentEditor } from '@/components/editors/ContentEditor';
|
||||
import { StandardEditor } from '@/components/editors/StandardEditor';
|
||||
import { useWebSocket } from '@/hooks/useWebSocket';
|
||||
import { AgentState, Tabs, ClientMessageType, ServerMessageType } from '@/constants';
|
||||
import { useWebSocket, WebSocketState } from '@/hooks/useWebSocket';
|
||||
|
||||
export const LlmState = {
|
||||
NO_OUTPUT: 'NO_OUTPUT',
|
||||
INFERENCE: 'INFERENCE',
|
||||
OUTPUT: 'OUTPUT'
|
||||
};
|
||||
|
||||
const App = () => {
|
||||
// Editor content state
|
||||
const [originalContext, setOriginalContext] = useState('');
|
||||
const [generatedContext, setGeneratedContext] = useState('');
|
||||
const [modifiedContext, setModifiedContext] = useState('');
|
||||
const [originalResponse, setOriginalResponse] = useState('');
|
||||
const [modifiedResponse, setModifiedResponse] = useState('');
|
||||
const [validationError, setValidationError] = useState(null);
|
||||
const [contextDirty, setContextDirty] = useState(false);
|
||||
const [generatedResponses, setGeneratedResponses] = useState({});
|
||||
const [modifiedResponses, setModifiedResponses] = useState({});
|
||||
const [input, setInput] = useState('');
|
||||
const [output, setOutput] = useState('');
|
||||
|
||||
// LLM state
|
||||
const [llms, setLlms] = useState({});
|
||||
const [activeLlm, setActiveLlm] = useState(null);
|
||||
|
||||
// UI state
|
||||
const [activeTab, setActiveTab] = useState(Tabs.CONTEXT);
|
||||
const [showDiff, setShowDiff] = useState(false);
|
||||
const [showSidebar, setShowSidebar] = useState(false);
|
||||
const [agentState, setAgentState] = useState(AgentState.CONTEXT_APPROVAL);
|
||||
const [autoApproverConfig, setAutoApproverConfig] = useState({
|
||||
context_enabled: false,
|
||||
response_enabled: false,
|
||||
context_timeout: 5.0,
|
||||
response_timeout: 10.0
|
||||
});
|
||||
|
||||
const { wsState, sendMessage, addMessageHandler } = useWebSocket(`${window.location.protocol === 'https:' ? 'wss:' : 'ws:'}//${window.location.host}/ws`);
|
||||
// WebSocket connections
|
||||
const wsRoot = `${window.location.protocol === 'https:' ? 'wss:' : 'ws:'}//${window.location.host}/ws`
|
||||
const llmWs = useWebSocket(`${wsRoot}/llm`);
|
||||
const contextWs = useWebSocket(`${wsRoot}/context`);
|
||||
const tokenWs = useWebSocket(`${wsRoot}/token`);
|
||||
const stdoutWs = useWebSocket(`${wsRoot}/stdout`);
|
||||
const [wsState, setWsState] = useState(WebSocketState.DISCONNECTED);
|
||||
|
||||
// Handle llm state changes
|
||||
useEffect(() => {
|
||||
addMessageHandler((message) => {
|
||||
switch(message.type) {
|
||||
case ServerMessageType.STATE_CHANGE:
|
||||
setAgentState(message.state);
|
||||
break;
|
||||
case ServerMessageType.CONTEXT_UPDATE:
|
||||
setOriginalContext(message.context);
|
||||
setModifiedContext(message.context);
|
||||
break;
|
||||
case ServerMessageType.RESPONSE_UPDATE:
|
||||
setOriginalResponse(message.response);
|
||||
setModifiedResponse(message.response);
|
||||
setValidationError(message.validation_error);
|
||||
break;
|
||||
case ServerMessageType.OUTPUT_UPDATE:
|
||||
setOutput(message.output);
|
||||
break;
|
||||
case ServerMessageType.AUTO_APPROVER_CONFIG:
|
||||
setAutoApproverConfig(message.config);
|
||||
break;
|
||||
llmWs.addMessageHandler((data) => {
|
||||
setLlms(prevLlms => ({
|
||||
...prevLlms,
|
||||
[data.llm]: data.state
|
||||
}));
|
||||
if (data.state === LlmState.NO_OUTPUT) {
|
||||
setGeneratedResponses(prev => {
|
||||
const updated = { ...prev };
|
||||
delete updated[data.llm];
|
||||
return updated;
|
||||
});
|
||||
setModifiedResponses(prev => {
|
||||
const updated = { ...prev };
|
||||
delete updated[data.llm];
|
||||
return updated;
|
||||
});
|
||||
}
|
||||
if (activeLlm === null) {
|
||||
setActiveLlm(data.llm);
|
||||
}
|
||||
});
|
||||
}, []);
|
||||
|
||||
// Handle context changes
|
||||
useEffect(() => {
|
||||
contextWs.addMessageHandler((data) => {
|
||||
setContextDirty(false);
|
||||
setModifiedContext(data.context);
|
||||
if (data.generated) {
|
||||
setGeneratedContext(data.context);
|
||||
}
|
||||
});
|
||||
}, [addMessageHandler]);
|
||||
}, []);
|
||||
|
||||
const handleAutoApproverConfigChange = (config) => {
|
||||
sendMessage({
|
||||
type: ClientMessageType.AUTO_APPROVER_CONFIG,
|
||||
config: config
|
||||
// Handle incoming tokens
|
||||
useEffect(() => {
|
||||
tokenWs.addMessageHandler((data) => {
|
||||
setGeneratedResponses(prev => ({
|
||||
...prev,
|
||||
[data.llm]: (prev[data.llm] || '') + data.token
|
||||
}));
|
||||
setModifiedResponses(prev => ({
|
||||
...prev,
|
||||
[data.llm]: (prev[data.llm] || '') + data.token
|
||||
}));
|
||||
});
|
||||
}, []);
|
||||
|
||||
// Handle stdout changes
|
||||
useEffect(() => {
|
||||
stdoutWs.addMessageHandler((data) => {
|
||||
setOutput(data.output);
|
||||
});
|
||||
}, []);
|
||||
|
||||
const handleNextLlm = () => {
|
||||
const llmNames = Object.keys(llms);
|
||||
const currentIndex = llmNames.indexOf(activeLlm);
|
||||
const nextIndex = (currentIndex + 1) % llmNames.length;
|
||||
setActiveLlm(llmNames[nextIndex]);
|
||||
};
|
||||
|
||||
const handleApprove = () => {
|
||||
if (agentState === AgentState.CONTEXT_APPROVAL) {
|
||||
sendMessage({
|
||||
"type": ClientMessageType.APPROVE_CONTEXT,
|
||||
"context": modifiedContext
|
||||
const state = llms[activeLlm];
|
||||
if (state === LlmState.NO_OUTPUT) {
|
||||
if (contextDirty) {
|
||||
fetch('/api/context', {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ context: modifiedContext })
|
||||
});
|
||||
}
|
||||
fetch(`/api/inference/${activeLlm}`, {
|
||||
method: 'POST',
|
||||
});
|
||||
} else if (agentState === AgentState.RESPONSE_APPROVAL) {
|
||||
sendMessage({
|
||||
"type": ClientMessageType.APPROVE_RESPONSE,
|
||||
"response": modifiedResponse
|
||||
} else if (state === LlmState.OUTPUT) {
|
||||
fetch(`/api/approve/${activeLlm}`, {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ response: modifiedResponses[activeLlm] })
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
const handleSendInput = () => {
|
||||
if (input.trim()) {
|
||||
sendMessage({
|
||||
"type": ClientMessageType.SEND_INPUT,
|
||||
"input": input
|
||||
fetch('/api/input', {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ input: input })
|
||||
});
|
||||
setInput('');
|
||||
}
|
||||
};
|
||||
|
||||
const handleClearOutput = () => {
|
||||
sendMessage({
|
||||
"type": ClientMessageType.CLEAR_OUTPUT
|
||||
});
|
||||
fetch('/api/clear', { method: 'POST' });
|
||||
};
|
||||
|
||||
useEffect(() => {
|
||||
if (agentState === AgentState.APPROVE_RESPONSE) {
|
||||
sendMessage({
|
||||
"type": ClientMessageType.MODIFY_RESPONSE,
|
||||
"response": modifiedResponse
|
||||
});
|
||||
}
|
||||
}, [modifiedResponse]);
|
||||
const handleContextEdit = (context) => {
|
||||
setContextDirty(true);
|
||||
setGeneratedResponses({});
|
||||
setModifiedResponses({});
|
||||
const resetLlms = {};
|
||||
for (const llm of Object.keys(llms)) {
|
||||
resetLlms[llm] = LlmState.NO_OUTPUT;
|
||||
} setLlms(resetLlms);
|
||||
setModifiedContext(context);
|
||||
}
|
||||
|
||||
useEffect(() => {
|
||||
switch (agentState) {
|
||||
case AgentState.UPDATE: // fall through
|
||||
case AgentState.CONTEXT_APPROVAL:
|
||||
const state = llms[activeLlm];
|
||||
switch (state) {
|
||||
case LlmState.NO_OUTPUT:
|
||||
setActiveTab(Tabs.CONTEXT);
|
||||
setShowDiff(false);
|
||||
setShowSidebar(false);
|
||||
break;
|
||||
case AgentState.INFERENCE: // fall through
|
||||
case AgentState.RESPONSE_APPROVAL:
|
||||
default:
|
||||
setActiveTab(Tabs.RESPONSE);
|
||||
setShowDiff(false);
|
||||
setShowSidebar(false);
|
||||
break;
|
||||
}
|
||||
setShowDiff(false);
|
||||
setShowSidebar(false);
|
||||
}, [agentState]);
|
||||
}, [llms[activeLlm]]);
|
||||
|
||||
useEffect(() => {
|
||||
if (llmWs.wsState === WebSocketState.CONNECTED && contextWs.wsState === WebSocketState.CONNECTED && tokenWs.wsState === WebSocketState.CONNECTED && stdoutWs.wsState === WebSocketState.CONNECTED) {
|
||||
setWsState(WebSocketState.CONNECTED);
|
||||
} else if (llmWs.wsState === WebSocketState.CONNECTING || contextWs.wsState === WebSocketState.CONNECTING || tokenWs.wsState === WebSocketState.CONNECTING || stdoutWs.wsState === WebSocketState.CONNECTING) {
|
||||
setWsState(WebSocketState.CONNECTING);
|
||||
} else {
|
||||
setWsState(WebSocketState.DISCONNECTED);
|
||||
}
|
||||
}, [llmWs.wsState, contextWs.wsState, tokenWs.wsState, stdoutWs.wsState]);
|
||||
|
||||
return (
|
||||
<div className="h-screen bg-gray-100 flex flex-col overflow-hidden">
|
||||
<Header
|
||||
wsState={wsState}
|
||||
agentState={agentState}
|
||||
agentState={llms[activeLlm]}
|
||||
activeTab={activeTab}
|
||||
setActiveTab={setActiveTab}
|
||||
/>
|
||||
@@ -131,19 +191,22 @@ const App = () => {
|
||||
{activeTab === Tabs.CONTEXT && (
|
||||
<ContentEditor
|
||||
showDiff={showDiff}
|
||||
originalContent={originalContext}
|
||||
originalContent={generatedContext}
|
||||
modifiedContent={modifiedContext}
|
||||
onChange={value => setModifiedContext(value)}
|
||||
onChange={handleContextEdit}
|
||||
validationError={null}
|
||||
/>
|
||||
)}
|
||||
{activeTab === Tabs.RESPONSE && (
|
||||
<ContentEditor
|
||||
showDiff={showDiff}
|
||||
originalContent={originalResponse}
|
||||
modifiedContent={modifiedResponse}
|
||||
onChange={value => setModifiedResponse(value)}
|
||||
validationError={validationError}
|
||||
originalContent={generatedResponses[activeLlm] || ''}
|
||||
modifiedContent={modifiedResponses[activeLlm] || ''}
|
||||
onChange={value => setModifiedResponses(prev => ({
|
||||
...prev,
|
||||
[activeLlm]: value
|
||||
}))}
|
||||
validationError={null}
|
||||
/>
|
||||
)}
|
||||
{activeTab === Tabs.INPUT && (
|
||||
@@ -164,17 +227,18 @@ const App = () => {
|
||||
|
||||
<Sidebar
|
||||
showSidebar={showSidebar}
|
||||
activeTab={activeTab}
|
||||
agentState={agentState}
|
||||
validationError={validationError}
|
||||
llms={llms}
|
||||
activeLlm={activeLlm}
|
||||
llmState={llms[activeLlm]}
|
||||
showDiff={showDiff}
|
||||
input={input}
|
||||
output={output}
|
||||
onApprove={handleApprove}
|
||||
onToggleDiff={() => setShowDiff(!showDiff)}
|
||||
onSendInput={handleSendInput}
|
||||
onClearOutput={handleClearOutput}
|
||||
autoApproverConfig={autoApproverConfig}
|
||||
onAutoApproverConfigChange={handleAutoApproverConfigChange}
|
||||
onLlmChange={setActiveLlm}
|
||||
onNextLlm={handleNextLlm}
|
||||
/>
|
||||
|
||||
<FloatingButton onClick={() => setShowSidebar(!showSidebar)} />
|
||||
@@ -183,4 +247,4 @@ const App = () => {
|
||||
);
|
||||
};
|
||||
|
||||
export default App;
|
||||
export default App;
|
||||
|
||||
Reference in New Issue
Block a user