Merge pull request #57 from smallcloudai/big-switch

Big switch
smallcloudai · Aug 19, 2024 · a48abe5 · a48abe5
2 parents 54befef + 9af5998
commit a48abe5
Show file tree

Hide file tree

Showing 9 changed files with 65 additions and 76 deletions.
diff --git a/src/components/Chat/Chat.tsx b/src/components/Chat/Chat.tsx
@@ -49,9 +49,6 @@ export type ChatProps = {
   onSetSystemPrompt: ChatFormProps["onSetSystemPrompt"];
   selectedSystemPrompt: ChatFormProps["selectedSystemPrompt"];
   requestPreviewFiles: ChatFormProps["requestPreviewFiles"];
-  canUseTools: ChatFormProps["canUseTools"];
-  setUseTools: ChatFormProps["setUseTools"];
-  useTools: ChatFormProps["useTools"];
 };
 
 export const Chat: React.FC<ChatProps> = ({
@@ -76,9 +73,6 @@ export const Chat: React.FC<ChatProps> = ({
   onSetSystemPrompt,
   selectedSystemPrompt,
   requestPreviewFiles,
-  canUseTools,
-  setUseTools,
-  useTools,
 }) => {
   const chatContentRef = useRef<HTMLDivElement>(null);
   const activeFile = useAppSelector(selectActiveFile);
@@ -245,9 +239,6 @@ export const Chat: React.FC<ChatProps> = ({
         onSetSystemPrompt={onSetSystemPrompt}
         selectedSystemPrompt={selectedSystemPrompt}
         requestPreviewFiles={requestPreviewFiles}
-        canUseTools={canUseTools}
-        setUseTools={setUseTools}
-        useTools={useTools}
       />
       <Flex justify="between" pl="1" pr="1" pt="1">
         {messages.length > 0 && (

diff --git a/src/components/ChatForm/ChatControls.tsx b/src/components/ChatForm/ChatControls.tsx
@@ -9,9 +9,10 @@ import { PromptSelect, PromptSelectProps } from "./PromptSelect";
 import { Checkbox } from "../Checkbox";
 import { QuestionMarkCircledIcon } from "@radix-ui/react-icons";
 import { useTourRefs } from "../../features/Tour";
+import { ToolUseSwitch } from "./ToolUseSwitch";
+import { ToolUse, selectToolUse, setToolUse } from "../../features/Chat";
 import { useCanUseTools } from "../../hooks/useCanUseTools";
 import { useAppDispatch, useAppSelector } from "../../app/hooks";
-import { selectUseTools, setUseTools } from "../../features/Chat/chatThread";
 
 type CapsSelectProps = {
   value: string;
@@ -149,9 +150,9 @@ export const ChatControls: React.FC<ChatControlsProps> = ({
   const refs = useTourRefs();
   const canUseTools = useCanUseTools();
   const dispatch = useAppDispatch();
-  const useTools = useAppSelector(selectUseTools);
-  const onSetUseTools = useCallback(
-    (value: boolean | string) => dispatch(setUseTools(!!value)),
+  const toolUse = useAppSelector(selectToolUse);
+  const onSetToolUse = useCallback(
+    (value: ToolUse) => dispatch(setToolUse(value)),
     [dispatch],
   );
 
@@ -163,23 +164,6 @@ export const ChatControls: React.FC<ChatControlsProps> = ({
       direction="column"
       className={classNames(styles.controls)}
     >
-      {canUseTools && (
-        <Flex
-          ref={(x) => refs.setUseTools(x)}
-          style={{ alignSelf: "flex-start" }}
-        >
-          <ChatContolCheckBox
-            name="use_tools"
-            checked={useTools}
-            onCheckChange={onSetUseTools}
-            label="Allow model to use tools"
-            infoText="Turn on when asking about your codebase. When tuned on the model can autonomously call functions to gather the best context."
-            href="https://docs.refact.ai/features/ai-chat/"
-            linkText="documentation"
-          />
-        </Flex>
-      )}
-
       {Object.entries(checkboxes).map(([key, checkbox]) => {
         if (host === "web" && checkbox.name === "file_upload") {
           return null;
@@ -203,6 +187,15 @@ export const ChatControls: React.FC<ChatControlsProps> = ({
           </Flex>
         );
       })}
+      {canUseTools && showControls && (
+        <Flex
+          ref={(x) => refs.setUseTools(x)}
+          style={{ alignSelf: "flex-start" }}
+        >
+          <ToolUseSwitch toolUse={toolUse} setToolUse={onSetToolUse} />
+        </Flex>
+      )}
+
       {showControls && (
         <Flex style={{ alignSelf: "flex-start" }}>
           <CapsSelect {...selectProps} />

diff --git a/src/components/ChatForm/ChatForm.stories.tsx b/src/components/ChatForm/ChatForm.stories.tsx
@@ -174,9 +174,6 @@ const meta: Meta<typeof ChatForm> = {
     prompts: SYSTEM_PROMPTS,
     onSetSystemPrompt: noop,
     // selectedSystemPrompt: null,
-    useTools: true,
-    canUseTools: true,
-    setUseTools: noop,
   },
   decorators: [
     (Children) => {

diff --git a/src/components/ChatForm/ChatForm.test.tsx b/src/components/ChatForm/ChatForm.test.tsx
@@ -55,9 +55,6 @@ const App: React.FC<Partial<ChatFormProps>> = ({ ...props }) => {
     prompts: SYSTEM_PROMPTS,
     onSetSystemPrompt: noop,
     selectedSystemPrompt: {},
-    canUseTools: false,
-    setUseTools: noop,
-    useTools: false,
     ...props,
   };
 

diff --git a/src/components/ChatForm/ChatForm.tsx b/src/components/ChatForm/ChatForm.tsx
@@ -20,7 +20,7 @@ import { ChatControls } from "./ChatControls";
 import { addCheckboxValuesToInput } from "./utils";
 import { usePreviewFileRequest } from "./usePreviewFileRequest";
 import { useAppDispatch, useAppSelector, useConfig } from "../../app/hooks";
-import { type Snippet } from "../../features/Chat/selectedSnippet";
+import type { Snippet } from "../../features/Chat";
 import { getErrorMessage, clearError } from "../../features/Errors/errorsSlice";
 import { useTourRefs } from "../../features/Tour";
 import { useCheckboxes } from "./useCheckBoxes";
@@ -55,9 +55,6 @@ export type ChatFormProps = {
   onSetSystemPrompt: (prompt: SystemPrompts) => void;
   selectedSystemPrompt: SystemPrompts;
   chatId: string;
-  canUseTools: boolean;
-  setUseTools: (value: boolean) => void;
-  useTools: boolean;
 };
 
 export const ChatForm: React.FC<ChatFormProps> = ({

diff --git a/src/components/ChatForm/ToolUseSwitch.tsx b/src/components/ChatForm/ToolUseSwitch.tsx
@@ -0,0 +1,26 @@
+import { Flex, SegmentedControl, Text } from "@radix-ui/themes";
+import { ToolUse } from "../../features/Chat";
+
+type ToolUseSwitchProps = {
+  toolUse: ToolUse;
+  setToolUse: (toolUse: ToolUse) => void;
+};
+
+export const ToolUseSwitch = ({ toolUse, setToolUse }: ToolUseSwitchProps) => {
+  return (
+    <Flex direction="column" gap="1" align="start">
+      <Text size="1">How fast do you want the answer:</Text>
+      <SegmentedControl.Root
+        defaultValue="quick"
+        value={toolUse}
+        onValueChange={(x) => {
+          setToolUse(x as ToolUse);
+        }}
+      >
+        <SegmentedControl.Item value="quick">Quick</SegmentedControl.Item>
+        <SegmentedControl.Item value="explore">Explore</SegmentedControl.Item>
+        <SegmentedControl.Item value="agent">Agent</SegmentedControl.Item>
+      </SegmentedControl.Root>
+    </Flex>
+  );
+};
diff --git a/src/features/Chat/Chat.tsx b/src/features/Chat/Chat.tsx
@@ -1,12 +1,10 @@
-import React, { useMemo } from "react";
-// import { useEventBusForChat } from "../hooks/useEventBusForChat";
+import React from "react";
 import type { Config } from "../Config/configSlice";
-import { CodeChatModel, SystemPrompts } from "../../services/refact";
+import { SystemPrompts } from "../../services/refact";
 import { Chat as ChatComponent } from "../../components/Chat";
 import {
   useGetCapsQuery,
   useGetPromptsQuery,
-  useGetToolsQuery,
   useGetCommandCompletionQuery,
   useGetCommandPreviewQuery,
 } from "../../app/hooks";
@@ -15,10 +13,7 @@ import { useAppDispatch, useAppSelector } from "../../app/hooks";
 import {
   getSelectedSystemPrompt,
   setSystemPrompt,
-  setUseTools,
   selectMessages,
-  selectModel,
-  selectUseTools,
 } from "./chatThread";
 
 export type ChatProps = {
@@ -56,7 +51,6 @@ export const Chat: React.FC<ChatProps> = ({
   // state,
 }) => {
   const capsRequest = useGetCapsQuery();
-  const chatModel = useAppSelector(selectModel);
 
   // TODO: these could be lower in the component tree
   const promptsRequest = useGetPromptsQuery();
@@ -65,12 +59,8 @@ export const Chat: React.FC<ChatProps> = ({
   const onSetSelectedSystemPrompt = (prompt: SystemPrompts) =>
     dispatch(setSystemPrompt(prompt));
 
-  const useTools = useAppSelector(selectUseTools);
-  const onSetUseTools = (value: boolean) => dispatch(setUseTools(value));
   const messages = useAppSelector(selectMessages);
 
-  // TODO: don't make this request if there are no caps
-  const toolsRequest = useGetToolsQuery();
   // const chatRequest = useSendChatRequest();
 
   // commands should be a selector, and calling the hook ?
@@ -119,19 +109,6 @@ export const Chat: React.FC<ChatProps> = ({
   const maybeSendToSideBar =
     host === "vscode" && tabbed ? sendToSideBar : undefined;
 
-  // make this a hook ?
-  const canUseTools = useMemo(() => {
-    if (!capsRequest.data) return false;
-    if (!toolsRequest.data) return false;
-    if (toolsRequest.data.length === 0) return false;
-    const modelName = chatModel || capsRequest.data.code_chat_default_model;
-
-    if (!(modelName in capsRequest.data.code_chat_models)) return false;
-    const model: CodeChatModel = capsRequest.data.code_chat_models[modelName];
-    if ("supports_tools" in model && model.supports_tools) return true;
-    return false;
-  }, [capsRequest.data, toolsRequest.data, chatModel]);
-
   // can be a selector
   const unCalledTools = React.useMemo(() => {
     if (messages.length === 0) return false;
@@ -191,9 +168,6 @@ export const Chat: React.FC<ChatProps> = ({
       onSetSystemPrompt={onSetSelectedSystemPrompt}
       selectedSystemPrompt={selectedSystemPrompt}
       requestPreviewFiles={() => ({})}
-      canUseTools={canUseTools}
-      setUseTools={onSetUseTools}
-      useTools={useTools}
       // openSettings={openSettings}
     />
   );

diff --git a/src/features/Chat/chatThread.ts b/src/features/Chat/chatThread.ts
@@ -36,6 +36,8 @@ export type ChatThread = {
   updatedAt?: string;
 };
 
+export type ToolUse = "quick" | "explore" | "agent";
+
 export type Chat = {
   streaming: boolean;
   thread: ChatThread;
@@ -45,7 +47,7 @@ export type Chat = {
   waiting_for_response: boolean;
   cache: Record<string, ChatThread>;
   system_prompt: SystemPrompts;
-  use_tools: boolean;
+  tool_use: ToolUse;
   send_immediately: boolean;
 };
 
@@ -69,7 +71,7 @@ const createInitialState = (): Chat => {
     waiting_for_response: false,
     cache: {},
     system_prompt: {},
-    use_tools: true,
+    tool_use: "agent",
     send_immediately: false,
   };
 };
@@ -124,11 +126,11 @@ export const clearChatError = createAction<PayloadWIthId>(
 
 export const enableSend = createAction<PayloadWIthId>("chatThread/enableSend");
 
-export const setUseTools = createAction<boolean>("chatThread/setUseTools");
+export const setToolUse = createAction<ToolUse>("chatThread/setToolUse");
 
 export const chatReducer = createReducer(initialState, (builder) => {
-  builder.addCase(setUseTools, (state, action) => {
-    state.use_tools = action.payload;
+  builder.addCase(setToolUse, (state, action) => {
+    state.tool_use = action.payload;
   });
 
   builder.addCase(enableSend, (state, action) => {
@@ -371,7 +373,7 @@ export const selectThread = (state: RootState) => state.chat.thread;
 export const selectChatId = (state: RootState) => state.chat.thread.id;
 export const selectModel = (state: RootState) => state.chat.thread.model;
 export const selectMessages = (state: RootState) => state.chat.thread.messages;
-export const selectUseTools = (state: RootState) => state.chat.use_tools;
+export const selectToolUse = (state: RootState) => state.chat.tool_use;
 export const selectIsWaiting = (state: RootState) =>
   state.chat.waiting_for_response;
 export const selectIsStreaming = (state: RootState) => state.chat.streaming;
@@ -394,6 +396,7 @@ export const useSendChatRequest = () => {
   const currentMessages = useAppSelector(selectMessages);
   const systemPrompt = useAppSelector(getSelectedSystemPrompt);
   const sendImmediately = useAppSelector(selectSendImmediately);
+  const toolUse = useAppSelector(selectToolUse);
 
   const messagesWithSystemPrompt = useMemo(() => {
     const prompts = Object.entries(systemPrompt);
@@ -409,7 +412,17 @@ export const useSendChatRequest = () => {
 
   const sendMessages = useCallback(
     (messages: ChatMessages) => {
-      const tools = toolsRequest.data ?? null;
+      let tools = toolsRequest.data ?? null;
+      if (toolUse === "quick") {
+        tools = [];
+      } else if (toolUse === "explore") {
+        tools = tools?.filter((t) => !t.function.agentic) ?? [];
+      }
+      tools =
+        tools?.map((t) => {
+          const { agentic: _, ...remaining } = t.function;
+          return { ...t, function: { ...remaining } };
+        }) ?? [];
       dispatch(backUpMessages({ id: chatId, messages }));
       dispatch(chatAskedQuestion({ id: chatId }));
 
@@ -422,7 +435,7 @@ export const useSendChatRequest = () => {
       const dispatchedAction = dispatch(action);
       abortRef.current = dispatchedAction.abort;
     },
-    [chatId, dispatch, toolsRequest.data],
+    [chatId, dispatch, toolsRequest.data, toolUse],
   );
 
   const submit = useCallback(

diff --git a/src/services/refact/tools.ts b/src/services/refact/tools.ts
@@ -41,6 +41,7 @@ export type ToolParams = {
 };
 
 export type ToolFunction = {
+  agentic?: boolean;
   name: string;
   description: string;
   parameters: ToolParams[];