Skip to content

Commit 5d34593

Browse files
bekmuradovbekmuradov
committed
feat: present retrieved chunks
1 parent 2b4f9bb commit 5d34593

File tree

5 files changed

+197
-67
lines changed

5 files changed

+197
-67
lines changed

‎src/braindrive-plugin/pluginTypes.ts‎

Lines changed: 48 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,26 @@ export interface Document {
4848
chunk_count: number;
4949
}
5050

51+
export interface DocumentChunkMetadata {
52+
created_at?: string;
53+
start_char?: number;
54+
end_char?: number;
55+
token_count?: number;
56+
document_filename?: string;
57+
document_type?: string;
58+
chunk_token_count?: number;
59+
chunk_char_count?: number;
60+
processing_method?: string;
61+
context_prefix?: string;
62+
has_context?: boolean;
63+
bm25_score?: number;
64+
rrf_score?: number;
65+
fusion_method?: string;
66+
found_in_vector?: boolean;
67+
found_in_bm25?: boolean;
68+
[key: string]: any;
69+
}
70+
5171
export interface DocumentChunk {
5272
id: string;
5373
document_id: string;
@@ -56,7 +76,7 @@ export interface DocumentChunk {
5676
chunk_index: number;
5777
chunk_type: string;
5878
parent_chunk_id?: string;
59-
metadata: object;
79+
metadata: DocumentChunkMetadata;
6080
embedding_vector?: number[];
6181
}
6282

@@ -104,6 +124,15 @@ export interface ChatSession {
104124
message_count?: number;
105125
}
106126

127+
// Interface for web search
128+
export interface SearchResult {
129+
title: string;
130+
url: string;
131+
content: string;
132+
engine?: string;
133+
score?: number;
134+
}
135+
107136
export interface ChatMessage {
108137
id: string;
109138
sender: 'user' | 'ai';
@@ -117,13 +146,31 @@ export interface ChatMessage {
117146
canContinue?: boolean;
118147
canRegenerate?: boolean;
119148
isCutOff?: boolean;
149+
150+
// Search results
151+
isSearchResults?: boolean;
152+
searchData?: {
153+
query: string;
154+
results: SearchResult[];
155+
scrapedContent?: any;
156+
totalResults: number;
157+
successfulScrapes?: number;
158+
};
120159

121160
// Document context
122161
isDocumentContext?: boolean;
123162
documentData?: {
124163
results: DocumentProcessingResult[];
125164
context: string;
126165
};
166+
// Retrieved chunks context (collection retrieval)
167+
isRetrievedContext?: boolean;
168+
retrievalData?: {
169+
chunks: DocumentChunk[];
170+
context: string;
171+
intent?: IntentResponse | null;
172+
metadata?: Record<string, any>;
173+
};
127174
// Markdown toggle
128175
showRawMarkdown?: boolean;
129176
}

‎src/collection-chat-view/CollectionChatViewShell.tsx‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -2308,22 +2308,22 @@ export class CollectionChatViewShell extends React.Component<CollectionChatProps
23082308
}).join(", ");
23092309
if (relevantContext) {
23102310
enhancedPrompt = `\n${relevantContext}\n\nUser Question: ${prompt}`;
2311-
// Add document context to state
23122311
this.setState({ documentContext: relevantContext }, () => {
2313-
// Add a message to show the documents were processed
2314-
const documentMessage: ChatMessage = {
2315-
id: generateId('documents'),
2312+
// Add retrieved chunks preview message to UI
2313+
const retrievedMessage: ChatMessage = {
2314+
id: generateId('retrieval'),
23162315
sender: 'ai',
23172316
content: '',
23182317
timestamp: new Date().toISOString(),
2319-
isDocumentContext: true,
2320-
documentData: {
2321-
results: [],
2318+
isRetrievedContext: true,
2319+
retrievalData: {
2320+
chunks: contextRetrievalResult.chunks as any,
23222321
context: relevantContext,
2323-
}
2322+
intent: contextRetrievalResult.intent,
2323+
metadata: contextRetrievalResult.metadata,
2324+
},
23242325
};
2325-
2326-
this.addMessageToChat(documentMessage);
2326+
this.addMessageToChat(retrievedMessage);
23272327
});
23282328
}
23292329
}

‎src/collection-chat-view/components/ChatHistory.tsx‎

