Merge pull request #7 from abhigyanpatwari/embeddings_pipeline

highlight tool fixes

src/components/RightPanel.tsx

@@ -383,9 +383,9 @@ export const RightPanel = () => {
                               
                               if (isInline) {
                                 return (
-                                  <code className="px-1 py-0.5 bg-surface rounded text-accent font-mono text-xs" {...props}>
-                                    {children}
-                                  </code>
+                                <code className="px-1 py-0.5 bg-surface rounded text-accent font-mono text-xs" {...props}>
+                                  {children}
+                                </code>
                                 );
                               }
                               

src/components/ToolCallCard.tsx

@@ -5,7 +5,7 @@
  * Shows the tool name, status, and when expanded, the query/args and result.
  */
 
-import { useState, useCallback } from 'react';
+import { useState, useCallback, useMemo } from 'react';
 import { ChevronDown, ChevronRight, Sparkles, Check, Loader2, AlertCircle, Eye, EyeOff } from 'lucide-react';
 import type { ToolCallInfo } from '../core/llm/types';
 import { useAppState } from '../hooks/useAppState';
@@ -111,17 +111,38 @@ const extractHighlightNodeIds = (result: string | undefined): string[] => {
 
 export const ToolCallCard = ({ toolCall, defaultExpanded = false }: ToolCallCardProps) => {
   const [isExpanded, setIsExpanded] = useState(defaultExpanded);
-  const { highlightedNodeIds, setHighlightedNodeIds } = useAppState();
+  const { highlightedNodeIds, setHighlightedNodeIds, graph } = useAppState();
   const status = getStatusDisplay(toolCall.status);
   const formattedArgs = formatArgs(toolCall.args);
   
   // Check if this is a highlight tool and extract node IDs
   const isHighlightTool = toolCall.name === 'highlight_in_graph';
-  const highlightNodeIds = isHighlightTool ? extractHighlightNodeIds(toolCall.result) : [];
+  const rawHighlightNodeIds = isHighlightTool ? extractHighlightNodeIds(toolCall.result) : [];
+  
+  // Resolve raw IDs to actual graph node IDs (handles partial ID matching)
+  const resolvedNodeIds = useMemo(() => {
+    if (rawHighlightNodeIds.length === 0 || !graph) return rawHighlightNodeIds;
+    
+    const graphNodeIds = graph.nodes.map(n => n.id);
+    const resolved: string[] = [];
+    
+    for (const rawId of rawHighlightNodeIds) {
+      if (graphNodeIds.includes(rawId)) {
+        resolved.push(rawId);
+      } else {
+        // Try partial match - find node whose ID ends with the raw ID
+        const found = graphNodeIds.find(gid => 
+          gid.endsWith(rawId) || gid.endsWith(':' + rawId)
+        );
+        if (found) resolved.push(found);
+      }
+    }
+    return resolved;
+  }, [rawHighlightNodeIds, graph]);
   
   // Check if these specific nodes are currently highlighted
-  const isHighlightActive = highlightNodeIds.length > 0 && 
-    highlightNodeIds.some(id => highlightedNodeIds.has(id));
+  const isHighlightActive = resolvedNodeIds.length > 0 && 
+    resolvedNodeIds.some(id => highlightedNodeIds.has(id));
   
   // Toggle highlight on/off
   const toggleHighlight = useCallback((e: React.MouseEvent) => {
@@ -131,9 +152,9 @@ export const ToolCallCard = ({ toolCall, defaultExpanded = false }: ToolCallCard
       setHighlightedNodeIds(new Set());
     } else {
       // Turn on - set these nodes as highlighted
-      setHighlightedNodeIds(new Set(highlightNodeIds));
+      setHighlightedNodeIds(new Set(resolvedNodeIds));
     }
-  }, [isHighlightActive, highlightNodeIds, setHighlightedNodeIds]);
+  }, [isHighlightActive, resolvedNodeIds, setHighlightedNodeIds]);
   
   return (
     <div className={`rounded-lg border ${status.borderColor} ${status.bgColor} overflow-hidden transition-all`}>
@@ -153,7 +174,7 @@ export const ToolCallCard = ({ toolCall, defaultExpanded = false }: ToolCallCard
         </span>
         
         {/* Highlight toggle button - only for highlight_in_graph tool with results */}
-        {isHighlightTool && highlightNodeIds.length > 0 && (
+        {isHighlightTool && resolvedNodeIds.length > 0 && (
           <button
             onClick={toggleHighlight}
             className={`flex items-center gap-1 px-2 py-0.5 rounded text-xs transition-colors ${

src/core/llm/tools.ts

@@ -600,9 +600,17 @@ export const createGraphRAGTools = (
     },
     {
       name: 'highlight_in_graph',
-      description: 'Highlight specific nodes in the visual knowledge graph that the user can see alongside this chat. Use this to visually show the user which code elements you are discussing. Pass the node IDs from previous search/query results.',
+      description: `Highlight specific nodes in the visual knowledge graph. Pass the EXACT node IDs from your query results.
+
+IMPORTANT: Node IDs include a label prefix! Format is: Label:filepath:name
+Examples:
+- Class:src/agents/base.py:BaseAgent
+- Function:src/utils.ts:calculateSum
+- File:src/main.py
+
+Copy the ID EXACTLY as it appears in query results (the "classId", "fnId", "fileId", etc. columns).`,
       schema: z.object({
-        nodeIds: z.array(z.string()).describe('Array of node IDs to highlight (from search results or queries)'),
+        nodeIds: z.array(z.string()).describe('Array of EXACT node IDs to highlight - must include the label prefix like "Class:" or "Function:"'),
         description: z.string().optional().nullable().describe('Brief description of what these nodes represent'),
       }),
     }

src/hooks/useAppState.tsx

@@ -527,9 +527,35 @@ export const AppStateProvider = ({ children }: { children: ReactNode }) => {
               if (tc.result) {
                 const highlightMatch = tc.result.match(/\[HIGHLIGHT_NODES:([^\]]+)\]/);
                 if (highlightMatch) {
-                  const nodeIds = highlightMatch[1].split(',').map((id: string) => id.trim()).filter(Boolean);
-                  if (nodeIds.length > 0) {
-                    setHighlightedNodeIds(new Set(nodeIds));
+                  const rawIds = highlightMatch[1].split(',').map((id: string) => id.trim()).filter(Boolean);
+                  if (rawIds.length > 0 && graph) {
+                    // Try to match IDs against actual graph nodes
+                    // This handles cases where the LLM passes partial IDs without the label prefix
+                    const matchedIds = new Set<string>();
+                    const graphNodeIds = graph.nodes.map(n => n.id);
+                    
+                    for (const rawId of rawIds) {
+                      // First try exact match
+                      if (graphNodeIds.includes(rawId)) {
+                        matchedIds.add(rawId);
+                      } else {
+                        // Try to find a node whose ID ends with the raw ID
+                        // e.g., "src/path:ClassName" should match "Class:src/path:ClassName"
+                        const found = graphNodeIds.find(gid => 
+                          gid.endsWith(rawId) || gid.endsWith(':' + rawId)
+                        );
+                        if (found) {
+                          matchedIds.add(found);
+                        }
+                      }
+                    }
+                    
+                    if (matchedIds.size > 0) {
+                      setHighlightedNodeIds(matchedIds);
+                    }
+                  } else if (rawIds.length > 0) {
+                    // Fallback: just use the IDs directly if no graph available
+                    setHighlightedNodeIds(new Set(rawIds));
                   }
                 }
               }