Enable multiple llms

2024-11-22 15:05:54 +01:00
parent bbb88f39e8
commit 8766a945c0
28 changed files with 972 additions and 525 deletions
--- a/web/package.json
+++ b/web/package.json
@@ -14,6 +14,7 @@
    "@monaco-editor/react": "^4.6.0",
    "@radix-ui/react-alert-dialog": "^1.0.5",
    "@radix-ui/react-scroll-area": "^1.0.5",
+    "@radix-ui/react-select": "^2.0.0",
    "@radix-ui/react-tabs": "^1.0.4",
    "class-variance-authority": "^0.7.0",
    "clsx": "^2.1.0",
@@ -39,4 +40,4 @@
    "tailwindcss": "^3.4.1",
    "vite": "^5.1.0"
  }
-}
+}
--- a/web/src/components/App.jsx
+++ b/web/src/components/App.jsx
@@ -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;
--- a/web/src/components/Header.jsx
+++ b/web/src/components/Header.jsx
@@ -1,5 +1,12 @@
 import React from 'react';
-import { WebSocketState, AgentState, Tabs } from '../constants';
+import { WebSocketState } from '../hooks/useWebSocket';
+
+export const Tabs = {
+  CONTEXT: 'CONTEXT',
+  RESPONSE: 'RESPONSE',
+  INPUT: 'INPUT',
+  OUTPUT: 'OUTPUT'
+};

 export const Header = ({ 
  wsState, 
--- a/web/src/components/Sidebar.jsx
+++ b/web/src/components/Sidebar.jsx
@@ -1,72 +1,51 @@
 import React from 'react';
 import { Button } from '@/components/ui/button';
-import { Switch } from '@/components/ui/switch';
-import { Input } from '@/components/ui/input';
-import { AgentState, Tabs } from '../constants';
+import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from '@/components/ui/select';
+import { LlmState } from '@/components/App';

 export const Sidebar = ({
  showSidebar,
-  agentState,
+  llms,
+  activeLlm,
+  llmState,
  showDiff,
  input,
-  autoApproverConfig,
+  output,
  onApprove,
  onToggleDiff,
  onSendInput,
  onClearOutput,
-  onAutoApproverConfigChange,
+  onLlmChange,
+  onNextLlm,
 }) => (
  <aside className={`
    fixed top-0 right-0 h-screen w-64 bg-white shadow-lg 
    transition-transform duration-200 ease-in-out z-50
-    md:sticky md:h-[calc(100vh-4rem)] md:top-16
+    md:sticky md:h-[calc(100vh-4rem)]
    ${showSidebar ? 'translate-x-0' : 'translate-x-full md:translate-x-0'}
  `}>
-    <div className="p-4 space-y-4 md:pt-0">
-      <div className="space-y-2">
-        <div className="flex items-center justify-between">
-          <label>Auto-approve Context</label>
-          <Switch 
-            checked={autoApproverConfig.context_enabled}
-            onCheckedChange={(checked) => onAutoApproverConfigChange({
-              ...autoApproverConfig,
-              context_enabled: checked
-            })}
-          />
-        </div>
-        <Input
-          type="number"
-          value={autoApproverConfig.context_timeout}
-          onChange={(e) => onAutoApproverConfigChange({
-            ...autoApproverConfig,
-            context_timeout: parseFloat(e.target.value)
-          })}
-          disabled={autoApproverConfig.context_enabled}
-        />
-      </div>
-
-      <div className="space-y-2">
-        <div className="flex items-center justify-between">
-          <label>Auto-approve Response</label>
-          <Switch
-            checked={autoApproverConfig.response_enabled}
-            onCheckedChange={(checked) => onAutoApproverConfigChange({
-              ...autoApproverConfig,
-              response_enabled: checked
-            })}
-          />
-        </div>
-        <Input
-          type="number"
-          value={autoApproverConfig.response_timeout}
-          onChange={(e) => onAutoApproverConfigChange({
-            ...autoApproverConfig,
-            response_timeout: parseFloat(e.target.value)
-          })}
-          disabled={autoApproverConfig.response_enabled}
-        />
-      </div>
+    <div className="p-4 space-y-4 overflow-y-auto h-full">
+      <Select value={activeLlm} onValueChange={onLlmChange}>
+        <SelectTrigger>
+          <SelectValue placeholder="Select LLM" />
+        </SelectTrigger>
+        <SelectContent>
+          {Object.keys(llms).map(llm => (
+            <SelectItem key={llm} value={llm}>
+              {llm} ({llms[llm]})
+            </SelectItem>
+          ))}
+        </SelectContent>
+      </Select>

+      <Button
+        variant="outline"
+        onClick={onNextLlm}
+        className="w-full"
+      >
+        Next LLM
+      </Button>
+      
      <Button
        variant="outline"
        onClick={onToggleDiff}
@@ -77,10 +56,14 @@ export const Sidebar = ({

      <Button 
        onClick={onApprove}
-        disabled={(agentState !== AgentState.CONTEXT_APPROVAL) && (agentState !== AgentState.RESPONSE_APPROVAL)}
+        disabled={llmState === LlmState.INFERENCE}
        className="w-full"
      >
-        Approve
+        {
+          llmState === LlmState.NO_OUTPUT ? 'Inference' :
+          llmState === LlmState.OUTPUT ? 'Approve' :
+          'In Progress'
+        }
      </Button>

      <Button 
@@ -93,7 +76,7 @@ export const Sidebar = ({

      <Button 
        onClick={onClearOutput}
-        variant="outline"
+        disabled={!output.trim()}
        className="w-full"
      >
        Clear Output
--- a/web/src/components/editors/ContextEditor.jsx
+++ b/web/src/components/editors/ContextEditor.jsx
@@ -1,15 +0,0 @@
-import React from 'react';
-import { BaseEditor } from './BaseEditor';
-
-export const ContextEditor = ({
-  originalContent,
-  modifiedContent,
-  showOriginal,
-  onChange,
-}) => (
-  <BaseEditor
-    content={showOriginal ? originalContent : modifiedContent}
-    onChange={onChange}
-    readOnly={showOriginal}
-  />
-);
--- a/web/src/components/editors/InputEditor.jsx
+++ b/web/src/components/editors/InputEditor.jsx
@@ -1,10 +0,0 @@
-import React from 'react';
-import { BaseEditor } from './BaseEditor';
-
-export const InputEditor = ({ content, onChange }) => (
-  <BaseEditor
-    content={content}
-    onChange={onChange}
-    language="plaintext"
-  />
-);
--- a/web/src/components/editors/OutputEditor.jsx
+++ b/web/src/components/editors/OutputEditor.jsx
@@ -1,10 +0,0 @@
-import React from 'react';
-import { BaseEditor } from './BaseEditor';
-
-export const OutputEditor = ({ content }) => (
-  <BaseEditor
-    content={content}
-    readOnly={true}
-    language="plaintext"
-  />
-);
--- a/web/src/components/editors/ResponseEditor.jsx
+++ b/web/src/components/editors/ResponseEditor.jsx
@@ -1,17 +0,0 @@
-import React from 'react';
-import { BaseEditor } from './BaseEditor';
-
-export const ResponseEditor = ({
-  originalContent,
-  modifiedContent,
-  showOriginal,
-  onChange,
-  validationError,
-}) => (
-  <BaseEditor
-    content={showOriginal ? originalContent : modifiedContent}
-    onChange={onChange}
-    readOnly={showOriginal}
-    validationError={validationError}
-  />
-);
--- a/web/src/components/ui/select.jsx
+++ b/web/src/components/ui/select.jsx
@@ -0,0 +1,52 @@
+import * as React from "react"
+import * as SelectPrimitive from "@radix-ui/react-select"
+import { Check, ChevronDown } from "lucide-react"
+
+const Select = SelectPrimitive.Root
+
+const SelectTrigger = React.forwardRef(({ className, children, ...props }, ref) => (
+  <SelectPrimitive.Trigger
+    ref={ref}
+    className="flex h-10 w-full items-center justify-between rounded-md border border-input bg-background px-3 py-2 text-sm ring-offset-background placeholder:text-muted-foreground focus:outline-none focus:ring-2 focus:ring-ring focus:ring-offset-2 disabled:cursor-not-allowed disabled:opacity-50"
+    {...props}>
+    {children}
+    <SelectPrimitive.Icon asChild>
+      <ChevronDown className="h-4 w-4 opacity-50" />
+    </SelectPrimitive.Icon>
+  </SelectPrimitive.Trigger>
+))
+SelectTrigger.displayName = "SelectTrigger"
+
+const SelectContent = React.forwardRef(({ className, children, position = "popper", ...props }, ref) => (
+  <SelectPrimitive.Portal>
+    <SelectPrimitive.Content
+      ref={ref}
+      className="relative z-50 min-w-[8rem] overflow-hidden rounded-md border bg-popover text-popover-foreground shadow-md data-[state=open]:animate-in data-[state=closed]:animate-out data-[state=closed]:fade-out-0 data-[state=open]:fade-in-0 data-[state=closed]:zoom-out-95 data-[state=open]:zoom-in-95 data-[side=bottom]:slide-in-from-top-2 data-[side=left]:slide-in-from-right-2 data-[side=right]:slide-in-from-left-2 data-[side=top]:slide-in-from-bottom-2"
+      position={position}
+      {...props}>
+      <SelectPrimitive.Viewport className="p-1">
+        {children}
+      </SelectPrimitive.Viewport>
+    </SelectPrimitive.Content>
+  </SelectPrimitive.Portal>
+))
+SelectContent.displayName = "SelectContent"
+
+const SelectItem = React.forwardRef(({ className, children, ...props }, ref) => (
+  <SelectPrimitive.Item
+    ref={ref}
+    className="relative flex w-full cursor-default select-none items-center rounded-sm py-1.5 pl-8 pr-2 text-sm outline-none focus:bg-accent focus:text-accent-foreground data-[disabled]:pointer-events-none data-[disabled]:opacity-50"
+    {...props}>
+    <span className="absolute left-2 flex h-3.5 w-3.5 items-center justify-center">
+      <SelectPrimitive.ItemIndicator>
+        <Check className="h-4 w-4" />
+      </SelectPrimitive.ItemIndicator>
+    </span>
+    <SelectPrimitive.ItemText>{children}</SelectPrimitive.ItemText>
+  </SelectPrimitive.Item>
+))
+SelectItem.displayName = "SelectItem"
+
+const SelectValue = SelectPrimitive.Value
+
+export { Select, SelectContent, SelectItem, SelectTrigger, SelectValue }
--- a/web/src/constants.jsx
+++ b/web/src/constants.jsx
@@ -1,36 +0,0 @@
-export const WebSocketState = {
-  CONNECTING: 'CONNECTING',
-  CONNECTED: 'CONNECTED',
-  DISCONNECTED: 'DISCONNECTED',
-};
-
-export const AgentState = {
-  UPDATE: 'UPDATE',
-  CONTEXT_APPROVAL: 'CONTEXT_APPROVAL',
-  INFERENCE: 'INFERENCE',
-  RESPONSE_APPROVAL: 'RESPONSE_APPROVAL',
-};
-
-export const Tabs = {
-  CONTEXT: 'context',
-  RESPONSE: 'response',
-  INPUT: 'input',
-  OUTPUT: 'output'
-};
-
-export const ClientMessageType = {
-  APPROVE_CONTEXT: 'APPROVE_CONTEXT',
-  MODIFY_RESPONSE: 'MODIFY_RESPONSE',
-  APPROVE_RESPONSE: 'APPROVE_RESPONSE',
-  SEND_INPUT: 'SEND_INPUT',
-  CLEAR_OUTPUT: 'CLEAR_OUTPUT',
-  AUTO_APPROVER_CONFIG: 'AUTO_APPROVER_CONFIG',
-};
-
-export const ServerMessageType = {
-  STATE_CHANGE: 'STATE_CHANGE',
-  CONTEXT_UPDATE: 'CONTEXT_UPDATE',
-  RESPONSE_UPDATE: 'RESPONSE_UPDATE',
-  OUTPUT_UPDATE: 'OUTPUT_UPDATE',
-  AUTO_APPROVER_CONFIG: 'AUTO_APPROVER_CONFIG',
-};
--- a/web/src/hooks/useWebSocket.jsx
+++ b/web/src/hooks/useWebSocket.jsx
@@ -1,5 +1,10 @@
 import { useState, useEffect, useRef, useCallback } from 'react';
-import { WebSocketState, MessageType } from '../constants';
+
+export const WebSocketState = {
+  CONNECTING: 'CONNECTING',
+  CONNECTED: 'CONNECTED',
+  DISCONNECTED: 'DISCONNECTED',
+};

 export const useWebSocket = (url) => {
  const [wsState, setWsState] = useState(WebSocketState.DISCONNECTED);