Lines changed: 68 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
import React from 'react';
22
import ReactMarkdown from 'react-markdown';
33
import remarkGfm from 'remark-gfm';
4-
import { ChatMessage, SearchResult } from '../../types';
4+
import { ChatMessage, SearchResult } from '../../braindrive-plugin/pluginTypes';
5+
import RetrievedChunksPreview from './RetrievedChunksPreview';
56
import { formatTimestamp } from '../../utils';
67
import EnhancedCodeBlock from './EnhancedCodeBlock';
78
import ThinkingBlock from './ThinkingBlock';
@@ -50,14 +51,16 @@ interface ChatHistoryProps {
5051
interface ChatHistoryState {
5152
expandedSearchResults: Set<string>;
5253
expandedDocumentContext: Set<string>;
54+
expandedRetrievedContext: Set<string>;
5355
}
5456

5557
class ChatHistory extends React.Component<ChatHistoryProps, ChatHistoryState> {
5658
constructor(props: ChatHistoryProps) {
5759
super(props);
5860
this.state = {
5961
expandedSearchResults: new Set(),
60-
expandedDocumentContext: new Set()
62+
expandedDocumentContext: new Set(),
63+
expandedRetrievedContext: new Set()
6164
};
6265
}
6366

@@ -115,6 +118,21 @@ class ChatHistory extends React.Component<ChatHistoryProps, ChatHistoryState> {
115118
});
116119
};
117120

121+
/**
122+
* Toggle retrieved context expansion
123+
*/
124+
toggleRetrievedContext = (messageId: string) => {
125+
this.setState(prevState => {
126+
const newSet = new Set(prevState.expandedRetrievedContext);
127+
if (newSet.has(messageId)) {
128+
newSet.delete(messageId);
129+
} else {
130+
newSet.add(messageId);
131+
}
132+
return { expandedRetrievedContext: newSet } as Pick<ChatHistoryState, keyof ChatHistoryState>;
133+
});
134+
};
135+
118136
/**
119137
* Render search result item
120138
*/
@@ -197,6 +215,49 @@ class ChatHistory extends React.Component<ChatHistoryProps, ChatHistoryState> {
197215
);
198216
};
199217

218+
/**
219+
* Render retrieved chunks context message
220+
*/
221+
renderRetrievedContext = (message: ChatMessage) => {
222+
const { retrievalData } = message;
223+
if (!retrievalData) return null;
224+
225+
const isExpanded = this.state.expandedRetrievedContext.has(message.id);
226+
const { chunks } = retrievalData;
227+
228+
return (
229+
<div key={message.id} className="message message-ai message-retrieved-context">
230+
<div className="message-bubble">
231+
<div className="message-body">
232+
<div className="retrieved-context-header">
233+
<div className="retrieved-context-summary">
234+
<span className="retrieved-icon">🔎</span>
235+
<span className="retrieved-count">{chunks.length} retrieved excerpt{chunks.length === 1 ? '' : 's'}</span>
236+
</div>
237+
<button
238+
onClick={() => this.toggleRetrievedContext(message.id)}
239+
className="retrieved-toggle-btn"
240+
title={isExpanded ? 'Hide retrieved context' : 'Show retrieved context'}
241+
>
242+
<span className="retrieved-toggle-text">{isExpanded ? 'Hide' : 'Show'}</span>
243+
<span className={`retrieved-toggle-icon ${isExpanded ? 'expanded' : ''}`}>{isExpanded ? '▼' : '▶'}</span>
244+
</button>
245+
</div>
246+
247+
{isExpanded && (
248+
<div className="retrieved-context-content">
249+
<RetrievedChunksPreview chunks={chunks} intent={message.retrievalData?.intent} metadata={message.retrievalData?.metadata} />
250+
</div>
251+
)}
252+
</div>
253+
<div className="message-meta">
254+
<span className="message-timestamp">{formatTimestamp(message.timestamp)}</span>
255+
</div>
256+
</div>
257+
</div>
258+
);
259+
};
260+
200261

201262

