From dcb4d35d07bf44c9f08cd8a21f0003d0c4d951f7 Mon Sep 17 00:00:00 2001
From: James Ketrenos <james_git@ketrenos.com>
Date: Mon, 31 Mar 2025 14:42:16 -0700
Subject: [PATCH] Making UI progress; seeding lkml RAG

---
 src/ketr-chat/src/App.tsx | 442 +++++++++++++++++++-------------------
 src/server.py             |  21 +-
 2 files changed, 234 insertions(+), 229 deletions(-)

diff --git a/src/ketr-chat/src/App.tsx b/src/ketr-chat/src/App.tsx
index bece700..9566b49 100644
--- a/src/ketr-chat/src/App.tsx
+++ b/src/ketr-chat/src/App.tsx
@@ -1,4 +1,4 @@
-import React, { useState, useEffect, useRef } from 'react';
+import React, { useState, useEffect, useRef, useCallback } from 'react';
 import FormGroup from '@mui/material/FormGroup';
 import FormControlLabel from '@mui/material/FormControlLabel';
 import Switch from '@mui/material/Switch';
@@ -62,228 +62,38 @@ type Tool = {
 type SeverityType = 'error' | 'info' | 'success' | 'warning' | undefined;
 
 interface ControlsParams {
-  sessionId: string,
-  connectionBase: string,
-  setSnack: (snackMessage: string, snackSeverity?: SeverityType) => void
-  onClearHistory: () => void
+  tools: Tool[],
+  rags: Tool[],
+  systemPrompt: string,
+  toggleTool: (tool: Tool) => void,
+  toggleRag: (tool: Tool) => void,
+  setRags: (rags: Tool[]) => void,
+  setSystemPrompt: (prompt: string) => void,
+  reset: (types: ("rags" | "tools" | "history" | "system-prompt")[]) => Promise<void>
 };
 
-const Controls = ({ sessionId, connectionBase, setSnack, onClearHistory }: ControlsParams) => {
-  const [systemPrompt, setSystemPrompt] = useState<string>("");
+const Controls = ({ tools, rags, systemPrompt, toggleTool, toggleRag, setSystemPrompt, reset }: ControlsParams) => {
   const [editSystemPrompt, setEditSystemPrompt] = useState<string>(systemPrompt);
-  const [tools, setTools] = useState<Tool[]>([]);
-  const [rags, setRags] = useState<Tool[]>([]);
 
   useEffect(() => {
-    if (systemPrompt !== "") {
-      return;
-    }
-    const fetchSystemPrompt = async () => {
-      // Make the fetch request with proper headers
-      const response = await fetch(connectionBase + `/api/system-prompt/${sessionId}`, {
-        method: 'GET',
-        headers: {
-          'Content-Type': 'application/json',
-          'Accept': 'application/json',
-        },
-      });
-      const data = await response.json();
-      const systemPrompt = data["system-prompt"].trim();
-      setSystemPrompt(systemPrompt);
-      setEditSystemPrompt(systemPrompt);
-    }
+    setEditSystemPrompt(systemPrompt);
+  }, [systemPrompt, setEditSystemPrompt]);
 
-    fetchSystemPrompt();
-  }, [sessionId, systemPrompt, setSystemPrompt, setEditSystemPrompt, connectionBase]);
-
-  useEffect(() => {
-    if (tools.length) {
-      return;
-    }
-    const fetchTools = async () => {
-      try {
-        // Make the fetch request with proper headers
-        const response = await fetch(connectionBase + `/api/tools/${sessionId}`, {
-          method: 'GET',
-          headers: {
-            'Content-Type': 'application/json',
-            'Accept': 'application/json',
-          },
-        });
-        if (!response.ok) {
-          throw Error();
-        }
-        const tools = await response.json();
-        setTools(tools);
-      } catch (error: any) {
-        setSnack("Unable to fetch tools", "error");
-        console.error(error);
-      }
-    }
-
-    fetchTools();
-  }, [sessionId, tools, setTools, setSnack, connectionBase]);
-
-  useEffect(() => {
-    if (rags.length) {
-      return;
-    }
-    const fetchRags = async () => {
-      try {
-        // Make the fetch request with proper headers
-        const response = await fetch(connectionBase + `/api/rags/${sessionId}`, {
-          method: 'GET',
-          headers: {
-            'Content-Type': 'application/json',
-            'Accept': 'application/json',
-          },
-        });
-        if (!response.ok) {
-          throw Error();
-        }
-        const rags = await response.json();
-        setRags(rags);
-      } catch (error: any) {
-        setSnack("Unable to fetch RAGs", "error");
-        console.error(error);
-      }
-    }
-
-    fetchRags();
-  }, [sessionId, rags, setRags, setSnack, connectionBase]);
-
-  const toggleTool = async (type: string, index: number) => {
+  const toggle = async (type: string, index: number) => {
     switch (type) {
       case "rag":
-        setSnack("RAG backend not yet implemented", "warning");
-        // rags[index].enabled = !rags[index].enabled
-        // setRags([...rags])
+        toggleRag(rags[index])
         break;
       case "tool":
-        const tool = tools[index];
-        tool.enabled = !tool.enabled
-        try {
-          const response = await fetch(connectionBase + `/api/tools/${sessionId}`, {
-            method: 'PUT',
-            headers: {
-              'Content-Type': 'application/json',
-              'Accept': 'application/json',
-            },
-            body: JSON.stringify({ "tool": tool?.function?.name, "enabled": tool.enabled }),
-          });
-
-          const tools = await response.json();
-          setTools([...tools])
-          setSnack(`${tool?.function?.name} ${tool.enabled ? "enabled" : "disabled"}`);
-        } catch (error) {
-          console.error('Fetch error:', error);
-          setSnack(`${tool?.function?.name} ${tool.enabled ? "enabling" : "disabling"} failed.`, "error");
-          tool.enabled = !tool.enabled
-        }
-    };
-  };
-
-  const sendSystemPrompt = async () => {
-    if (!editSystemPrompt.trim()) {
-      setEditSystemPrompt(systemPrompt)
-      return;
-    }
-
-    try {
-      const response = await fetch(connectionBase + `/api/system-prompt/${sessionId}`, {
-        method: 'PUT',
-        headers: {
-          'Content-Type': 'application/json',
-          'Accept': 'application/json',
-        },
-        body: JSON.stringify({ "system-prompt": editSystemPrompt.trim() }),
-      });
-
-      const data = await response.json();
-      if (data["system-prompt"] !== systemPrompt) {
-        setSystemPrompt(data["system-prompt"].trim());
-        setSnack("System prompt updated", "success");
-      }
-    } catch (error) {
-      console.error('Fetch error:', error);
-      setSnack("System prompt update failed", "error");
+        toggleTool(tools[index]);
     }
   };
 
-  const onResetSystemPrompt = async () => {
-    try {
-      const response = await fetch(connectionBase + `/api/system-prompt/${sessionId}`, {
-        method: 'PUT',
-        headers: {
-          'Content-Type': 'application/json',
-          'Accept': 'application/json',
-        },
-        body: JSON.stringify({ "system-prompt": "" }),
-      });
-
-      const data = await response.json();
-      const systemPrompt = data["system-prompt"].trim();
-      setSystemPrompt(systemPrompt);
-      setEditSystemPrompt(systemPrompt)
-    } catch (error) {
-      console.error('Fetch error:', error);
-      setSnack("Unable to fetch system prompt", "error");
-    }
-  };
-
-  const onResetToDefaults = async (event: any) => {
-    try {
-      const response = await fetch(connectionBase + `/api/reset/${sessionId}`, {
-        method: 'PUT',
-        headers: {
-          'Content-Type': 'application/json',
-          'Accept': 'application/json',
-        },
-        body: JSON.stringify({ "reset": ["rag", "tools"] }),
-      });
-
-      if (response.ok) {
-        await response.json();
-        setSnack("Defaults restored", "success");
-      } else {
-        throw Error(`${{ status: response.status, message: response.statusText }}`);
-      }
-    } catch (error) {
-      console.error('Fetch error:', error);
-      setSnack("Unable to restore defaults", "error");
-    }
-  };
-
-  const onClearContext = async (event: any) => {
-    try {
-      const response = await fetch(connectionBase + `/api/reset/${sessionId}`, {
-        method: 'PUT',
-        headers: {
-          'Content-Type': 'application/json',
-          'Accept': 'application/json',
-        },
-        body: JSON.stringify({ "reset": ["history"] }),
-      });
-
-      if (response.ok) {
-        await response.json();
-        setSnack("Chat history cleared", "success");
-        onClearHistory();
-      } else {
-        throw Error(`${{ status: response.status, message: response.statusText }}`);
-      }
-    } catch (error) {
-      console.error('Fetch error:', error);
-      setSnack("Unable to clear chat history", "error");
-    }
-  };
-
-
   const handleKeyPress = (event: any) => {
     if (event.key === 'Enter' && event.ctrlKey) {
       switch (event.target.id) {
         case 'SystemPromptInput':
-          sendSystemPrompt();
+          setSystemPrompt(editSystemPrompt);
           break;
       }
     }
@@ -314,8 +124,8 @@ const Controls = ({ sessionId, connectionBase, setSnack, onClearHistory }: Contr
           id="SystemPromptInput"
         />
         <div style={{ display: "flex", flexDirection: "row", gap: "8px", paddingTop: "8px" }}>
-          <Button variant="contained" disabled={editSystemPrompt === systemPrompt} onClick={sendSystemPrompt}>Set</Button>
-          <Button variant="outlined" onClick={onResetSystemPrompt} color="error">Reset</Button>
+          <Button variant="contained" disabled={editSystemPrompt === systemPrompt} onClick={() => { setSystemPrompt(editSystemPrompt); }}>Set</Button>
+          <Button variant="outlined" onClick={() => { reset(["system-prompt"]); }} color="error">Reset</Button>
         </div>
       </AccordionActions>
     </Accordion>
@@ -334,7 +144,7 @@ const Controls = ({ sessionId, connectionBase, setSnack, onClearHistory }: Contr
             tools.map((tool, index) =>
               <Box key={index}>
                 <Divider />
-                <FormControlLabel control={<Switch checked={tool.enabled} />} onChange={() => toggleTool("tool", index)} label={tool?.function?.name} />
+                <FormControlLabel control={<Switch checked={tool.enabled} />} onChange={() => toggle("tool", index)} label={tool?.function?.name} />
                 <Typography>{tool?.function?.description}</Typography>
               </Box>
             )
@@ -356,15 +166,15 @@ const Controls = ({ sessionId, connectionBase, setSnack, onClearHistory }: Contr
             rags.map((rag, index) =>
               <Box key={index}>
                 <Divider />
-                <FormControlLabel control={<Switch checked={rag.enabled} />} onChange={() => toggleTool("rag", index)} label={rag?.name} />
+                <FormControlLabel control={<Switch checked={rag.enabled} />} onChange={() => toggle("rag", index)} label={rag?.name} />
                 <Typography>{rag?.description}</Typography>
               </Box>
             )
           }</FormGroup>
       </AccordionActions>
     </Accordion>
-    <Button onClick={onClearContext}>Clear Chat History</Button>
-    <Button onClick={onResetToDefaults}>Reset to defaults</Button>
+    <Button onClick={() => { reset(["history"]); }}>Clear Chat History</Button>
+    <Button onClick={() => { reset(["rags", "tools", "system-prompt"]) }}>Reset to defaults</Button>
   </div>);
 }
 
@@ -380,6 +190,10 @@ const App = () => {
   const [snackOpen, setSnackOpen] = useState(false);
   const [snackMessage, setSnackMessage] = useState("");
   const [snackSeverity, setSnackSeverity] = useState<SeverityType>("success");
+  const [tools, setTools] = useState<Tool[]>([]);
+  const [rags, setRags] = useState<Tool[]>([]);
+  const [systemPrompt, setSystemPrompt] = useState<string>("");
+  const [serverSystemPrompt, setServerSystemPrompt] = useState<string>("");
 
   // Scroll to bottom of conversation when conversation updates
   useEffect(() => {
@@ -388,6 +202,7 @@ const App = () => {
     }
   }, [conversation]);
 
+  // Set the initial chat history to "loading" or the welcome message if loaded.
   useEffect(() => {
     if (sessionId === undefined) {
       setConversation([loadingMessage]);
@@ -396,6 +211,8 @@ const App = () => {
     }
   }, [sessionId, setConversation]);
 
+  // Extract the sessionId from the URL if present, otherwise
+  // request a sessionId from the server.
   useEffect(() => {
     const url = new URL(loc.href);
     const pathParts = url.pathname.split('/').filter(Boolean);
@@ -419,15 +236,199 @@ const App = () => {
       console.log(`Session id: ${pathParts[0]} -- existing session`)
       setSessionId(pathParts[0]);
     }
-    
+
   }, [setSessionId, loc]);
 
-  const setSnack = (message: string, severity: SeverityType = "success") => {
+  // If the systemPrompt has not been set, fetch it from the server
+  useEffect(() => {
+    if (serverSystemPrompt !== "" || sessionId === undefined) {
+      return;
+    }
+    const fetchSystemPrompt = async () => {
+      // Make the fetch request with proper headers
+      const response = await fetch(getConnectionBase(loc) + `/api/system-prompt/${sessionId}`, {
+        method: 'GET',
+        headers: {
+          'Content-Type': 'application/json',
+          'Accept': 'application/json',
+        },
+      });
+      const data = await response.json();
+      const serverSystemPrompt = data["system-prompt"].trim();
+      console.log("System prompt initialized to:", serverSystemPrompt);
+      setServerSystemPrompt(serverSystemPrompt);
+      setSystemPrompt(serverSystemPrompt);
+    }
+
+    fetchSystemPrompt();
+  }, [sessionId, serverSystemPrompt, setServerSystemPrompt, loc]);
+
+  // Set the snack pop-up and open it
+  const setSnack = useCallback((message: string, severity: SeverityType = "success") => {
     setSnackMessage(message);
     setSnackSeverity(severity);
     setSnackOpen(true);
+  }, []);
+
+  // If the tools have not been set, fetch them from the server
+  useEffect(() => {
+    if (tools.length || sessionId === undefined) {
+      return;
+    }
+    const fetchTools = async () => {
+      try {
+        // Make the fetch request with proper headers
+        const response = await fetch(getConnectionBase(loc) + `/api/tools/${sessionId}`, {
+          method: 'GET',
+          headers: {
+            'Content-Type': 'application/json',
+            'Accept': 'application/json',
+          },
+        });
+        if (!response.ok) {
+          throw Error();
+        }
+        const tools = await response.json();
+        setTools(tools);
+      } catch (error: any) {
+        setSnack("Unable to fetch tools", "error");
+        console.error(error);
+      }
+    }
+
+    fetchTools();
+  }, [sessionId, tools, setTools, setSnack, loc]);
+
+  // If the RAGs have not been set, fetch them from the server
+  useEffect(() => {
+    if (rags.length || sessionId === undefined) {
+      return;
+    }
+    const fetchRags = async () => {
+      try {
+        // Make the fetch request with proper headers
+        const response = await fetch(getConnectionBase(loc) + `/api/rags/${sessionId}`, {
+          method: 'GET',
+          headers: {
+            'Content-Type': 'application/json',
+            'Accept': 'application/json',
+          },
+        });
+        if (!response.ok) {
+          throw Error();
+        }
+        const rags = await response.json();
+        setRags(rags);
+      } catch (error: any) {
+        setSnack("Unable to fetch RAGs", "error");
+        console.error(error);
+      }
+    }
+
+    fetchRags();
+  }, [sessionId, rags, setRags, setSnack, loc]);
+
+  const toggleRag = async (tool: Tool) => {
+    setSnack("RAG is not yet implemented", "warning");
   }
 
+  const toggleTool = async (tool: Tool) => {
+    tool.enabled = !tool.enabled
+    try {
+      const response = await fetch(getConnectionBase(loc) + `/api/tools/${sessionId}`, {
+        method: 'PUT',
+        headers: {
+          'Content-Type': 'application/json',
+          'Accept': 'application/json',
+        },
+        body: JSON.stringify({ "tool": tool?.function?.name, "enabled": tool.enabled }),
+      });
+
+      const tools = await response.json();
+      setTools([...tools])
+      setSnack(`${tool?.function?.name} ${tool.enabled ? "enabled" : "disabled"}`);
+    } catch (error) {
+      console.error('Fetch error:', error);
+      setSnack(`${tool?.function?.name} ${tool.enabled ? "enabling" : "disabling"} failed.`, "error");
+      tool.enabled = !tool.enabled
+    }
+  };
+
+  useEffect(() => {
+    if (systemPrompt === serverSystemPrompt || !systemPrompt.trim() || sessionId === undefined) {
+      return;
+    }
+    const sendSystemPrompt = async (prompt: string) => {
+      try {
+        const response = await fetch(getConnectionBase(loc) + `/api/system-prompt/${sessionId}`, {
+          method: 'PUT',
+          headers: {
+            'Content-Type': 'application/json',
+            'Accept': 'application/json',
+          },
+          body: JSON.stringify({ "system-prompt": prompt }),
+        });
+
+        const data = await response.json();
+        const newPrompt = data["system-prompt"];
+        if (newPrompt !== serverSystemPrompt) {
+          setServerSystemPrompt(newPrompt);
+          setSystemPrompt(newPrompt)
+          setSnack("System prompt updated", "success");
+        }
+      } catch (error) {
+        console.error('Fetch error:', error);
+        setSnack("System prompt update failed", "error");
+      }
+    };
+
+    sendSystemPrompt(systemPrompt);
+
+  }, [systemPrompt, setServerSystemPrompt, serverSystemPrompt, loc, sessionId, setSnack]);
+
+  const reset = async (types: ("rags" | "tools" | "history" | "system-prompt")[]) => {
+    try {
+      const response = await fetch(getConnectionBase(loc) + `/api/reset/${sessionId}`, {
+        method: 'PUT',
+        headers: {
+          'Content-Type': 'application/json',
+          'Accept': 'application/json',
+        },
+        body: JSON.stringify({ "reset": types }),
+      });
+
+      if (response.ok) {
+        const data = await response.json();
+        if (data.error) {
+          throw Error()
+        }
+        for (const [key, value] of Object.entries(data)) {
+          switch (key) {
+            case "rags":
+              setRags(value as Tool[]);
+              break;
+            case "tools":
+              setTools(value as Tool[]);
+              break;
+            case "system-prompt":
+              setServerSystemPrompt((value as any)["system-prompt"].trim());
+              setSystemPrompt((value as any)["system-prompt"].trim());
+              break;
+            case "history":
+              setConversation([welcomeMessage]);
+              break;
+          }
+        }
+        setSnack("Update successful.", "success");
+      } else {
+        throw Error(`${{ status: response.status, message: response.statusText }}`);
+      }
+    } catch (error) {
+      console.error('Fetch error:', error);
+      setSnack("Unable to restore defaults", "error");
+    }
+  };
+
   const handleDrawerClose = () => {
     setIsClosing(true);
     setMobileOpen(false);
@@ -443,13 +444,9 @@ const App = () => {
     }
   };
 
-  const onClearHistory = () => {
-    setConversation([welcomeMessage]);
-  };
-
   const drawer = (
     <>
-      {sessionId !== undefined && <Controls onClearHistory={onClearHistory} setSnack={setSnack} sessionId={sessionId} connectionBase={getConnectionBase(loc)} />}
+      {sessionId !== undefined && <Controls {...{ tools, rags, reset, systemPrompt, toggleTool, toggleRag, setRags, setSystemPrompt }} />}
     </>
   );
 
@@ -479,6 +476,10 @@ const App = () => {
 
   type MessageList = Message[];
 
+  const onNew = async () => {
+    reset(["rags", "tools", "history", "system-prompt"]);
+  }
+
   const sendQuery = async () => {
     if (!query.trim()) return;
 
@@ -743,10 +744,11 @@ const App = () => {
           />
           <AccordionActions>
             <Button variant="contained" onClick={sendQuery}>Send</Button>
+            <Button variant="outlined" onClick={onNew}>New</Button>
           </AccordionActions>
         </Box>
       </Box>
-      <Snackbar open={snackOpen} autoHideDuration={snackSeverity === "success" ? 1500 : 6000} onClose={handleSnackClose}>
+      <Snackbar open={snackOpen} autoHideDuration={(snackSeverity === "success" || snackSeverity === "info") ? 1500 : 6000} onClose={handleSnackClose}>
         <Alert
           onClose={handleSnackClose}
           severity={snackSeverity}
diff --git a/src/server.py b/src/server.py
index 101d5ae..3c34b9d 100644
--- a/src/server.py
+++ b/src/server.py
@@ -87,7 +87,7 @@ When you receive a response from summarize_site, you must:
 1. Review the entire content returned by the second LLM
 2. Provide the URL used to obtain the information.
 3. Incorporate the information into your response as appropriate
-"""
+""".strip()
 
 tool_log = []
 command_log = []
@@ -376,25 +376,28 @@ class WebServer:
             context = self.upsert_context(context_id)
             data = await request.json()
             try:
+                response = {}
                 for reset in data["reset"]:
                     match reset:
                         case "system-prompt":
                             context["system"] = [{"role": "system", "content": system_message}]
-                            return JSONResponse(context["system"])
+                            response["system-prompt"] = { "system-prompt": system_message }
                         case "rag":
                             context["rag"] = rags.copy()
-                            return JSONResponse(context["rag"])
+                            response["rags"] = context["rag"]
                         case "tools":
                             context["tools"] = default_tools(tools)
-                            return JSONResponse(context["tools"])
+                            response["tools"] = context["tools"]
                         case "history":
                             context["history"] = []
-                            return JSONResponse(context["history"])
-
-                return JSONResponse({ "error": "Usage: { reset: rag|tools|history|system-prompt}"}), 405
+                            response["history"] = context["history"]
+                if not response:
+                    return JSONResponse({ "error": "Usage: { reset: rag|tools|history|system-prompt}"})
+                else:
+                    return JSONResponse(response);
         
             except:
-                return JSONResponse({ "error": "Usage: { reset: rag|tools|history|system-prompt}"}), 405
+                return JSONResponse({ "error": "Usage: { reset: rag|tools|history|system-prompt}"})
         
 
         @self.app.put('/api/system-prompt/{context_id}')
@@ -406,7 +409,7 @@ class WebServer:
             data = await request.json()
             system_prompt = data["system-prompt"].strip()
             if not system_prompt:
-                return JSONResponse({ "status": "error", "message": "System prompt can not be empty." }), 405
+                return JSONResponse({ "status": "error", "message": "System prompt can not be empty." })
             context["system"] =  [{"role": "system", "content": system_prompt}]
             return JSONResponse({ "system-prompt": system_prompt })