smallcloudai · alashchev17 · Feb 28, 2025 · Feb 23, 2025 · Feb 24, 2025 · Feb 26, 2025
diff --git a/refact-agent/engine/src/http/routers/v1/at_commands.rs b/refact-agent/engine/src/http/routers/v1/at_commands.rs
@@ -5,7 +5,7 @@ use indexmap::IndexMap;
 use serde::{Deserialize, Serialize};
 use std::sync::Arc;
 use std::sync::RwLock as StdRwLock;
-use serde_json::json;
+use serde_json::{json, Value};
 use tokio::sync::RwLock as ARwLock;
 use tokio::sync::Mutex as AMutex;
 use strsim::jaro_winkler;
@@ -17,13 +17,15 @@ use crate::at_commands::execute_at::run_at_commands_locally;
 use crate::cached_tokenizers;
 use crate::at_commands::at_commands::AtCommandsContext;
 use crate::at_commands::execute_at::{execute_at_commands_in_query, parse_words_from_line};
-use crate::call_validation::{PostprocessSettings, SubchatParameters};
+use crate::call_validation::{ChatMeta, PostprocessSettings, SubchatParameters};
 use crate::custom_error::ScratchError;
 use crate::global_context::try_load_caps_quickly_if_not_present;
 use crate::global_context::GlobalContext;
 use crate::call_validation::{ChatMessage, ChatContent, ContextEnum};
 use crate::at_commands::at_commands::filter_only_context_file_from_context_tool;
+use crate::http::routers::v1::chat::deserialize_messages_from_post;
 use crate::postprocessing::pp_context_files::postprocess_context_files;
+use crate::scratchpads::chat_utils_prompts::prepend_the_right_system_prompt_and_maybe_more_initial_messages;
 use crate::scratchpads::scratchpad_utils::max_tokens_for_rag_chat;
 use crate::scratchpads::scratchpad_utils::HasRagResults;
 
@@ -43,9 +45,12 @@ struct CommandCompletionResponse {
 
 #[derive(Serialize, Deserialize, Clone)]
 struct CommandPreviewPost {
-    query: String,
+    #[serde(default)]
+    pub messages: Vec<Value>,
     #[serde(default)]
     model: String,
+    #[serde(default)]
+    pub meta: ChatMeta,
 }
 
 #[derive(Serialize, Deserialize, Clone)]
@@ -124,13 +129,40 @@ pub async fn handle_v1_command_completion(
         .unwrap())
 }
 
+async fn count_tokens(tokenizer_arc: Arc<StdRwLock<Tokenizer>>, messages: &Vec<ChatMessage>) -> Result<u64, ScratchError> {
+    let mut accum: u64 = 0;
+
+    for message in messages {
+        accum += message.content.count_tokens(tokenizer_arc.clone(), &None)
+            .map_err(|e| ScratchError {
+                status_code: StatusCode::INTERNAL_SERVER_ERROR,
+                message: format!("v1_chat_token_counter: count_tokens failed: {}", e),
+                telemetry_skip: false})? as u64;
+    }
+    Ok(accum)
+}
+
 pub async fn handle_v1_command_preview(
     Extension(global_context): Extension<Arc<ARwLock<GlobalContext>>>,
     body_bytes: hyper::body::Bytes,
 ) -> Result<Response<Body>, ScratchError> {
     let post = serde_json::from_slice::<CommandPreviewPost>(&body_bytes)
         .map_err(|e| ScratchError::new(StatusCode::UNPROCESSABLE_ENTITY, format!("JSON problem: {}", e)))?;
-    let mut query = post.query.clone();
+    let mut messages = deserialize_messages_from_post(&post.messages)?;
+
+    let last_message = messages.pop().unwrap();
+    let mut query = match &last_message.content {
+        ChatContent::SimpleText(query) => query.clone(),
+        ChatContent::Multimodal(elements) => {
+            let mut query = String::new();
+            for element in elements {
+                if element.is_text() { // use last text, but expected to be only one
+                    query = element.m_content.clone();
+                }    
+            }
+            query
+        }
+    };
 
     let caps = crate::global_context::try_load_caps_quickly_if_not_present(global_context.clone(), 0).await?;
     let (model_name, recommended_model_record) = {
@@ -219,10 +251,31 @@ pub async fn handle_v1_command_preview(
             reason: h.reason.unwrap_or_default(),
         })
     }