202263
/**
@@ -276,6 +337,11 @@ class ChatHistory extends React.Component<ChatHistoryProps, ChatHistoryState> {
276337
return this.renderDocumentContext(message);
277338
}
278339

340+
// Handle retrieved context messages separately
341+
if (message.isRetrievedContext) {
342+
return this.renderRetrievedContext(message);
343+
}
344+
279345
// Handle thinking tags in content
280346
const { sender, content, timestamp, isStreaming, isEditable, isEdited, canRegenerate, canContinue, showRawMarkdown } = message;
281347

‎src/collection-chat-view/components/RetrievedChunksPreview.tsx‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
import React from 'react';
2+
import ReactMarkdown from 'react-markdown';
3+
import remarkGfm from 'remark-gfm';
4+
import type { DocumentChunk, IntentResponse } from '../../braindrive-plugin/pluginTypes';
5+
6+
interface RetrievedChunksPreviewProps {
7+
chunks: DocumentChunk[];
8+
intent?: IntentResponse | null;
9+
metadata?: Record<string, any>;
10+
}
11+
12+
const truncate = (text: string, max = 280) => {
13+
if (!text) return '';
14+
const cleaned = text.replace(/\s+/g, ' ').trim();
15+
if (cleaned.length <= max) return cleaned;
16+
return cleaned.slice(0, max) + '…';
17+
};
18+
19+
const getChunkLabel = (chunk: DocumentChunk, index: number) => {
20+
const idx = typeof chunk.chunk_index === 'number' ? `#${chunk.chunk_index}` : `#${index}`;
21+
return `Excerpt ${idx}`;
22+
};
23+
24+
const RetrievedChunksPreview: React.FC<RetrievedChunksPreviewProps> = ({ chunks, intent, metadata }) => {
25+
if (!Array.isArray(chunks) || chunks.length === 0) {
26+
return null;
27+
}
28+
29+
return (
30+
<div className="retrieved-chunks-preview">
31+
<div className="retrieved-chunks-header">
32+
{intent && (
33+
<div className="retrieved-intent">
34+
<div className="retrieved-intent-type">Intent: <strong>{intent.type}</strong></div>
35+
{intent.reasoning && (
36+
<div className="retrieved-intent-reasoning">{intent.reasoning}</div>
37+
)}
38+
</div>
39+
)}
40+
{metadata && Array.isArray(metadata.transformed_queries) && metadata.transformed_queries.length > 0 && (
41+
<div className="retrieved-transformed-queries">
42+
<div className="retrieved-queries-title">Transformed queries:</div>
43+
<ul className="retrieved-queries-list">
44+
{metadata.transformed_queries.slice(0, 5).map((q: string, i: number) => (
45+
<li key={i} className="retrieved-query-item">{q}</li>
46+
))}
47+
</ul>
48+
</div>
49+
)}
50+
</div>
51+
<div className="retrieved-chunks-list">
52+
{chunks.map((chunk, i) => (
53+
<div key={chunk.id || i} className="retrieved-chunk-item">
54+
<div className="retrieved-chunk-header">
55+
<span className="retrieved-chunk-source">{getChunkLabel(chunk, i)}</span>
56+
</div>
57+
<div className="retrieved-chunk-content">
58+
<ReactMarkdown remarkPlugins={[remarkGfm]}>
59+
{truncate(chunk.content, 1200)}
60+
</ReactMarkdown>
61+
</div>
62+
</div>
63+
))}
64+
</div>
65+
</div>
66+
);
67+
};
68+
69+
export default RetrievedChunksPreview;
70+
71+

‎src/types.ts‎

Lines changed: 0 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -1,57 +1,3 @@
1-
// Chat message types
2-
export interface ChatMessage {
3-
id: string;
4-
sender: 'user' | 'ai';
5-
content: string;
6-
timestamp: string;
7-
isStreaming?: boolean;
8-
// Search results
9-
isSearchResults?: boolean;
10-
searchData?: {
11-
query: string;
12-
results: SearchResult[];
13-
scrapedContent?: any;
14-
totalResults: number;
15-
successfulScrapes?: number;
16-
};
17-
// User control features
18-
isEditable?: boolean;
19-
isEdited?: boolean;
20-
originalContent?: string;
21-
canContinue?: boolean;
22-
canRegenerate?: boolean;
23-
isCutOff?: boolean;
24-
25-
// Document context
26-
isDocumentContext?: boolean;
27-
documentData?: {
28-
results: DocumentProcessingResult[];
29-
context: string;
30-
};
31-
// Markdown toggle
32-
showRawMarkdown?: boolean;
33-
34-
// Web search context (for user messages)
35-
hasSearchContext?: boolean;
36-
searchContextData?: {
37-
originalPrompt: string;
38-
searchQuery: string;
39-
searchResults: SearchResult[];
40-
scrapedContent?: any;
41-
totalResults: number;
42-
successfulScrapes?: number;
43-
};
44-
}
45-
46-
// Interface for web search
47-
export interface SearchResult {
48-
title: string;
49-
url: string;
50-
content: string;
51-
engine?: string;
52-
score?: number;
53-
}
54-
551
// Model information
562
export interface ModelInfo {
573
name: string;

0 commit comments

Comments
 (0)