feat(app,backend): RAG pipeline improvements and token estimation (#31)

Author: Antonio Maiolo

.changeset/rag-improvements-token-estimation.md

@@ -0,0 +1,38 @@
+---
+"think-app": patch
+---
+
+# RAG Pipeline Improvements & Token Estimation
+
+## Frontend Changes
+
+### Token Usage Estimation
+- Replaced API-provided token counts with client-side estimation (~4 chars/token)
+- Context usage indicator now shows approximate values with `~` prefix
+- More stable UI without flickering from API response timing
+
+### Model Selector Optimization
+- Added provider tracking to prevent unnecessary re-fetches during polling
+- Reduces API calls and eliminates visual flickering
+
+## Backend Changes
+
+### Improved Memory Filtering
+- Dynamic threshold-based filtering adapts to match quality
+- Tiered filtering: excellent matches allow more results, marginal matches are stricter
+- Skip RAG for very short messages (< 10 chars)
+
+### Enhanced Search Pipeline
+- Added match type tracking (vector/keyword/hybrid) and RRF scores
+- Graceful fallback to vector-only search if hybrid fails
+- Comprehensive logging throughout the search pipeline
+
+### Embedding Safety
+- Model-specific context windows (Ollama models have smaller limits than documented)
+- Intelligent text chunking with paragraph/sentence awareness
+- Parallel chunk processing with embedding averaging
+
+### Other Improvements
+- Blocked `all-minilm` embedding model (context too small)
+- Removed conversation history limit for fuller context
+- Added logging for query transformations

app/src/components/ContextUsageIndicator.tsx

@@ -3,23 +3,23 @@ import { cn } from "@/lib/utils";
 import type { TokenUsage } from "@/types/chat";
 
 interface ContextUsageIndicatorProps {
-  contextUsage: TokenUsage | null;  // Latest message (for context window %)
+  estimatedTokens: number;  // Estimated conversation tokens
   billingUsage: TokenUsage | null;  // Cumulative (for session totals)
   contextWindow: number;
   className?: string;
 }
 
 export function ContextUsageIndicator({
-  contextUsage,
+  estimatedTokens,
   billingUsage,
   contextWindow,
   className,
 }: ContextUsageIndicatorProps) {
   const [showPopover, setShowPopover] = useState(false);
 
-  if (!contextUsage) return null;
+  if (estimatedTokens === 0) return null;
 
-  const percentage = Math.min((contextUsage.total_tokens / contextWindow) * 100, 100);
+  const percentage = Math.min((estimatedTokens / contextWindow) * 100, 100);
 
   // Color based on usage level
   const getStrokeColor = () => {
@@ -78,9 +78,9 @@ export function ContextUsageIndicator({
           <div className="text-xs font-medium mb-2">Context Window</div>
           <div className="space-y-1.5 text-xs">
             <div className="flex justify-between">
-              <span className="text-muted-foreground">Current:</span>
+              <span className="text-muted-foreground">Conversation:</span>
               <span className="font-mono font-medium">
-                {contextUsage.total_tokens.toLocaleString()} / {contextWindow.toLocaleString()}
+                ~{estimatedTokens.toLocaleString()} / {contextWindow.toLocaleString()}
               </span>
             </div>
             <div className="flex justify-between text-muted-foreground">

app/src/components/ModelSelector.tsx

@@ -33,6 +33,7 @@ export function ModelSelector({ type = "chat", provider, selectedModel, onModelC
   const [pullingModel, setPullingModel] = useState<string | null>(null);
   const [pullProgress, setPullProgress] = useState(0);
   const dropdownRef = useRef<HTMLDivElement>(null);
+  const lastFetchedProviderRef = useRef<string | undefined>(undefined);
 
   // Use selectedModel prop if provided (controlled), otherwise use internal state
   const displayModel = selectedModel !== undefined ? selectedModel : currentModel;
@@ -61,14 +62,20 @@ export function ModelSelector({ type = "chat", provider, selectedModel, onModelC
 
   // Initial fetch on mount or when type/provider changes
   useEffect(() => {
+    // Skip if we already fetched for this provider (prevents flashing during polling)
+    if (lastFetchedProviderRef.current === provider && models.length > 0) {
+      return;
+    }
+    lastFetchedProviderRef.current = provider;
+
     setCurrentModel(""); // Clear stale model before fetch
     fetchModels().then((data) => {
       // Notify parent of the initial model (for controlled mode)
       if (data && type === "embedding" && selectedModel === undefined) {
         onModelChange?.(data.current_model);
       }
     });
-  }, [fetchModels, type, selectedModel, onModelChange]);
+  }, [fetchModels, type, selectedModel, onModelChange, provider, models.length]);
 
   // Close dropdown when clicking outside
   useEffect(() => {

app/src/contexts/ConversationContext.tsx

@@ -8,7 +8,7 @@ interface ConversationContextType {
   allSources: SourceMemory[];
   isLoadingMessages: boolean;
   pendingMessage: string | null;
-  contextUsage: TokenUsage | null;  // Latest message only (for context window %)
+  estimatedTokens: number;  // Estimated conversation tokens (stable, grows with messages)
   billingUsage: TokenUsage | null;  // Cumulative (for cost tracking)
   contextWindow: number;
   selectConversation: (conversation: Conversation | null) => void;
@@ -18,42 +18,34 @@ interface ConversationContextType {
   updateMessage: (id: string | number, updates: Partial<ChatMessage>) => void;
   clearMessages: () => void;
   setPendingMessage: (message: string | null) => void;
-  updateUsage: (usage: TokenUsage | null, contextWindow?: number) => void;
+  updateContextWindow: (contextWindow: number) => void;
 }
 
 const ConversationContext = createContext<ConversationContextType | null>(null);
 
+// Estimate tokens from text (~4 chars per token is a common approximation)
+const SYSTEM_PROMPT_TOKENS = 80; // Approximate tokens for the system prompt
+
+function estimateTokens(text: string): number {
+  return Math.ceil(text.length / 4);
+}
+
 export function ConversationProvider({ children }: { children: ReactNode }) {
   const [currentConversationId, setCurrentConversationId] = useState<number | null>(null);
   const [messages, setMessages] = useState<ChatMessage[]>([]);
   const [isLoadingMessages, setIsLoadingMessages] = useState(false);
   const [pendingMessage, setPendingMessage] = useState<string | null>(null);
-  const [contextUsage, setContextUsage] = useState<TokenUsage | null>(null);  // Latest only
   const [billingUsage, setBillingUsage] = useState<TokenUsage | null>(null);  // Cumulative
   const [contextWindow, setContextWindow] = useState(128000);
 
-  const updateUsage = useCallback((newUsage: TokenUsage | null, newContextWindow?: number) => {
-    if (newUsage) {
-      // Context usage = latest message only (for context window %)
-      setContextUsage(newUsage);
-
-      // Billing usage = accumulate across conversation (for cost tracking)
-      setBillingUsage((prev) =>
-        prev
-          ? {
-              prompt_tokens: prev.prompt_tokens + newUsage.prompt_tokens,
-              completion_tokens: prev.completion_tokens + newUsage.completion_tokens,
-              total_tokens: prev.total_tokens + newUsage.total_tokens,
-            }
-          : newUsage
-      );
-    } else {
-      setContextUsage(null);
-      setBillingUsage(null);
-    }
-    if (newContextWindow) {
-      setContextWindow(newContextWindow);
-    }
+  // Estimate conversation tokens from messages (stable, grows with conversation)
+  const estimatedTokens = useMemo(() => {
+    const messageTokens = messages.reduce((acc, msg) => acc + estimateTokens(msg.content), 0);
+    return SYSTEM_PROMPT_TOKENS + messageTokens;
+  }, [messages]);
+
+  const updateContextWindow = useCallback((newContextWindow: number) => {
+    setContextWindow(newContextWindow);
   }, []);
 
   const loadConversation = useCallback(async (conversationId: number) => {
@@ -70,24 +62,10 @@ export function ConversationProvider({ children }: { children: ReactNode }) {
           }))
         );
 
-        // Get usage from messages
+        // Billing usage = sum of all assistant message tokens (for cost tracking)
         const assistantMessages = data.messages.filter(
           (m) => m.role === "assistant" && m.total_tokens
         );
-
-        // Context usage = last assistant message only (current context window state)
-        const lastAssistant = assistantMessages[assistantMessages.length - 1];
-        if (lastAssistant) {
-          setContextUsage({
-            prompt_tokens: lastAssistant.prompt_tokens || 0,
-            completion_tokens: lastAssistant.completion_tokens || 0,
-            total_tokens: lastAssistant.total_tokens || 0,
-          });
-        } else {
-          setContextUsage(null);
-        }
-
-        // Billing usage = sum of all (total tokens consumed)
         const totalBilling = assistantMessages.reduce(
           (acc, m) => ({
             prompt_tokens: acc.prompt_tokens + (m.prompt_tokens || 0),
@@ -110,7 +88,7 @@ export function ConversationProvider({ children }: { children: ReactNode }) {
     }
   }, []);
 
-  // Aggregate all sources from messages, deduplicated by id
+  // Aggregate all sources from conversation, deduplicated by id
   const allSources = useMemo(() => {
     const sourceMap = new Map<number, SourceMemory>();
     for (const msg of messages) {
@@ -133,7 +111,6 @@ export function ConversationProvider({ children }: { children: ReactNode }) {
       } else {
         setCurrentConversationId(null);
         setMessages([]);
-        setContextUsage(null);
         setBillingUsage(null);
       }
     },
@@ -143,7 +120,6 @@ export function ConversationProvider({ children }: { children: ReactNode }) {
   const startNewChat = useCallback(() => {
     setCurrentConversationId(null);
     setMessages([]);
-    setContextUsage(null);
     setBillingUsage(null);
   }, []);
 
@@ -169,7 +145,7 @@ export function ConversationProvider({ children }: { children: ReactNode }) {
         allSources,
         isLoadingMessages,
         pendingMessage,
-        contextUsage,
+        estimatedTokens,
         billingUsage,
         contextWindow,
         selectConversation,
@@ -179,7 +155,7 @@ export function ConversationProvider({ children }: { children: ReactNode }) {
         updateMessage,
         clearMessages,
         setPendingMessage,
-        updateUsage,
+        updateContextWindow,
       }}
     >
       {children}

app/src/pages/ChatPage.tsx

@@ -27,7 +27,7 @@ export default function ChatPage() {
     allSources,
     isLoadingMessages,
     pendingMessage,
-    contextUsage,
+    estimatedTokens,
     billingUsage,
     contextWindow,
     setCurrentConversationId,
@@ -36,7 +36,7 @@ export default function ChatPage() {
     selectConversation,
     startNewChat,
     setPendingMessage,
-    updateUsage,
+    updateContextWindow,
   } = useConversation();
 
   const { conversations } = useConversations();
@@ -152,9 +152,9 @@ export default function ChatPage() {
                 updateMessage(assistantMessageId, { content });
               } else if (data.type === "done") {
                 updateMessage(assistantMessageId, { isStreaming: false });
-                // Update usage from stream response
-                if (data.usage) {
-                  updateUsage(data.usage, data.context_window);
+                // Update context window from stream response
+                if (data.context_window) {
+                  updateContextWindow(data.context_window);
                 }
               } else if (data.type === "error") {
                 updateMessage(assistantMessageId, {
@@ -180,8 +180,8 @@ export default function ChatPage() {
             updateMessage(assistantMessageId, { content });
           } else if (data.type === "done") {
             updateMessage(assistantMessageId, { isStreaming: false });
-            if (data.usage) {
-              updateUsage(data.usage, data.context_window);
+            if (data.context_window) {
+              updateContextWindow(data.context_window);
             }
           }
         } catch {
@@ -200,7 +200,7 @@ export default function ChatPage() {
       // Clear pending conversation ref
       pendingConversationRef.current = null;
     }
-  }, [addMessage, updateMessage, setCurrentConversationId, updateUsage]);
+  }, [addMessage, updateMessage, setCurrentConversationId, updateContextWindow]);
 
   // Handler for manual chat input
   const handleChat = useCallback(() => {
@@ -267,7 +267,7 @@ export default function ChatPage() {
                 />
               </div>
               <ContextUsageIndicator
-                contextUsage={contextUsage}
+                estimatedTokens={estimatedTokens}
                 billingUsage={billingUsage}
                 contextWindow={contextWindow}
               />