+
+    messages.extend(preview.clone());
+    let mut new_last_message = last_message.clone();
+    {
+        match &mut new_last_message.content {
+            ChatContent::SimpleText(_) => {new_last_message.content = ChatContent::SimpleText(query.clone());}
+            ChatContent::Multimodal(elements) => {
+                for elem in elements {
+                    if elem.is_text() {
+                        elem.m_content = query.clone();
+                    }
+                }
+            }
+        }
+    }
+    messages.push(new_last_message);
+    let messages = prepend_the_right_system_prompt_and_maybe_more_initial_messages(
+        global_context.clone(), messages.clone(), &post.meta, &mut HasRagResults::new()).await;
+    let tokens_number = count_tokens(tokenizer_arc, &messages).await?;
+
     Ok(Response::builder()
         .status(StatusCode::OK)
         .body(Body::from(serde_json::to_string_pretty(
-            &json!({"messages": preview, "model": model_name, "highlight": highlights})
+            &json!({"messages": preview, "model": model_name, "highlight": highlights, 
+                "current_context": tokens_number, "number_context": recommended_model_record.n_ctx})
         ).unwrap()))
         .unwrap())
 }

diff --git a/refact-agent/gui/src/components/ChatContent/ChatContent.tsx b/refact-agent/gui/src/components/ChatContent/ChatContent.tsx
@@ -118,7 +118,7 @@ export const ChatContent: React.FC<ChatContentProps> = ({
         {messages.length === 0 && <PlaceHolderText />}
         {renderMessages(messages, onRetryWrapper)}
         <UncommittedChangesWarning />
-        {threadUsage && <UsageCounter usage={threadUsage} />}
+        {threadUsage && messages.length > 0 && <UsageCounter />}
 
         <Container py="4">
           <Spinner spinning={isWaiting && !isWaitingForConfirmation} />

diff --git a/refact-agent/gui/src/components/ChatForm/ChatForm.module.css b/refact-agent/gui/src/components/ChatForm/ChatForm.module.css
@@ -1,10 +1,20 @@
-.buttonGroup {
+.textareaInteractive {
   position: absolute;
   /* display: flex;
     gap: 4px; */
+  left: 2px;
   right: 6px;
   bottom: 8px;
   padding: 1px 6px;
+  justify-content: flex-end;
+  width: calc(100% - 6px - 2px);
+}
+.buttonGroup {
+  align-self: flex-end;
+  justify-self: flex-end;
+  &:first-child {
+    margin-left: auto;
+  }
 }
 
 .button {
@@ -17,6 +27,13 @@
   border-radius: 6px;
   flex-shrink: 0;
   overflow: hidden;
+  --interactive-group-height: 17px;
+  min-height: calc(100% + var(--interactive-group-height));
+  background-color: var(--color-surface);
+}
+
+.chatForm__form :global(.rt-TextAreaRoot) {
+  background-color: unset;
 }
 
 .file {
@@ -29,6 +46,7 @@
   overflow-wrap: break-word;
   word-break: break-word;
   white-space: normal;
+  max-height: 16px;
 }
 
 .file_name {

diff --git a/refact-agent/gui/src/components/ChatForm/ChatForm.tsx b/refact-agent/gui/src/components/ChatForm/ChatForm.tsx
@@ -48,12 +48,15 @@ import {
   selectIsWaiting,
   selectMessages,
   selectPreventSend,
+  selectThreadMaximumTokens,
   selectThreadToolUse,
   selectToolUse,
 } from "../../features/Chat";
 import { telemetryApi } from "../../services/refact";
 import { push } from "../../features/Pages/pagesSlice";
 import { AgentCapabilities } from "./AgentCapabilities";
+import { TokensPreview } from "./TokensPreview";
+import { useUsageCounter } from "../UsageCounter/useUsageCounter";
 
 export type ChatFormProps = {
   onSubmit: (str: string) => void;
@@ -65,7 +68,7 @@ export type ChatFormProps = {
 export const ChatForm: React.FC<ChatFormProps> = ({
   onSubmit,
   onClose,
-  className,
+  // className,
   unCalledTools,
 }) => {
   const dispatch = useAppDispatch();
@@ -87,6 +90,12 @@ export const ChatForm: React.FC<ChatFormProps> = ({
   const threadToolUse = useAppSelector(selectThreadToolUse);
   const messages = useAppSelector(selectMessages);
   const preventSend = useAppSelector(selectPreventSend);
+  const currentThreadMaximumContextTokens = useAppSelector(
+    selectThreadMaximumTokens,
+  );
+
+  const { isOverflown: arePromptTokensBiggerThanContext, currentThreadUsage } =
+    useUsageCounter();
 
   const shouldAgentCapabilitiesBeShown = useMemo(() => {
     return threadToolUse === "agent" && toolUse === "agent";
@@ -109,9 +118,26 @@ export const ChatForm: React.FC<ChatFormProps> = ({
   const disableSend = useMemo(() => {
     // TODO: if interrupting chat some errors can occur
     if (allDisabled) return true;
+    if (
+      currentThreadMaximumContextTokens &&
+      currentThreadUsage?.prompt_tokens &&
+      currentThreadUsage.prompt_tokens > currentThreadMaximumContextTokens
+    )
+      return false;
+    if (arePromptTokensBiggerThanContext) return true;
     if (messages.length === 0) return false;
     return isWaiting || isStreaming || !isOnline || preventSend;
-  }, [isOnline, isStreaming, isWaiting, preventSend, messages, allDisabled]);
+  }, [
+    isOnline,
+    isStreaming,
+    isWaiting,
+    arePromptTokensBiggerThanContext,
+    currentThreadMaximumContextTokens,
+    currentThreadUsage?.prompt_tokens,
+    preventSend,
+    messages,
+    allDisabled,
+  ]);
 
   const { processAndInsertImages } = useAttachedImages();
   const handlePastingFile = useCallback(
@@ -323,7 +349,7 @@ export const ChatForm: React.FC<ChatFormProps> = ({
         {shouldAgentCapabilitiesBeShown && <AgentCapabilities />}
         <Form
           disabled={disableSend}
-          className={className}
+          className={styles.chatForm__form}
           onSubmit={handleSubmit}
         >
           <FilesPreview files={previewFiles} />
@@ -352,33 +378,40 @@ export const ChatForm: React.FC<ChatFormProps> = ({
               />
             )}
           />
-          <Flex gap="2" className={styles.buttonGroup}>
-            {toolUse === "agent" && (
-              <AgentIntegrationsButton
-                title="Set up Agent Integrations"
-                size="1"
-                type="button"
-                onClick={handleAgentIntegrationsClick}
-                ref={(x) => refs.setSetupIntegrations(x)}
-              />
-            )}
-            {onClose && (
-              <BackToSideBarButton
-                disabled={isStreaming}
-                title="Return to sidebar"
+          <Flex
+            className={styles.textareaInteractive}
+            align="center"
+            justify="between"
+          >
+            <TokensPreview />
+            <Flex gap="2" className={styles.buttonGroup}>
+              {toolUse === "agent" && (
+                <AgentIntegrationsButton
+                  title="Set up Agent Integrations"
+                  size="1"
+                  type="button"
+                  onClick={handleAgentIntegrationsClick}
+                  ref={(x) => refs.setSetupIntegrations(x)}
+                />
+              )}
+              {onClose && (
+                <BackToSideBarButton
+                  disabled={isStreaming}
+                  title="Return to sidebar"
+                  size="1"
+                  onClick={onClose}
+                />
+              )}
+              {config.features?.images !== false &&
+                isMultimodalitySupportedForCurrentModel && <AttachFileButton />}
+              {/* TODO: Reserved space for microphone button coming later on */}
+              <PaperPlaneButton
+                disabled={disableSend || disableInput}
+                title="Send message"
                 size="1"
-                onClick={onClose}
+                type="submit"
               />
-            )}
-            {config.features?.images !== false &&
-              isMultimodalitySupportedForCurrentModel && <AttachFileButton />}
-            {/* TODO: Reserved space for microphone button coming later on */}
-            <PaperPlaneButton
-              disabled={disableSend || disableInput}
-              title="Send message"
-              size="1"
-              type="submit"
-            />
+            </Flex>
           </Flex>
         </Form>
       </Flex>

diff --git a/refact-agent/gui/src/components/ChatForm/SuggestNewChat/SuggestNewChat.tsx b/refact-agent/gui/src/components/ChatForm/SuggestNewChat/SuggestNewChat.tsx
@@ -1,6 +1,6 @@
 import { Box, Flex, IconButton, Text } from "@radix-ui/themes";
 import { Cross2Icon } from "@radix-ui/react-icons";
-import { useCallback, useEffect, useState } from "react";
+import { useCallback, useEffect, useMemo, useState } from "react";
 
 import { clearPauseReasonsAndHandleToolsStatus } from "../../../features/ToolConfirmation/confirmationSlice";
 import { useAppDispatch, useAppSelector } from "../../../hooks";
@@ -9,13 +9,15 @@ import { telemetryApi } from "../../../services/refact";
 import {
   newChatAction,
   selectChatId,
+  selectThreadNewChatSuggested,
   setIsNewChatSuggestionRejected,
 } from "../../../features/Chat";
 
 import { Link } from "../../Link";
 
 import styles from "./SuggestNewChat.module.css";
 import classNames from "classnames";
+import { useUsageCounter } from "../../UsageCounter/useUsageCounter";
 
 type SuggestNewChatProps = {
   shouldBeVisible?: boolean;
@@ -26,8 +28,12 @@ export const SuggestNewChat = ({
 }: SuggestNewChatProps) => {
   const dispatch = useAppDispatch();
   const chatId = useAppSelector(selectChatId);
+  const { isMandatory } = useAppSelector(selectThreadNewChatSuggested);
   const [sendTelemetryEvent] =
     telemetryApi.useLazySendTelemetryChatEventQuery();
+
+  const { isWarning, isOverflown: isContextOverflown } = useUsageCounter();
+
   const [isRendered, setIsRendered] = useState(shouldBeVisible);
   const [isAnimating, setIsAnimating] = useState(false);
 
@@ -79,6 +85,14 @@ export const SuggestNewChat = ({
     });
   }, [dispatch, sendTelemetryEvent]);
 
+  const tipText = useMemo(() => {
+    if (isWarning)
+      return "Long chats cause you to reach your usage limits faster.";
+    if (isContextOverflown)
+      return "Maximum available context for this chat is exceeded. Consider starting a new chat.";
+    return "Agent performs better for chats which don't switch topics often.";
+  }, [isWarning, isContextOverflown]);
+
   return (
     <Box
       py="3"
@@ -92,22 +106,23 @@ export const SuggestNewChat = ({
     >
       <Flex align="center" justify="between" gap="2">
         <Text size="1">
-          <Text weight="bold">Tip:</Text> Long chats cause you to reach your
-          usage limits faster.
+          <Text weight="bold">Tip:</Text> {tipText}
         </Text>
         <Flex align="center" gap="3" flexShrink="0">
           <Link size="1" onClick={onCreateNewChat} color="indigo">
             Start a new chat
           </Link>
-          <IconButton
-            asChild
-            variant="ghost"
-            color="violet"
-            size="1"
-            onClick={handleClose}
-          >
-            <Cross2Icon />
-          </IconButton>
+          {!isMandatory && (
+            <IconButton
+              asChild
+              variant="ghost"
+              color="violet"
+              size="1"
+              onClick={handleClose}
+            >
+              <Cross2Icon />
+            </IconButton>
+          )}
         </Flex>
       </Flex>
     </Box>