Skip to content

Commit dacf991

Browse files
Backlog/v12 socai assistant (#2506)
* feat[frontend](assistant): added chat history on assistant chats * fix[plugins](socai): added tool deduplication, chat history and improved context window compaction * fix[backend](socai): added chat history on assitant requests
1 parent 936df10 commit dacf991

6 files changed

Lines changed: 176 additions & 30 deletions

File tree

backend/modules/socai/handler/chat.go

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -67,10 +67,16 @@ func NewChatHandler(client socAIStreamer) *ChatHandler {
6767
return &ChatHandler{client: client}
6868
}
6969

70+
type chatTurn struct {
71+
Role string `json:"role"`
72+
Content string `json:"content"`
73+
}
74+
7075
type chatRequest struct {
71-
Task string `json:"task" binding:"required"`
72-
Page string `json:"page"`
73-
Lang string `json:"lang"`
76+
Task string `json:"task" binding:"required"`
77+
Page string `json:"page"`
78+
Lang string `json:"lang"`
79+
History []chatTurn `json:"history,omitempty"`
7480
}
7581

7682
// Chat godoc
@@ -101,7 +107,7 @@ func (h *ChatHandler) Chat(c *gin.Context) {
101107
return
102108
}
103109

104-
body, err := json.Marshal(map[string]string{"task": req.Task, "page": req.Page, "lang": req.Lang})
110+
body, err := json.Marshal(req)
105111
if err != nil {
106112
c.JSON(http.StatusInternalServerError, gin.H{"status": "error", "message": err.Error()})
107113
return

frontend/src/features/soc-ai/SocAiProvider.tsx

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,11 @@
11
import { createContext, useCallback, useContext, useMemo, useRef, useState, type Dispatch, type ReactNode, type SetStateAction } from 'react'
22
import { useLocation } from 'react-router-dom'
33
import { useTranslation } from 'react-i18next'
4-
import { extractNavigation, streamChat, type NavAction } from './lib/chat-stream'
4+
import { extractNavigation, streamChat, type ChatHistoryTurn, type NavAction } from './lib/chat-stream'
5+
6+
// How many prior text turns to replay to the backend as chat memory. Server-side
7+
// compaction will still trim if this exceeds the model context window.
8+
const HISTORY_LIMIT = 10
59

610
export interface ToolStep {
711
tool: string
@@ -114,8 +118,13 @@ export function SocAiProvider({ children }: { children: ReactNode }) {
114118
const page = pageContext(location.pathname)
115119
const lang = (i18n.language || 'en').split('-')[0]
116120

121+
const history: ChatHistoryTurn[] = current
122+
.filter((m) => m.text && !m.error && !m.pending)
123+
.slice(-HISTORY_LIMIT)
124+
.map((m) => ({ role: m.role === 'user' ? 'user' : 'assistant', content: m.text }))
125+
117126
streamChat(
118-
{ task: text, page, lang },
127+
{ task: text, page, lang, history },
119128
(ev) => {
120129
patchMsg(scope, aiId, (msg) => {
121130
switch (ev.kind) {

frontend/src/features/soc-ai/lib/chat-stream.ts

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,13 +32,21 @@ export interface NavAction {
3232
time?: string
3333
}
3434

35+
/** A single prior chat turn replayed to the backend as context. Only user and
36+
* assistant text turns are forwarded — tool_use/tool_result blocks are internal
37+
* to a single Run() on the server and must not be replayed. */
38+
export interface ChatHistoryTurn {
39+
role: 'user' | 'assistant'
40+
content: string
41+
}
42+
3543
/**
3644
* Streams the SOC-AI chat agent over SSE. The backend (/soc-ai/chat) proxies the
3745
* plugin's agent and emits tool_call / tool_result / final / error events. Uses
3846
* fetch + ReadableStream because the shared axios client can't stream.
3947
*/
4048
export async function streamChat(
41-
body: { task: string; page?: string; lang?: string },
49+
body: { task: string; page?: string; lang?: string; history?: ChatHistoryTurn[] },
4250
onEvent: (e: ChatEvent) => void,
4351
signal?: AbortSignal,
4452
): Promise<void> {

plugins/soc-ai/internal/agent/loop.go

Lines changed: 107 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,8 @@ const (
1414
defaultMaxIters = 12
1515
compactionThreshold = 0.80
1616
summaryMaxTokens = 400 // ~200 words + slack
17+
keepTailMessages = 4 // messages kept raw when compacting
18+
genericErrorMsg = "An error has occurred while processing your request."
1719
)
1820

1921
var modelContextWindow = []struct {
@@ -69,6 +71,7 @@ func (s EventSink) emit(e Event) {
6971
type RunTask struct {
7072
System string // system prompt
7173
Input string // the user turn (alert JSON for triage, free task for ops)
74+
History []Message
7275
EnabledGroups []string
7376
AlwaysAllow []string
7477
MaxIters int
@@ -103,7 +106,10 @@ func (a *Agent) Broker() *ToolBroker { return a.broker }
103106
func (a *Agent) Run(ctx context.Context, task RunTask, sink EventSink) (RunResult, error) {
104107
specs, err := a.broker.ListSpecs(ctx)
105108
if err != nil {
106-
sink.emit(Event{Kind: EventError, Text: "could not load tools: " + err.Error()})
109+
_ = catcher.Error("could not load tools", err, map[string]any{
110+
"process": "plugin_com.utmstack.soc-ai",
111+
})
112+
sink.emit(Event{Kind: EventError, Text: genericErrorMsg})
107113
return RunResult{}, fmt.Errorf("list tools: %w", err)
108114
}
109115
allowed := filterTools(specs, task)
@@ -117,14 +123,19 @@ func (a *Agent) Run(ctx context.Context, task RunTask, sink EventSink) (RunResul
117123
maxIters = defaultMaxIters
118124
}
119125

120-
msgs := []Message{{Role: RoleUser, Content: task.Input}}
126+
msgs := append([]Message{}, task.History...)
127+
msgs = append(msgs, Message{Role: RoleUser, Content: task.Input})
121128
result := RunResult{}
122129

130+
// Skipped in-batch dedup; upgrade to singleflight if same-batch duplicates become measurable.
131+
toolCache := map[string]tcOut{}
132+
var cacheMu sync.Mutex
133+
123134
for step := 1; step <= maxIters; step++ {
124135
result.Steps = step
125136

126137
if a.contextWindow > 0 && len(msgs) > 1 &&
127-
estimateTokens(task.System, msgs) >= int(compactionThreshold*float64(a.contextWindow)) {
138+
estimateTokens(task.System, msgs, allowed) >= int(compactionThreshold*float64(a.contextWindow)) {
128139
newMsgs, cErr := a.compact(ctx, task.Input, msgs)
129140
if cErr != nil {
130141
_ = catcher.Error("context compaction failed, continuing with full history", cErr, map[string]any{
@@ -144,7 +155,10 @@ func (a *Agent) Run(ctx context.Context, task RunTask, sink EventSink) (RunResul
144155
MaxTokens: a.maxTokens,
145156
})
146157
if err != nil {
147-
sink.emit(Event{Kind: EventError, Text: err.Error()})
158+
_ = catcher.Error("llm completion failed", err, map[string]any{
159+
"process": "plugin_com.utmstack.soc-ai",
160+
})
161+
sink.emit(Event{Kind: EventError, Text: genericErrorMsg})
148162
return result, err
149163
}
150164

@@ -156,33 +170,80 @@ func (a *Agent) Run(ctx context.Context, task RunTask, sink EventSink) (RunResul
156170

157171
msgs = append(msgs, Message{Role: RoleAssistant, Content: resp.Content, ToolCalls: resp.ToolCalls})
158172

159-
for _, tc := range resp.ToolCalls {
173+
outs := make([]tcOut, len(resp.ToolCalls))
174+
var wg sync.WaitGroup
175+
for i, tc := range resp.ToolCalls {
160176
result.ToolCalls++
161177
sink.emit(Event{Kind: EventToolCall, Step: step, Tool: tc.Name, Args: tc.Args})
162178

163179
if !allowedSet[tc.Name] {
164-
const msg = "tool not permitted in this mode"
165-
msgs = append(msgs, Message{Role: RoleTool, ToolResult: &ToolResult{ID: tc.ID, Name: tc.Name, Content: msg, IsError: true}})
166-
sink.emit(Event{Kind: EventToolResult, Step: step, Tool: tc.Name, Output: msg, IsError: true})
180+
outs[i] = tcOut{out: "tool not permitted in this mode", isErr: true}
167181
continue
168182
}
169183

170-
out, isErr, callErr := a.broker.Call(ctx, tc.Name, tc.Args)
171-
if callErr != nil {
172-
out = callErr.Error()
173-
isErr = true
184+
key := tc.Name + "|" + string(tc.Args)
185+
cacheMu.Lock()
186+
cached, ok := toolCache[key]
187+
cacheMu.Unlock()
188+
if ok {
189+
outs[i] = cached
190+
continue
174191
}
175-
msgs = append(msgs, Message{Role: RoleTool, ToolResult: &ToolResult{ID: tc.ID, Name: tc.Name, Content: out, IsError: isErr}})
176-
sink.emit(Event{Kind: EventToolResult, Step: step, Tool: tc.Name, Output: out, IsError: isErr})
192+
193+
wg.Add(1)
194+
go func(i int, tc ToolCall, key string) {
195+
defer wg.Done()
196+
out, isErr, callErr := a.broker.Call(ctx, tc.Name, tc.Args)
197+
if callErr != nil {
198+
out = callErr.Error()
199+
isErr = true
200+
}
201+
r := tcOut{out: out, isErr: isErr}
202+
outs[i] = r
203+
cacheMu.Lock()
204+
toolCache[key] = r
205+
cacheMu.Unlock()
206+
}(i, tc, key)
207+
}
208+
wg.Wait()
209+
210+
for i, tc := range resp.ToolCalls {
211+
r := outs[i]
212+
msgs = append(msgs, Message{Role: RoleTool, ToolResult: &ToolResult{ID: tc.ID, Name: tc.Name, Content: r.out, IsError: r.isErr}})
213+
sink.emit(Event{Kind: EventToolResult, Step: step, Tool: tc.Name, Output: r.out, IsError: r.isErr})
177214
}
178215
}
179216

180-
const exhausted = "Reached the maximum number of tool iterations before finishing."
181-
sink.emit(Event{Kind: EventFinal, Text: exhausted})
182-
result.Final = exhausted
217+
// Loop exhausted: give the model one last chance to finalize with no tools.
218+
msgs = append(msgs, Message{
219+
Role: RoleUser,
220+
Content: "You have reached the maximum number of tool iterations. Do not call any more tools. Provide your final assessment now based on what you have gathered so far.",
221+
})
222+
finalResp, ferr := a.llm.Complete(ctx, CompletionRequest{
223+
System: task.System,
224+
Messages: msgs,
225+
Model: a.model,
226+
MaxTokens: a.maxTokens,
227+
})
228+
if ferr != nil {
229+
_ = catcher.Error("max-iters finalization llm call failed", ferr, map[string]any{
230+
"process": "plugin_com.utmstack.soc-ai",
231+
})
232+
const msg = "Reached the maximum number of tool iterations and could not finalize."
233+
sink.emit(Event{Kind: EventFinal, Text: msg})
234+
result.Final = msg
235+
return result, nil
236+
}
237+
sink.emit(Event{Kind: EventFinal, Text: finalResp.Content})
238+
result.Final = finalResp.Content
183239
return result, nil
184240
}
185241

242+
type tcOut struct {
243+
out string
244+
isErr bool
245+
}
246+
186247
func filterTools(specs []ToolSpec, task RunTask) []ToolSpec {
187248
enabled := make(map[string]bool, len(task.EnabledGroups))
188249
for _, g := range task.EnabledGroups {
@@ -205,8 +266,16 @@ func filterTools(specs []ToolSpec, task RunTask) []ToolSpec {
205266
return out
206267
}
207268

208-
func estimateTokens(system string, msgs []Message) int {
269+
func estimateTokens(system string, msgs []Message, tools []ToolSpec) int {
209270
n := len(system)
271+
for _, t := range tools {
272+
n += len(t.Name) + len(t.Description)
273+
if t.InputSchema != nil {
274+
if b, err := json.Marshal(t.InputSchema); err == nil {
275+
n += len(b)
276+
}
277+
}
278+
}
210279
for _, m := range msgs {
211280
n += len(m.Content)
212281
for _, tc := range m.ToolCalls {
@@ -220,8 +289,24 @@ func estimateTokens(system string, msgs []Message) int {
220289
}
221290

222291
func (a *Agent) compact(ctx context.Context, userInput string, msgs []Message) ([]Message, error) {
292+
// Keep the last keepTailMessages raw. Advance the cut point forward past any
293+
// tool messages so the preserved tail never starts with an orphan tool_result
294+
// (which would reference an assistant tool_call left in the summarized head).
295+
cut := len(msgs) - keepTailMessages
296+
if cut < 1 {
297+
cut = 1
298+
}
299+
for cut < len(msgs) && msgs[cut].Role == RoleTool {
300+
cut++
301+
}
302+
head := msgs[:cut]
303+
var tail []Message
304+
if cut < len(msgs) {
305+
tail = msgs[cut:]
306+
}
307+
223308
var b strings.Builder
224-
for _, m := range msgs {
309+
for _, m := range head {
225310
fmt.Fprintf(&b, "[%s] %s\n", m.Role, m.Content)
226311
for _, tc := range m.ToolCalls {
227312
fmt.Fprintf(&b, " tool_call %s(%s)\n", tc.Name, string(tc.Args))
@@ -242,10 +327,11 @@ func (a *Agent) compact(ctx context.Context, userInput string, msgs []Message) (
242327
if strings.TrimSpace(resp.Content) == "" {
243328
return msgs, fmt.Errorf("empty summary")
244329
}
245-
return []Message{{
330+
summary := Message{
246331
Role: RoleUser,
247332
Content: "Original task:\n" + userInput + "\n\nProgress so far (summary of prior context):\n" + resp.Content + "\n\nContinue the task.",
248-
}}, nil
333+
}
334+
return append([]Message{summary}, tail...), nil
249335
}
250336

251337
type registry struct {

plugins/soc-ai/internal/agent/prompt.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@ func OpsPrompt(page, lang string, enabledGroups []string) string {
3737
}
3838
langLine := "Answer in the same language as the user's message."
3939
if strings.TrimSpace(lang) != "" {
40-
langLine = `Always write your reply in the user's interface language, identified by the code "` + strings.TrimSpace(lang) + `" (e.g. es=Spanish, pt=Portuguese, en=English), regardless of the language of their message.`
40+
langLine = `Always write your reply in the user's interface language, identified by the ISO code "` + strings.TrimSpace(lang) + `", regardless of the language of their message.`
4141
}
4242
return `You are the UTMStack operations agent — an autonomous SOC assistant embedded in the UTMStack SIEM. The user chats with you, and you operate the SIEM on their behalf through the available tools (alerts, incidents, log/alert search, SOAR response actions, datasources, compliance, and more).
4343
@@ -52,7 +52,7 @@ Use this to choose the most relevant tools and to craft navigation. For example,
5252
` + permissionsBlock(enabledGroups) + `
5353
5454
## How to work
55-
- Plan briefly, then act. Carry the task end to end.
55+
- Carry the task end to end.
5656
- Use tools ONLY when you need data or actions you don't already have. Many messages need few or no tools — do not over-call; prefer the smallest set of tools that answers the question.
5757
- Prefer read-only tools to investigate before any mutating or response action. Mutating/response actions (changing status, creating incidents, running SOAR jobs, etc.) take effect immediately — only perform them when the task clearly asks for them.
5858
- Never invent data; rely on tool results. If a tool fails, adapt or report it plainly.

plugins/soc-ai/internal/api/server.go

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,16 @@ type AgentTaskRequest struct {
2222
// Lang is the user's interface language code (en/es/pt/…) so the agent
2323
// replies in that language regardless of the message language.
2424
Lang string `json:"lang"`
25+
// History is prior chat turns from the client (text only). Only user and
26+
// assistant roles are accepted; tool_use/tool_result turns are internal to
27+
// a single Run() and must not be replayed.
28+
History []AgentTurn `json:"history,omitempty"`
29+
}
30+
31+
// AgentTurn is a single prior chat message forwarded by the client.
32+
type AgentTurn struct {
33+
Role string `json:"role"` // "user" | "assistant"
34+
Content string `json:"content"`
2535
}
2636

2737
// AnalyzeRequest represents the request body for manual alert analysis
@@ -214,11 +224,38 @@ func handleAgentTask(w http.ResponseWriter, r *http.Request) {
214224
_, _ = ag.Run(r.Context(), agent.RunTask{
215225
System: agent.OpsPrompt(req.Page, req.Lang, capabilities),
216226
Input: req.Task,
227+
History: toHistory(req.History),
217228
EnabledGroups: capabilities,
218229
MaxIters: maxIters,
219230
}, sink)
220231
}
221232

233+
// toHistory converts client-supplied turns into agent.Message. Unknown roles
234+
// and empty content are dropped so a malformed client can't inject tool turns
235+
// or blank rows.
236+
func toHistory(turns []AgentTurn) []agent.Message {
237+
if len(turns) == 0 {
238+
return nil
239+
}
240+
out := make([]agent.Message, 0, len(turns))
241+
for _, t := range turns {
242+
if t.Content == "" {
243+
continue
244+
}
245+
var role agent.Role
246+
switch t.Role {
247+
case "user":
248+
role = agent.RoleUser
249+
case "assistant":
250+
role = agent.RoleAssistant
251+
default:
252+
continue
253+
}
254+
out = append(out, agent.Message{Role: role, Content: t.Content})
255+
}
256+
return out
257+
}
258+
222259
func writeJSONError(w http.ResponseWriter, status int, msg string) {
223260
w.Header().Set("Content-Type", "application/json")
224261
w.WriteHeader(status)

0 commit comments

Comments
 (0)