Skip to content

Commit 09ecd11

Browse files
reakaleekreakaleek
committed
Add ability to switch between LLM Gateway and Agent Builder
1 parent f1cf415 commit 09ecd11

30 files changed

+1808
-371
lines changed

‎actions/update-link-index/README.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
<!-- Generated by https://github.com/reakaleek/gh-action-readme -->
12
<!--
23
this documentation was generated by https://github.com/reakaleek/gh-action-readme
34
with the command `VERSION=main gh action-readme update`

‎src/Elastic.Documentation.Site/Assets/eui-icons-cache.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@ import { icon as EuiIconCopy } from '@elastic/eui/es/components/icon/assets/copy
77
import { icon as EuiIconCopyClipboard } from '@elastic/eui/es/components/icon/assets/copy_clipboard'
88
import { icon as EuiIconCross } from '@elastic/eui/es/components/icon/assets/cross'
99
import { icon as EuiIconDocument } from '@elastic/eui/es/components/icon/assets/document'
10+
import { icon as EuiIconDot } from '@elastic/eui/es/components/icon/assets/dot'
11+
import { icon as EuiIconEmpty } from '@elastic/eui/es/components/icon/assets/empty'
1012
import { icon as EuiIconError } from '@elastic/eui/es/components/icon/assets/error'
1113
import { icon as EuiIconFaceHappy } from '@elastic/eui/es/components/icon/assets/face_happy'
1214
import { icon as EuiIconFaceSad } from '@elastic/eui/es/components/icon/assets/face_sad'
@@ -32,6 +34,8 @@ appendIconComponentCache({
3234
arrowLeft: EuiIconArrowLeft,
3335
arrowRight: EuiIconArrowRight,
3436
document: EuiIconDocument,
37+
dot: EuiIconDot,
38+
empty: EuiIconEmpty,
3539
search: EuiIconSearch,
3640
trash: EuiIconTrash,
3741
user: EuiIconUser,

‎src/Elastic.Documentation.Site/Assets/web-components/SearchOrAskAi/AskAi/AiProviderSelector.tsx‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
/** @jsxImportSource @emotion/react */
2+
import { EuiRadioGroup } from '@elastic/eui'
3+
import type { EuiRadioGroupOption } from '@elastic/eui'
4+
import { css } from '@emotion/react'
5+
import { useAiProviderStore } from './aiProviderStore'
6+
7+
const containerStyles = css`
8+
padding: 1rem;
9+
display: flex;
10+
justify-content: center;
11+
`
12+
13+
const options: EuiRadioGroupOption[] = [
14+
{
15+
id: 'LlmGateway',
16+
label: 'LLM Gateway',
17+
},
18+
{
19+
id: 'AgentBuilder',
20+
label: 'Agent Builder',
21+
},
22+
]
23+
24+
export const AiProviderSelector = () => {
25+
const { provider, setProvider } = useAiProviderStore()
26+
27+
return (
28+
<div css={containerStyles}>
29+
<EuiRadioGroup
30+
options={options}
31+
idSelected={provider}
32+
onChange={(id) => setProvider(id as 'AgentBuilder' | 'LlmGateway')}
33+
name="aiProvider"
34+
legend={{
35+
children: 'AI Provider',
36+
display: 'visible',
37+
}}
38+
/>
39+
</div>
40+
)
41+
}

‎src/Elastic.Documentation.Site/Assets/web-components/SearchOrAskAi/AskAi/AskAiEvent.ts‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
// Canonical AskAI event types - matches backend AskAiEvent records
2+
import * as z from 'zod'
3+
4+
// Event type constants for type-safe referencing
5+
export const EventTypes = {
6+
CONVERSATION_START: 'conversation_start',
7+
CHUNK: 'chunk',
8+
CHUNK_COMPLETE: 'chunk_complete',
9+
SEARCH_TOOL_CALL: 'search_tool_call',
10+
TOOL_CALL: 'tool_call',
11+
TOOL_RESULT: 'tool_result',
12+
REASONING: 'reasoning',
13+
CONVERSATION_END: 'conversation_end',
14+
ERROR: 'error',
15+
} as const
16+
17+
// Individual event schemas
18+
export const ConversationStartEventSchema = z.object({
19+
type: z.literal(EventTypes.CONVERSATION_START),
20+
id: z.string(),
21+
timestamp: z.number(),
22+
conversationId: z.string(),
23+
})
24+
25+
export const ChunkEventSchema = z.object({
26+
type: z.literal(EventTypes.CHUNK),
27+
id: z.string(),
28+
timestamp: z.number(),
29+
content: z.string(),
30+
})
31+
32+
export const ChunkCompleteEventSchema = z.object({
33+
type: z.literal(EventTypes.CHUNK_COMPLETE),
34+
id: z.string(),
35+
timestamp: z.number(),
36+
fullContent: z.string(),
37+
})
38+
39+
export const SearchToolCallEventSchema = z.object({
40+
type: z.literal(EventTypes.SEARCH_TOOL_CALL),
41+
id: z.string(),
42+
timestamp: z.number(),
43+
toolCallId: z.string(),
44+
searchQuery: z.string(),
45+
})
46+
47+
export const ToolCallEventSchema = z.object({
48+
type: z.literal(EventTypes.TOOL_CALL),
49+
id: z.string(),
50+
timestamp: z.number(),
51+
toolCallId: z.string(),
52+
toolName: z.string(),
53+
arguments: z.string(),
54+
})
55+
56+
export const ToolResultEventSchema = z.object({
57+
type: z.literal(EventTypes.TOOL_RESULT),
58+
id: z.string(),
59+
timestamp: z.number(),
60+
toolCallId: z.string(),
61+
result: z.string(),
62+
})
63+
64+
export const ReasoningEventSchema = z.object({
65+
type: z.literal(EventTypes.REASONING),
66+
id: z.string(),
67+
timestamp: z.number(),
68+
message: z.string().nullable(),
69+
})
70+
71+
export const ConversationEndEventSchema = z.object({
72+
type: z.literal(EventTypes.CONVERSATION_END),
73+
id: z.string(),
74+
timestamp: z.number(),
75+
})
76+
77+
export const ErrorEventSchema = z.object({
78+
type: z.literal(EventTypes.ERROR),
79+
id: z.string(),
80+
timestamp: z.number(),
81+
message: z.string(),
82+
})
83+
84+
// Discriminated union of all event types
85+
export const AskAiEventSchema = z.discriminatedUnion('type', [
86+
ConversationStartEventSchema,
87+
ChunkEventSchema,
88+
ChunkCompleteEventSchema,
89+
SearchToolCallEventSchema,
90+
ToolCallEventSchema,
91+
ToolResultEventSchema,
92+
ReasoningEventSchema,
93+
ConversationEndEventSchema,
94+
ErrorEventSchema,
95+
])
96+
97+
// Infer TypeScript types from schemas
98+
export type ConversationStartEvent = z.infer<
99+
typeof ConversationStartEventSchema
100+
>
101+
export type ChunkEvent = z.infer<typeof ChunkEventSchema>
102+
export type ChunkCompleteEvent = z.infer<typeof ChunkCompleteEventSchema>
103+
export type SearchToolCallEvent = z.infer<typeof SearchToolCallEventSchema>
104+
export type ToolCallEvent = z.infer<typeof ToolCallEventSchema>
105+
export type ToolResultEvent = z.infer<typeof ToolResultEventSchema>
106+
export type ReasoningEvent = z.infer<typeof ReasoningEventSchema>
107+
export type ConversationEndEvent = z.infer<typeof ConversationEndEventSchema>
108+
export type ErrorEvent = z.infer<typeof ErrorEventSchema>
109+
export type AskAiEvent = z.infer<typeof AskAiEventSchema>

‎src/Elastic.Documentation.Site/Assets/web-components/SearchOrAskAi/AskAi/Chat.tsx‎

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
/** @jsxImportSource @emotion/react */
22
import { AskAiSuggestions } from './AskAiSuggestions'
3+
import { AiProviderSelector } from './AiProviderSelector'
34
import { ChatMessageList } from './ChatMessageList'
45
import { useChatActions, useChatMessages } from './chat.store'
56
import {
@@ -137,12 +138,16 @@ export const Chat = () => {
137138
<h2>Hi! I'm the Elastic Docs AI Assistant</h2>
138139
}
139140
body={
140-
<p>
141-
I can help answer your questions about
142-
Elastic documentation. <br />
143-
Ask me anything about Elasticsearch, Kibana,
144-
Observability, Security, and more.
145-
</p>
141+
<>
142+
<p>
143+
I can help answer your questions about
144+
Elastic documentation. <br />
145+
Ask me anything about Elasticsearch, Kibana,
146+
Observability, Security, and more.
147+
</p>
148+
<EuiSpacer size="m" />
149+
<AiProviderSelector />
150+
</>
146151
}
147152
footer={
148153
<>

‎src/Elastic.Documentation.Site/Assets/web-components/SearchOrAskAi/AskAi/ChatMessage.tsx‎

Lines changed: 79 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
11
import { initCopyButton } from '../../../copybutton'
22
import { hljs } from '../../../hljs'
3+
import { AskAiEvent, EventTypes } from './AskAiEvent'
34
import { GeneratingStatus } from './GeneratingStatus'
45
import { References } from './RelatedResources'
56
import { ChatMessage as ChatMessageType } from './chat.store'
6-
import { LlmGatewayMessage } from './useLlmGateway'
7+
import { useStatusMinDisplay } from './useStatusMinDisplay'
78
import {
89
EuiButtonIcon,
910
EuiCallOut,
@@ -56,16 +57,16 @@ const markedInstance = createMarkedInstance()
5657

5758
interface ChatMessageProps {
5859
message: ChatMessageType
59-
llmMessages?: LlmGatewayMessage[]
60+
events?: AskAiEvent[]
6061
streamingContent?: string
6162
error?: Error | null
6263
onRetry?: () => void
6364
}
6465

65-
const getAccumulatedContent = (messages: LlmGatewayMessage[]) => {
66+
const getAccumulatedContent = (messages: AskAiEvent[]) => {
6667
return messages
67-
.filter((m) => m.type === 'ai_message_chunk')
68-
.map((m) => m.data.content)
68+
.filter((m) => m.type === 'chunk')
69+
.map((m) => m.content)
6970
.join('')
7071
}
7172

@@ -100,57 +101,86 @@ const getMessageState = (message: ChatMessageType) => ({
100101
hasError: message.status === 'error',
101102
})
102103

103-
// Helper functions for computing AI status
104-
const getToolCallSearchQuery = (
105-
messages: LlmGatewayMessage[]
106-
): string | null => {
107-
const toolCallMessage = messages.find((m) => m.type === 'tool_call')
108-
if (!toolCallMessage) return null
104+
// Status message constants
105+
const STATUS_MESSAGES = {
106+
THINKING: 'Thinking',
107+
ANALYZING: 'Analyzing results',
108+
GATHERING: 'Gathering resources',
109+
GENERATING: 'Generating',
110+
} as const
109111

112+
// Helper to extract search query from tool call arguments
113+
const tryParseSearchQuery = (argsJson: string): string | null => {
110114
try {
111-
const toolCalls = toolCallMessage.data?.toolCalls
112-
if (toolCalls && toolCalls.length > 0) {
113-
const firstToolCall = toolCalls[0]
114-
return firstToolCall.args?.searchQuery || null
115-
}
116-
} catch (e) {
117-
console.error('Error extracting search query from tool call:', e)
115+
const args = JSON.parse(argsJson)
116+
return args.searchQuery || args.query || null
117+
} catch {
118+
return null
118119
}
119-
120-
return null
121120
}
122121

123-
const hasContentStarted = (messages: LlmGatewayMessage[]): boolean => {
124-
return messages.some((m) => m.type === 'ai_message_chunk' && m.data.content)
125-
}
122+
// Helper to get tool call status message
123+
const getToolCallStatus = (event: AskAiEvent): string => {
124+
if (event.type !== EventTypes.TOOL_CALL) {
125+
return STATUS_MESSAGES.THINKING
126+
}
126127

127-
const hasReachedReferences = (messages: LlmGatewayMessage[]): boolean => {
128-
const accumulatedContent = messages
129-
.filter((m) => m.type === 'ai_message_chunk')
130-
.map((m) => m.data.content)
131-
.join('')
132-
return accumulatedContent.includes('<!--REFERENCES')
128+
const query = tryParseSearchQuery(event.arguments)
129+
return query ? `Searching for "${query}"` : `Using ${event.toolName}`
133130
}
134131

132+
// Helper function for computing AI status - time-based latest status
135133
const computeAiStatus = (
136-
llmMessages: LlmGatewayMessage[],
134+
events: AskAiEvent[],
137135
isComplete: boolean
138136
): string | null => {
139137
if (isComplete) return null
140138

141-
const searchQuery = getToolCallSearchQuery(llmMessages)
142-
const contentStarted = hasContentStarted(llmMessages)
143-
const reachedReferences = hasReachedReferences(llmMessages)
139+
// Get events sorted by timestamp (most recent last)
140+
const statusEvents = events
141+
.filter(
142+
(m) =>
143+
m.type === EventTypes.REASONING ||
144+
m.type === EventTypes.SEARCH_TOOL_CALL ||
145+
m.type === EventTypes.TOOL_CALL ||
146+
m.type === EventTypes.TOOL_RESULT ||
147+
m.type === EventTypes.CHUNK
148+
)
149+
.sort((a, b) => a.timestamp - b.timestamp)
144150

145-
if (reachedReferences) {
146-
return 'Gathering resources'
147-
} else if (contentStarted) {
148-
return 'Generating'
149-
} else if (searchQuery) {
150-
return `Searching for "${searchQuery}"`
151-
}
151+
// Get the most recent status-worthy event
152+
const latestEvent = statusEvents[statusEvents.length - 1]
153+
154+
if (!latestEvent) return STATUS_MESSAGES.THINKING
155+
156+
switch (latestEvent.type) {
157+
case EventTypes.REASONING:
158+
return latestEvent.message || STATUS_MESSAGES.THINKING
152159

153-
return 'Thinking'
160+
case EventTypes.SEARCH_TOOL_CALL:
161+
return `Searching Elastic's Docs for "${latestEvent.searchQuery}"`
162+
163+
case EventTypes.TOOL_CALL:
164+
return getToolCallStatus(latestEvent)
165+
166+
case EventTypes.TOOL_RESULT:
167+
return STATUS_MESSAGES.ANALYZING
168+
169+
case EventTypes.CHUNK: {
170+
const allContent = events
171+
.filter((m) => m.type === EventTypes.CHUNK)
172+
.map((m) => m.content)
173+
.join('')
174+
175+
if (allContent.includes('<!--REFERENCES')) {
176+
return STATUS_MESSAGES.GATHERING
177+
}
178+
return STATUS_MESSAGES.GENERATING
179+
}
180+
181+
default:
182+
return STATUS_MESSAGES.THINKING
183+
}
154184
}
155185

156186
// Action bar for complete AI messages
@@ -215,7 +245,7 @@ const ActionBar = ({
215245

216246
export const ChatMessage = ({
217247
message,
218-
llmMessages = [],
248+
events = [],
219249
streamingContent,
220250
error,
221251
onRetry,
@@ -251,9 +281,7 @@ export const ChatMessage = ({
251281

252282
const content =
253283
streamingContent ||
254-
(llmMessages.length > 0
255-
? getAccumulatedContent(llmMessages)
256-
: message.content)
284+
(events.length > 0 ? getAccumulatedContent(events) : message.content)
257285

258286
const hasError = message.status === 'error' || !!error
259287

@@ -279,11 +307,14 @@ export const ChatMessage = ({
279307
return DOMPurify.sanitize(html)
280308
}, [mainContent])
281309

282-
const aiStatus = useMemo(
283-
() => computeAiStatus(llmMessages, isComplete),
284-
[llmMessages, isComplete]
310+
const rawAiStatus = useMemo(
311+
() => computeAiStatus(events, isComplete),
312+
[events, isComplete]
285313
)
286314

315+
// Apply minimum display time to prevent status flickering
316+
const aiStatus = useStatusMinDisplay(rawAiStatus)
317+
287318
const ref = React.useRef<HTMLDivElement>(null)
288319

289320
useEffect(() => {

0 commit comments

Comments
 (0)