mirror of
https://github.com/rcourtman/Pulse.git
synced 2026-09-10 02:25:56 +00:00
feat(ai): add knowledge accumulation and enhance agentic loop
- Introduce KnowledgeAccumulator to persist facts across turns - Enhance AgenticLoop to support knowledge injection and final text summaries - Update chat service to wire up knowledge components - Frontend updates to support enhanced chat capabilities
This commit is contained in:
@@ -96,6 +96,18 @@ export function useChat(options: UseChatOptions = {}) {
|
||||
return '';
|
||||
};
|
||||
|
||||
const extractTokens = (data: unknown): { input: number; output: number } | null => {
|
||||
if (!data || typeof data !== 'object') return null;
|
||||
const record = data as Record<string, unknown>;
|
||||
const input = Number(record.input_tokens ?? record.inputTokens ?? record.input);
|
||||
const output = Number(record.output_tokens ?? record.outputTokens ?? record.output);
|
||||
if (!Number.isFinite(input) && !Number.isFinite(output)) return null;
|
||||
return {
|
||||
input: Number.isFinite(input) && input > 0 ? input : 0,
|
||||
output: Number.isFinite(output) && output > 0 ? output : 0,
|
||||
};
|
||||
};
|
||||
|
||||
const processEvent = (
|
||||
assistantId: string,
|
||||
event: StreamEvent
|
||||
@@ -281,6 +293,10 @@ export function useChat(options: UseChatOptions = {}) {
|
||||
}
|
||||
|
||||
case 'done': {
|
||||
const tokens = extractTokens(event.data);
|
||||
if (tokens && (tokens.input > 0 || tokens.output > 0)) {
|
||||
return { ...msg, isStreaming: false, pendingTools: [], tokens };
|
||||
}
|
||||
return { ...msg, isStreaming: false, pendingTools: [] };
|
||||
}
|
||||
|
||||
|
||||
@@ -104,6 +104,7 @@ export const AIChat: Component<AIChatProps> = (props) => {
|
||||
// Chat hook
|
||||
const chat = useChat({ model: initialModelSelections[DEFAULT_SESSION_KEY] || '' });
|
||||
|
||||
|
||||
const defaultModelLabel = createMemo(() => {
|
||||
const fallback = defaultModel().trim();
|
||||
if (!fallback) return '';
|
||||
|
||||
@@ -131,3 +131,9 @@ export interface StreamCompleteData {
|
||||
output_tokens: number;
|
||||
tool_calls?: ToolExecution[];
|
||||
}
|
||||
|
||||
export interface StreamDoneData {
|
||||
session_id?: string;
|
||||
input_tokens?: number;
|
||||
output_tokens?: number;
|
||||
}
|
||||
|
||||
+220
-15
@@ -40,6 +40,9 @@ type AgenticLoop struct {
|
||||
// Per-session FSMs for workflow enforcement (set before each execution)
|
||||
sessionFSM *SessionFSM
|
||||
|
||||
// Knowledge accumulator for fact extraction across turns
|
||||
knowledgeAccumulator *KnowledgeAccumulator
|
||||
|
||||
// Budget checker called after each turn to enforce token spending limits
|
||||
budgetChecker func() error
|
||||
}
|
||||
@@ -78,8 +81,10 @@ func (a *AgenticLoop) ExecuteWithTools(ctx context.Context, sessionID string, me
|
||||
}
|
||||
|
||||
func (a *AgenticLoop) executeWithTools(ctx context.Context, sessionID string, messages []Message, tools []providers.Tool, callback StreamCallback) ([]Message, error) {
|
||||
// Track this session for potential abort
|
||||
// Snapshot maxTurns under the lock — callers may override via SetMaxTurns
|
||||
// before calling ExecuteWithTools, and this avoids races with concurrent sessions.
|
||||
a.mu.Lock()
|
||||
maxTurns := a.maxTurns
|
||||
a.aborted[sessionID] = false
|
||||
a.mu.Unlock()
|
||||
defer func() {
|
||||
@@ -113,11 +118,11 @@ func (a *AgenticLoop) executeWithTools(ctx context.Context, sessionID string, me
|
||||
|
||||
// Track where each turn's messages begin in providerMessages for compaction.
|
||||
// We keep the last N turns' tool results in full; older ones get compacted.
|
||||
const compactionKeepTurns = 3 // Keep last 3 turns' tool results in full
|
||||
const compactionMinChars = 500 // Only compact results longer than this
|
||||
const compactionKeepTurns = 2 // Keep last 2 turns' tool results in full (KA preserves key facts)
|
||||
const compactionMinChars = 300 // Only compact results longer than this
|
||||
currentTurnStartIndex := len(providerMessages) // Initial messages are never compacted
|
||||
|
||||
for turn < a.maxTurns {
|
||||
for turn < maxTurns {
|
||||
// === CONTEXT COMPACTION: Compact old tool results to prevent context blowout ===
|
||||
if turn > 0 {
|
||||
compactOldToolResults(providerMessages, currentTurnStartIndex, compactionKeepTurns, compactionMinChars)
|
||||
@@ -168,11 +173,23 @@ func (a *AgenticLoop) executeWithTools(ctx context.Context, sessionID string, me
|
||||
// 2. It's the first turn
|
||||
// 3. The user's message indicates they need live data or an action
|
||||
// This prevents forcing tool calls on conceptual questions like "What is TCP?"
|
||||
if writeCompletedLastTurn {
|
||||
forcedTextOnly := false
|
||||
if turn >= maxTurns-1 {
|
||||
// Last turn before hitting the limit — force a text-only response so
|
||||
// the model summarizes its findings instead of silently stopping.
|
||||
req.ToolChoice = &providers.ToolChoice{Type: providers.ToolChoiceNone}
|
||||
forcedTextOnly = true
|
||||
log.Warn().
|
||||
Int("turn", turn).
|
||||
Int("max_turns", maxTurns).
|
||||
Str("session_id", sessionID).
|
||||
Msg("[AgenticLoop] Approaching max turns — forcing text-only response for summary")
|
||||
} else if writeCompletedLastTurn {
|
||||
// A write action completed successfully on the previous turn.
|
||||
// Force text-only response so the model summarizes the result instead of
|
||||
// making more tool calls (which often return stale cached data and cause loops).
|
||||
req.ToolChoice = &providers.ToolChoice{Type: providers.ToolChoiceNone}
|
||||
forcedTextOnly = true
|
||||
writeCompletedLastTurn = false
|
||||
log.Debug().
|
||||
Str("session_id", sessionID).
|
||||
@@ -182,6 +199,7 @@ func (a *AgenticLoop) executeWithTools(ctx context.Context, sessionID string, me
|
||||
// The model already has the data it gathered — force it to produce a text
|
||||
// response instead of continuing to call tools that will just be blocked again.
|
||||
req.ToolChoice = &providers.ToolChoice{Type: providers.ToolChoiceNone}
|
||||
forcedTextOnly = true
|
||||
toolBlockedLastTurn = false
|
||||
log.Debug().
|
||||
Str("session_id", sessionID).
|
||||
@@ -248,7 +266,11 @@ func (a *AgenticLoop) executeWithTools(ctx context.Context, sessionID string, me
|
||||
// Format input for frontend display
|
||||
// For control tools, show a human-readable summary instead of raw JSON to avoid "hallucination" look
|
||||
inputStr := "{}"
|
||||
rawInput := ""
|
||||
if data.Input != nil {
|
||||
if inputBytes, err := json.Marshal(data.Input); err == nil {
|
||||
rawInput = string(inputBytes)
|
||||
}
|
||||
// Special handling for command execution tools to avoid showing raw JSON
|
||||
if data.Name == "pulse_control" || data.Name == "pulse_run_command" || data.Name == "control" {
|
||||
if cmd, ok := data.Input["command"].(string); ok {
|
||||
@@ -277,9 +299,10 @@ func (a *AgenticLoop) executeWithTools(ctx context.Context, sessionID string, me
|
||||
}
|
||||
}
|
||||
jsonData, _ := json.Marshal(ToolStartData{
|
||||
ID: data.ID,
|
||||
Name: data.Name,
|
||||
Input: inputStr,
|
||||
ID: data.ID,
|
||||
Name: data.Name,
|
||||
Input: inputStr,
|
||||
RawInput: rawInput,
|
||||
})
|
||||
callback(StreamEvent{Type: "tool_start", Data: jsonData})
|
||||
}
|
||||
@@ -323,6 +346,18 @@ func (a *AgenticLoop) executeWithTools(ctx context.Context, sessionID string, me
|
||||
return resultMessages, fmt.Errorf("provider error: %w", err)
|
||||
}
|
||||
|
||||
// Guard: if we forced text-only but the model still returned tool calls
|
||||
// (some providers like Gemini can hallucinate function calls from conversation
|
||||
// history even when tools are not offered in the request), strip them so the
|
||||
// model's text content is treated as the final response.
|
||||
if forcedTextOnly && len(toolCalls) > 0 {
|
||||
log.Warn().
|
||||
Str("session_id", sessionID).
|
||||
Int("stripped_tool_calls", len(toolCalls)).
|
||||
Msg("[AgenticLoop] Model returned tool calls despite ToolChoiceNone — stripping them")
|
||||
toolCalls = nil
|
||||
}
|
||||
|
||||
// Check mid-run budget after each turn completes
|
||||
if a.budgetChecker != nil {
|
||||
if budgetErr := a.budgetChecker(); budgetErr != nil {
|
||||
@@ -475,6 +510,7 @@ func (a *AgenticLoop) executeWithTools(ctx context.Context, sessionID string, me
|
||||
}
|
||||
|
||||
log.Debug().Msg("Agentic loop complete - no tool calls")
|
||||
resultMessages = a.ensureFinalTextResponse(ctx, sessionID, resultMessages, providerMessages, callback)
|
||||
return resultMessages, nil
|
||||
}
|
||||
|
||||
@@ -611,6 +647,58 @@ func (a *AgenticLoop) executeWithTools(ctx context.Context, sessionID string, me
|
||||
continue
|
||||
}
|
||||
|
||||
// === KNOWLEDGE GATE: Return cached facts for redundant tool calls ===
|
||||
if a.knowledgeAccumulator != nil {
|
||||
if keys := PredictFactKeys(tc.Name, tc.Input); len(keys) > 0 {
|
||||
var cachedParts []string
|
||||
for _, key := range keys {
|
||||
if value, found := a.knowledgeAccumulator.Lookup(key); found {
|
||||
cachedParts = append(cachedParts, value)
|
||||
}
|
||||
}
|
||||
if len(cachedParts) > 0 {
|
||||
cachedResult := fmt.Sprintf("Already known (from earlier investigation, matched keys: %s): %s. If you need fresh data, use a different query or approach.", strings.Join(keys, ","), strings.Join(cachedParts, "; "))
|
||||
|
||||
log.Info().
|
||||
Str("tool", tc.Name).
|
||||
Str("session_id", sessionID).
|
||||
Strs("matched_keys", keys).
|
||||
Int("cached_parts", len(cachedParts)).
|
||||
Msg("[AgenticLoop] Knowledge gate: returning cached fact instead of re-executing tool")
|
||||
|
||||
jsonData, _ := json.Marshal(ToolEndData{
|
||||
ID: tc.ID,
|
||||
Name: tc.Name,
|
||||
Input: "",
|
||||
Output: cachedResult,
|
||||
Success: true,
|
||||
})
|
||||
callback(StreamEvent{Type: "tool_end", Data: jsonData})
|
||||
|
||||
toolResultMsg := Message{
|
||||
ID: uuid.New().String(),
|
||||
Role: "user",
|
||||
Timestamp: time.Now(),
|
||||
ToolResult: &ToolResult{
|
||||
ToolUseID: tc.ID,
|
||||
Content: cachedResult,
|
||||
IsError: false,
|
||||
},
|
||||
}
|
||||
resultMessages = append(resultMessages, toolResultMsg)
|
||||
providerMessages = append(providerMessages, providers.Message{
|
||||
Role: "user",
|
||||
ToolResult: &providers.ToolResult{
|
||||
ToolUseID: tc.ID,
|
||||
Content: cachedResult,
|
||||
IsError: false,
|
||||
},
|
||||
})
|
||||
continue
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Execute the tool
|
||||
result, err := a.executor.ExecuteTool(ctx, tc.Name, tc.Input)
|
||||
|
||||
@@ -628,6 +716,20 @@ func (a *AgenticLoop) executeWithTools(ctx context.Context, sessionID string, me
|
||||
log.Debug().
|
||||
Str("tool", tc.Name).
|
||||
Msg("[AgenticLoop] Tool succeeded - toolsSucceededThisEpisode set to true")
|
||||
|
||||
// Extract and accumulate knowledge facts
|
||||
if a.knowledgeAccumulator != nil {
|
||||
a.knowledgeAccumulator.SetTurn(turn)
|
||||
facts := ExtractFacts(tc.Name, tc.Input, resultText)
|
||||
for _, f := range facts {
|
||||
a.knowledgeAccumulator.AddFact(f.Category, f.Key, f.Value)
|
||||
log.Debug().
|
||||
Str("tool", tc.Name).
|
||||
Str("fact_key", f.Key).
|
||||
Int("value_len", len(f.Value)).
|
||||
Msg("[AgenticLoop] Stored knowledge fact")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -786,7 +888,11 @@ func (a *AgenticLoop) executeWithTools(ctx context.Context, sessionID string, me
|
||||
// Send tool_end event
|
||||
// Convert input to JSON string for frontend display
|
||||
inputStr := ""
|
||||
rawInput := ""
|
||||
if tc.Input != nil {
|
||||
if inputBytes, err := json.Marshal(tc.Input); err == nil {
|
||||
rawInput = string(inputBytes)
|
||||
}
|
||||
// Special handling for command execution tools to avoid showing raw JSON
|
||||
if tc.Name == "pulse_control" || tc.Name == "pulse_run_command" || tc.Name == "control" {
|
||||
if cmd, ok := tc.Input["command"].(string); ok {
|
||||
@@ -814,11 +920,12 @@ func (a *AgenticLoop) executeWithTools(ctx context.Context, sessionID string, me
|
||||
}
|
||||
}
|
||||
jsonData, _ := json.Marshal(ToolEndData{
|
||||
ID: tc.ID,
|
||||
Name: tc.Name,
|
||||
Input: inputStr,
|
||||
Output: resultText,
|
||||
Success: !isError,
|
||||
ID: tc.ID,
|
||||
Name: tc.Name,
|
||||
Input: inputStr,
|
||||
RawInput: rawInput,
|
||||
Output: resultText,
|
||||
Success: !isError,
|
||||
})
|
||||
callback(StreamEvent{Type: "tool_end", Data: jsonData})
|
||||
|
||||
@@ -951,10 +1058,93 @@ func (a *AgenticLoop) executeWithTools(ctx context.Context, sessionID string, me
|
||||
turn++
|
||||
}
|
||||
|
||||
log.Warn().Int("max_turns", a.maxTurns).Msg("Agentic loop hit max turns limit")
|
||||
log.Warn().Int("max_turns", maxTurns).Str("session_id", sessionID).Msg("Agentic loop hit max turns limit")
|
||||
resultMessages = a.ensureFinalTextResponse(ctx, sessionID, resultMessages, providerMessages, callback)
|
||||
return resultMessages, nil
|
||||
}
|
||||
|
||||
// ensureFinalTextResponse checks if the result messages contain any assistant text.
|
||||
// If not, it makes one last text-only LLM call to force the model to summarize its findings.
|
||||
// This prevents the loop from exiting silently after making tool calls without answering.
|
||||
func (a *AgenticLoop) ensureFinalTextResponse(
|
||||
ctx context.Context,
|
||||
sessionID string,
|
||||
resultMessages []Message,
|
||||
providerMessages []providers.Message,
|
||||
callback StreamCallback,
|
||||
) []Message {
|
||||
// Check if any assistant message has text content
|
||||
for i := len(resultMessages) - 1; i >= 0; i-- {
|
||||
if resultMessages[i].Role == "assistant" && strings.TrimSpace(resultMessages[i].Content) != "" {
|
||||
return resultMessages // Already has text — nothing to do
|
||||
}
|
||||
}
|
||||
|
||||
// No text content from the model. Make a final text-only call.
|
||||
log.Warn().Str("session_id", sessionID).Msg("[AgenticLoop] No text content produced — making final summary call")
|
||||
|
||||
// Build clean message history for the summary call:
|
||||
// 1. Strip any trailing empty assistant messages (the model already failed to produce
|
||||
// text with these, so including them would just get the same empty result).
|
||||
// 2. Append a user-role nudge to give the model a clear instruction.
|
||||
cleanMessages := make([]providers.Message, len(providerMessages))
|
||||
copy(cleanMessages, providerMessages)
|
||||
for len(cleanMessages) > 0 {
|
||||
last := cleanMessages[len(cleanMessages)-1]
|
||||
if last.Role == "assistant" && strings.TrimSpace(last.Content) == "" && len(last.ToolCalls) == 0 {
|
||||
cleanMessages = cleanMessages[:len(cleanMessages)-1]
|
||||
} else {
|
||||
break
|
||||
}
|
||||
}
|
||||
cleanMessages = append(cleanMessages, providers.Message{
|
||||
Role: "user",
|
||||
Content: "You've gathered information above using tools. Now provide your analysis and answer to the original question. Summarize your findings concisely.",
|
||||
})
|
||||
|
||||
summaryReq := providers.ChatRequest{
|
||||
Messages: cleanMessages,
|
||||
System: a.getSystemPrompt(),
|
||||
ToolChoice: &providers.ToolChoice{Type: providers.ToolChoiceNone},
|
||||
// No Tools field — completely omit tools to prevent hallucinated function calls
|
||||
}
|
||||
|
||||
var summaryBuilder strings.Builder
|
||||
|
||||
summaryErr := a.provider.ChatStream(ctx, summaryReq, func(event providers.StreamEvent) {
|
||||
switch event.Type {
|
||||
case "content":
|
||||
if data, ok := event.Data.(providers.ContentEvent); ok {
|
||||
summaryBuilder.WriteString(data.Text)
|
||||
jsonData, _ := json.Marshal(ContentData{Text: data.Text})
|
||||
callback(StreamEvent{Type: "content", Data: jsonData})
|
||||
}
|
||||
case "done":
|
||||
if data, ok := event.Data.(providers.DoneEvent); ok {
|
||||
a.totalInputTokens += data.InputTokens
|
||||
a.totalOutputTokens += data.OutputTokens
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
if summaryErr == nil && summaryBuilder.Len() > 0 {
|
||||
summaryMsg := Message{
|
||||
ID: uuid.New().String(),
|
||||
Role: "assistant",
|
||||
Content: summaryBuilder.String(),
|
||||
Timestamp: time.Now(),
|
||||
}
|
||||
resultMessages = append(resultMessages, summaryMsg)
|
||||
log.Info().Str("session_id", sessionID).Int("summary_len", summaryBuilder.Len()).Msg("[AgenticLoop] Final summary produced")
|
||||
} else if summaryErr != nil {
|
||||
log.Error().Err(summaryErr).Str("session_id", sessionID).Msg("[AgenticLoop] Final summary call failed")
|
||||
} else {
|
||||
log.Warn().Str("session_id", sessionID).Msg("[AgenticLoop] Final summary call returned empty content")
|
||||
}
|
||||
|
||||
return resultMessages
|
||||
}
|
||||
|
||||
// Abort aborts an ongoing session
|
||||
func (a *AgenticLoop) Abort(sessionID string) {
|
||||
a.mu.Lock()
|
||||
@@ -978,6 +1168,14 @@ func (a *AgenticLoop) SetSessionFSM(fsm *SessionFSM) {
|
||||
a.mu.Unlock()
|
||||
}
|
||||
|
||||
// SetKnowledgeAccumulator sets the knowledge accumulator for fact extraction.
|
||||
// This must be called before Execute to enable knowledge accumulation.
|
||||
func (a *AgenticLoop) SetKnowledgeAccumulator(ka *KnowledgeAccumulator) {
|
||||
a.mu.Lock()
|
||||
a.knowledgeAccumulator = ka
|
||||
a.mu.Unlock()
|
||||
}
|
||||
|
||||
// SetMaxTurns overrides the maximum number of agentic turns for this loop.
|
||||
func (a *AgenticLoop) SetMaxTurns(n int) {
|
||||
a.mu.Lock()
|
||||
@@ -1475,7 +1673,14 @@ confirmation prompt - you don't need to ask "Would you like me to...?" Just exec
|
||||
needed and the system will prompt the user to approve if required.`
|
||||
}
|
||||
|
||||
return a.baseSystemPrompt + modeContext
|
||||
prompt := a.baseSystemPrompt + modeContext
|
||||
|
||||
// Append accumulated knowledge facts to system prompt
|
||||
if ka := a.knowledgeAccumulator; ka != nil && ka.Len() > 0 {
|
||||
prompt += "\n\n" + ka.Render()
|
||||
}
|
||||
|
||||
return prompt
|
||||
}
|
||||
|
||||
// AnswerQuestion provides an answer to a pending question
|
||||
|
||||
@@ -11,6 +11,7 @@ import (
|
||||
"github.com/rcourtman/pulse-go-rewrite/internal/ai/approval"
|
||||
"github.com/rcourtman/pulse-go-rewrite/internal/ai/providers"
|
||||
"github.com/rcourtman/pulse-go-rewrite/internal/ai/tools"
|
||||
"github.com/rcourtman/pulse-go-rewrite/internal/models"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/stretchr/testify/require"
|
||||
@@ -164,6 +165,83 @@ func TestAgenticLoop(t *testing.T) {
|
||||
// In a real execution, loop would check aborted map
|
||||
assert.True(t, loop.aborted["abort-me"])
|
||||
})
|
||||
|
||||
t.Run("KnowledgeAccumulatorInjectedInSystemPrompt", func(t *testing.T) {
|
||||
state := models.StateSnapshot{
|
||||
VMs: []models.VM{{
|
||||
ID: "vm-100",
|
||||
VMID: 100,
|
||||
Name: "vm-one",
|
||||
Node: "node1",
|
||||
Status: "running",
|
||||
CPU: 0.5,
|
||||
CPUs: 4,
|
||||
Memory: models.Memory{
|
||||
Usage: 0.42,
|
||||
Used: 4 * 1024 * 1024 * 1024,
|
||||
Total: 8 * 1024 * 1024 * 1024,
|
||||
},
|
||||
}},
|
||||
}
|
||||
execWithState := tools.NewPulseToolExecutor(tools.ExecutorConfig{
|
||||
StateProvider: &mockStateProvider{state: state},
|
||||
})
|
||||
mockProvider := &MockProvider{}
|
||||
loop := NewAgenticLoop(mockProvider, execWithState, "You are a helper")
|
||||
ka := NewKnowledgeAccumulator()
|
||||
loop.SetKnowledgeAccumulator(ka)
|
||||
|
||||
ctx := context.Background()
|
||||
sessionID := "ka-session"
|
||||
messages := []Message{{Role: "user", Content: "Get VM status"}}
|
||||
|
||||
mockProvider.On("ChatStream", mock.Anything, mock.MatchedBy(func(req providers.ChatRequest) bool {
|
||||
return len(req.Messages) == 1
|
||||
}), mock.Anything).Return(nil).Run(func(args mock.Arguments) {
|
||||
callback := args.Get(2).(providers.StreamCallback)
|
||||
callback(providers.StreamEvent{
|
||||
Type: "tool_start",
|
||||
Data: providers.ToolStartEvent{ID: "call_1", Name: "pulse_query"},
|
||||
})
|
||||
callback(providers.StreamEvent{
|
||||
Type: "done",
|
||||
Data: providers.DoneEvent{
|
||||
ToolCalls: []providers.ToolCall{{
|
||||
ID: "call_1",
|
||||
Name: "pulse_query",
|
||||
Input: map[string]interface{}{
|
||||
"action": "get",
|
||||
"resource_type": "vm",
|
||||
"resource_id": "100",
|
||||
},
|
||||
}},
|
||||
},
|
||||
})
|
||||
}).Once()
|
||||
|
||||
var secondSystem string
|
||||
mockProvider.On("ChatStream", mock.Anything, mock.MatchedBy(func(req providers.ChatRequest) bool {
|
||||
return len(req.Messages) == 3
|
||||
}), mock.Anything).Return(nil).Run(func(args mock.Arguments) {
|
||||
req := args.Get(1).(providers.ChatRequest)
|
||||
secondSystem = req.System
|
||||
callback := args.Get(2).(providers.StreamCallback)
|
||||
callback(providers.StreamEvent{
|
||||
Type: "content",
|
||||
Data: providers.ContentEvent{Text: "Done."},
|
||||
})
|
||||
callback(providers.StreamEvent{
|
||||
Type: "done",
|
||||
Data: providers.DoneEvent{},
|
||||
})
|
||||
}).Once()
|
||||
|
||||
_, err := loop.Execute(ctx, sessionID, messages, func(event StreamEvent) {})
|
||||
require.NoError(t, err)
|
||||
require.NotEmpty(t, secondSystem)
|
||||
assert.Contains(t, secondSystem, "## Known Facts")
|
||||
assert.Contains(t, secondSystem, "vm-one")
|
||||
})
|
||||
}
|
||||
|
||||
func TestAgenticLoop_UpdateTools(t *testing.T) {
|
||||
|
||||
@@ -0,0 +1,216 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// FactCategory classifies a knowledge fact for grouping in the rendered output.
|
||||
type FactCategory string
|
||||
|
||||
const (
|
||||
FactCategoryResource FactCategory = "resource"
|
||||
FactCategoryStorage FactCategory = "storage"
|
||||
FactCategoryDiscovery FactCategory = "discovery"
|
||||
FactCategoryExec FactCategory = "exec"
|
||||
FactCategoryMetrics FactCategory = "metrics"
|
||||
FactCategoryFinding FactCategory = "finding"
|
||||
)
|
||||
|
||||
// Fact is a single extracted knowledge entry.
|
||||
type Fact struct {
|
||||
Category FactCategory
|
||||
Key string // Dedup key, e.g. "lxc:delly:106:status"
|
||||
Value string // Compact value, e.g. "running, Postfix, hostname=patrol-signal-test"
|
||||
ObservedAt time.Time
|
||||
Turn int
|
||||
}
|
||||
|
||||
const (
|
||||
defaultMaxEntries = 60
|
||||
defaultMaxChars = 2000
|
||||
maxValueLen = 200
|
||||
)
|
||||
|
||||
// KnowledgeAccumulator stores extracted facts from tool results.
|
||||
// Per-session, in-memory, bounded. Facts are keyed for upsert semantics.
|
||||
// Thread-safe: all methods are protected by a mutex.
|
||||
type KnowledgeAccumulator struct {
|
||||
mu sync.Mutex
|
||||
facts map[string]*Fact // key -> fact (upsert: same key updates value)
|
||||
order []string // insertion order for LRU eviction
|
||||
totalChars int
|
||||
maxEntries int
|
||||
maxChars int
|
||||
currentTurn int
|
||||
}
|
||||
|
||||
// NewKnowledgeAccumulator creates a new bounded accumulator.
|
||||
func NewKnowledgeAccumulator() *KnowledgeAccumulator {
|
||||
return &KnowledgeAccumulator{
|
||||
facts: make(map[string]*Fact),
|
||||
maxEntries: defaultMaxEntries,
|
||||
maxChars: defaultMaxChars,
|
||||
}
|
||||
}
|
||||
|
||||
// SetTurn updates the current turn number for new facts.
|
||||
func (ka *KnowledgeAccumulator) SetTurn(turn int) {
|
||||
ka.mu.Lock()
|
||||
ka.currentTurn = turn
|
||||
ka.mu.Unlock()
|
||||
}
|
||||
|
||||
// AddFact upserts a fact. If the key already exists, the value is updated.
|
||||
// If over budget after insertion, the oldest non-pinned fact is evicted.
|
||||
func (ka *KnowledgeAccumulator) AddFact(category FactCategory, key, value string) {
|
||||
if key == "" || value == "" {
|
||||
return
|
||||
}
|
||||
|
||||
ka.mu.Lock()
|
||||
defer ka.mu.Unlock()
|
||||
|
||||
// Truncate value
|
||||
if len(value) > maxValueLen {
|
||||
value = value[:maxValueLen]
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
|
||||
if existing, ok := ka.facts[key]; ok {
|
||||
// Upsert: update existing fact
|
||||
ka.totalChars -= len(existing.Value)
|
||||
existing.Value = value
|
||||
existing.Category = category
|
||||
existing.ObservedAt = now
|
||||
existing.Turn = ka.currentTurn
|
||||
ka.totalChars += len(value)
|
||||
} else {
|
||||
// New fact
|
||||
fact := &Fact{
|
||||
Category: category,
|
||||
Key: key,
|
||||
Value: value,
|
||||
ObservedAt: now,
|
||||
Turn: ka.currentTurn,
|
||||
}
|
||||
ka.facts[key] = fact
|
||||
ka.order = append(ka.order, key)
|
||||
ka.totalChars += len(value)
|
||||
}
|
||||
|
||||
// Evict until within budget
|
||||
ka.evict()
|
||||
}
|
||||
|
||||
// evict removes the oldest non-pinned facts until within budget.
|
||||
// Facts from the current or previous turn are soft-pinned (not evicted).
|
||||
// Caller must hold ka.mu.
|
||||
func (ka *KnowledgeAccumulator) evict() {
|
||||
for (len(ka.facts) > ka.maxEntries || ka.totalChars > ka.maxChars) && len(ka.order) > 0 {
|
||||
evicted := false
|
||||
for i, key := range ka.order {
|
||||
fact, ok := ka.facts[key]
|
||||
if !ok {
|
||||
// Stale key in order slice — remove it
|
||||
ka.order = append(ka.order[:i], ka.order[i+1:]...)
|
||||
evicted = true
|
||||
break
|
||||
}
|
||||
// Soft-pin: don't evict facts from current or previous turn
|
||||
if fact.Turn >= ka.currentTurn-1 {
|
||||
continue
|
||||
}
|
||||
// Evict this fact
|
||||
ka.totalChars -= len(fact.Value)
|
||||
delete(ka.facts, key)
|
||||
ka.order = append(ka.order[:i], ka.order[i+1:]...)
|
||||
evicted = true
|
||||
break
|
||||
}
|
||||
if !evicted {
|
||||
// All remaining facts are pinned — can't evict more
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Len returns the number of facts stored.
|
||||
func (ka *KnowledgeAccumulator) Len() int {
|
||||
ka.mu.Lock()
|
||||
defer ka.mu.Unlock()
|
||||
return len(ka.facts)
|
||||
}
|
||||
|
||||
// TotalChars returns the total character count of all fact values.
|
||||
func (ka *KnowledgeAccumulator) TotalChars() int {
|
||||
ka.mu.Lock()
|
||||
defer ka.mu.Unlock()
|
||||
return ka.totalChars
|
||||
}
|
||||
|
||||
// Lookup returns the fact value for a given key, or empty string if not found.
|
||||
func (ka *KnowledgeAccumulator) Lookup(key string) (string, bool) {
|
||||
ka.mu.Lock()
|
||||
defer ka.mu.Unlock()
|
||||
if fact, ok := ka.facts[key]; ok {
|
||||
return fact.Value, true
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
// Render produces the system prompt section with grouped facts.
|
||||
func (ka *KnowledgeAccumulator) Render() string {
|
||||
ka.mu.Lock()
|
||||
defer ka.mu.Unlock()
|
||||
|
||||
if len(ka.facts) == 0 {
|
||||
return ""
|
||||
}
|
||||
|
||||
// Group facts by category
|
||||
categoryOrder := []FactCategory{
|
||||
FactCategoryResource,
|
||||
FactCategoryStorage,
|
||||
FactCategoryDiscovery,
|
||||
FactCategoryExec,
|
||||
FactCategoryMetrics,
|
||||
FactCategoryFinding,
|
||||
}
|
||||
|
||||
categoryLabels := map[FactCategory]string{
|
||||
FactCategoryResource: "Resources",
|
||||
FactCategoryStorage: "Storage",
|
||||
FactCategoryDiscovery: "Discovery",
|
||||
FactCategoryExec: "Exec",
|
||||
FactCategoryMetrics: "Metrics",
|
||||
FactCategoryFinding: "Findings",
|
||||
}
|
||||
|
||||
grouped := make(map[FactCategory][]*Fact)
|
||||
for _, key := range ka.order {
|
||||
if fact, ok := ka.facts[key]; ok {
|
||||
grouped[fact.Category] = append(grouped[fact.Category], fact)
|
||||
}
|
||||
}
|
||||
|
||||
var sb strings.Builder
|
||||
sb.WriteString("## Known Facts (auto-collected — do NOT re-query unless user asks for fresh data)")
|
||||
|
||||
for _, cat := range categoryOrder {
|
||||
facts, ok := grouped[cat]
|
||||
if !ok || len(facts) == 0 {
|
||||
continue
|
||||
}
|
||||
label := categoryLabels[cat]
|
||||
sb.WriteString(fmt.Sprintf("\n%s:", label))
|
||||
for _, fact := range facts {
|
||||
sb.WriteString(fmt.Sprintf("\n- %s", fact.Value))
|
||||
}
|
||||
}
|
||||
|
||||
return sb.String()
|
||||
}
|
||||
@@ -0,0 +1,183 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestKnowledgeAccumulator_Basic(t *testing.T) {
|
||||
ka := NewKnowledgeAccumulator()
|
||||
assert.Equal(t, 0, ka.Len())
|
||||
assert.Equal(t, "", ka.Render())
|
||||
|
||||
ka.AddFact(FactCategoryResource, "lxc:delly:106:status", "running, Postfix")
|
||||
assert.Equal(t, 1, ka.Len())
|
||||
|
||||
rendered := ka.Render()
|
||||
assert.Contains(t, rendered, "Known Facts")
|
||||
assert.Contains(t, rendered, "Resources:")
|
||||
assert.Contains(t, rendered, "running, Postfix")
|
||||
}
|
||||
|
||||
func TestKnowledgeAccumulator_Upsert(t *testing.T) {
|
||||
ka := NewKnowledgeAccumulator()
|
||||
|
||||
ka.AddFact(FactCategoryResource, "lxc:delly:106:status", "running, CPU=5%")
|
||||
assert.Equal(t, 1, ka.Len())
|
||||
|
||||
// Upsert with same key should update value, not add a new entry
|
||||
ka.AddFact(FactCategoryResource, "lxc:delly:106:status", "stopped")
|
||||
assert.Equal(t, 1, ka.Len())
|
||||
|
||||
rendered := ka.Render()
|
||||
assert.Contains(t, rendered, "stopped")
|
||||
assert.NotContains(t, rendered, "CPU=5%")
|
||||
}
|
||||
|
||||
func TestKnowledgeAccumulator_EmptyKeyOrValue(t *testing.T) {
|
||||
ka := NewKnowledgeAccumulator()
|
||||
|
||||
ka.AddFact(FactCategoryResource, "", "some value")
|
||||
assert.Equal(t, 0, ka.Len())
|
||||
|
||||
ka.AddFact(FactCategoryResource, "some-key", "")
|
||||
assert.Equal(t, 0, ka.Len())
|
||||
}
|
||||
|
||||
func TestKnowledgeAccumulator_ValueTruncation(t *testing.T) {
|
||||
ka := NewKnowledgeAccumulator()
|
||||
|
||||
longValue := strings.Repeat("x", 300)
|
||||
ka.AddFact(FactCategoryExec, "exec:host:cmd", longValue)
|
||||
|
||||
fact := ka.facts["exec:host:cmd"]
|
||||
require.NotNil(t, fact)
|
||||
assert.Equal(t, maxValueLen, len(fact.Value))
|
||||
}
|
||||
|
||||
func TestKnowledgeAccumulator_MaxEntries(t *testing.T) {
|
||||
ka := NewKnowledgeAccumulator()
|
||||
ka.maxEntries = 5
|
||||
ka.maxChars = 100000 // High limit so entries are the constraint
|
||||
|
||||
// Insert 5 facts at turn 0
|
||||
ka.SetTurn(0)
|
||||
for i := 0; i < 5; i++ {
|
||||
ka.AddFact(FactCategoryResource, keyForIndex(i), valForIndex(i))
|
||||
}
|
||||
assert.Equal(t, 5, ka.Len())
|
||||
|
||||
// Move to turn 2 so turn 0 facts are evictable
|
||||
ka.SetTurn(2)
|
||||
ka.AddFact(FactCategoryResource, "new-key", "new-value")
|
||||
|
||||
// Should have evicted one old fact to make room
|
||||
assert.Equal(t, 5, ka.Len())
|
||||
_, hasNew := ka.facts["new-key"]
|
||||
assert.True(t, hasNew, "new fact should be present")
|
||||
}
|
||||
|
||||
func TestKnowledgeAccumulator_MaxChars(t *testing.T) {
|
||||
ka := NewKnowledgeAccumulator()
|
||||
ka.maxEntries = 1000
|
||||
ka.maxChars = 50 // Very low char budget
|
||||
|
||||
ka.SetTurn(0)
|
||||
ka.AddFact(FactCategoryResource, "k1", strings.Repeat("a", 20))
|
||||
ka.AddFact(FactCategoryResource, "k2", strings.Repeat("b", 20))
|
||||
|
||||
// Move forward so old facts can be evicted
|
||||
ka.SetTurn(2)
|
||||
ka.AddFact(FactCategoryResource, "k3", strings.Repeat("c", 20))
|
||||
|
||||
// Should have evicted oldest to stay within budget
|
||||
assert.LessOrEqual(t, ka.TotalChars(), 50+maxValueLen, "should stay near char budget")
|
||||
}
|
||||
|
||||
func TestKnowledgeAccumulator_SoftPinCurrentTurn(t *testing.T) {
|
||||
ka := NewKnowledgeAccumulator()
|
||||
ka.maxEntries = 2
|
||||
ka.maxChars = 100000
|
||||
|
||||
// All facts in current turn should be soft-pinned
|
||||
ka.SetTurn(5)
|
||||
ka.AddFact(FactCategoryResource, "k1", "val1")
|
||||
ka.AddFact(FactCategoryResource, "k2", "val2")
|
||||
ka.AddFact(FactCategoryResource, "k3", "val3")
|
||||
|
||||
// All 3 should survive even though maxEntries=2 because they're all from current turn
|
||||
assert.Equal(t, 3, ka.Len())
|
||||
}
|
||||
|
||||
func TestKnowledgeAccumulator_CategoryGrouping(t *testing.T) {
|
||||
ka := NewKnowledgeAccumulator()
|
||||
|
||||
ka.AddFact(FactCategoryResource, "res1", "LXC 106 running")
|
||||
ka.AddFact(FactCategoryStorage, "stor1", "PBS available, 42% used")
|
||||
ka.AddFact(FactCategoryExec, "exec1", "exit=0, service active")
|
||||
ka.AddFact(FactCategoryFinding, "find1", "warning: high CPU on vm101")
|
||||
|
||||
rendered := ka.Render()
|
||||
|
||||
// Verify category headers appear
|
||||
assert.Contains(t, rendered, "Resources:")
|
||||
assert.Contains(t, rendered, "Storage:")
|
||||
assert.Contains(t, rendered, "Exec:")
|
||||
assert.Contains(t, rendered, "Findings:")
|
||||
|
||||
// Verify values appear
|
||||
assert.Contains(t, rendered, "LXC 106 running")
|
||||
assert.Contains(t, rendered, "PBS available, 42% used")
|
||||
assert.Contains(t, rendered, "exit=0, service active")
|
||||
assert.Contains(t, rendered, "warning: high CPU on vm101")
|
||||
}
|
||||
|
||||
func TestKnowledgeAccumulator_RenderOrder(t *testing.T) {
|
||||
ka := NewKnowledgeAccumulator()
|
||||
|
||||
// Add in non-category order
|
||||
ka.AddFact(FactCategoryFinding, "f1", "finding")
|
||||
ka.AddFact(FactCategoryResource, "r1", "resource")
|
||||
|
||||
rendered := ka.Render()
|
||||
// Resources should appear before Findings in the rendered output
|
||||
resIdx := strings.Index(rendered, "Resources:")
|
||||
findIdx := strings.Index(rendered, "Findings:")
|
||||
assert.Greater(t, findIdx, resIdx, "Resources should come before Findings")
|
||||
}
|
||||
|
||||
func TestKnowledgeAccumulator_Lookup(t *testing.T) {
|
||||
ka := NewKnowledgeAccumulator()
|
||||
|
||||
ka.AddFact(FactCategoryDiscovery, "discovery:delly:106", "service=Postfix, hostname=patrol-signal-test")
|
||||
|
||||
val, found := ka.Lookup("discovery:delly:106")
|
||||
assert.True(t, found)
|
||||
assert.Equal(t, "service=Postfix, hostname=patrol-signal-test", val)
|
||||
|
||||
val, found = ka.Lookup("discovery:minipc:200")
|
||||
assert.False(t, found)
|
||||
assert.Equal(t, "", val)
|
||||
}
|
||||
|
||||
func TestKnowledgeAccumulator_SetTurn(t *testing.T) {
|
||||
ka := NewKnowledgeAccumulator()
|
||||
|
||||
ka.SetTurn(3)
|
||||
ka.AddFact(FactCategoryResource, "k1", "val1")
|
||||
|
||||
fact := ka.facts["k1"]
|
||||
require.NotNil(t, fact)
|
||||
assert.Equal(t, 3, fact.Turn)
|
||||
}
|
||||
|
||||
func keyForIndex(i int) string {
|
||||
return strings.Repeat("k", i+1)
|
||||
}
|
||||
|
||||
func valForIndex(i int) string {
|
||||
return strings.Repeat("v", i+1)
|
||||
}
|
||||
@@ -0,0 +1,582 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// FactEntry is the output of ExtractFacts — ready to feed into KnowledgeAccumulator.AddFact.
|
||||
type FactEntry struct {
|
||||
Category FactCategory
|
||||
Key string
|
||||
Value string
|
||||
}
|
||||
|
||||
// ExtractFacts deterministically extracts knowledge facts from a tool result.
|
||||
// No LLM calls. Parses the JSON text from FormatToolResult() output.
|
||||
// Returns empty slice on parse errors or unrecognized tools — never panics.
|
||||
//
|
||||
// Tool results use NewJSONResult (direct struct marshaling), NOT ToolResponse wrapper.
|
||||
// So the resultText is the JSON of the response struct directly (e.g. ResourceResponse).
|
||||
func ExtractFacts(toolName string, toolInput map[string]interface{}, resultText string) []FactEntry {
|
||||
switch toolName {
|
||||
case "pulse_query":
|
||||
return extractQueryFacts(toolInput, resultText)
|
||||
case "pulse_storage":
|
||||
return extractStorageFacts(toolInput, resultText)
|
||||
case "pulse_discovery":
|
||||
return extractDiscoveryFacts(toolInput, resultText)
|
||||
case "pulse_read", "pulse_run_command":
|
||||
return extractExecFacts(toolInput, resultText)
|
||||
case "pulse_metrics":
|
||||
return extractMetricsFacts(toolInput, resultText)
|
||||
case "patrol_report_finding":
|
||||
return extractFindingFacts(toolInput, resultText)
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
// --- pulse_query ---
|
||||
|
||||
func extractQueryFacts(input map[string]interface{}, resultText string) []FactEntry {
|
||||
action := strFromMap(input, "action")
|
||||
if action == "" {
|
||||
action = strFromMap(input, "type")
|
||||
}
|
||||
|
||||
switch action {
|
||||
case "get":
|
||||
return extractQueryGetFacts(input, resultText)
|
||||
case "search":
|
||||
return extractQuerySearchFacts(input, resultText)
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func extractQueryGetFacts(input map[string]interface{}, resultText string) []FactEntry {
|
||||
// Tool results are direct JSON (NewJSONResult), no ToolResponse wrapper.
|
||||
// ResourceResponse has nested CPU/Memory structs.
|
||||
var resource struct {
|
||||
Type string `json:"type"`
|
||||
Name string `json:"name"`
|
||||
Status string `json:"status"`
|
||||
Node string `json:"node"`
|
||||
ID string `json:"id"`
|
||||
VMID int `json:"vmid"`
|
||||
Host string `json:"host"`
|
||||
CPU struct {
|
||||
Percent float64 `json:"percent"`
|
||||
} `json:"cpu"`
|
||||
Memory struct {
|
||||
Percent float64 `json:"percent"`
|
||||
} `json:"memory"`
|
||||
// Error field for not-found responses
|
||||
Error string `json:"error"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(resultText), &resource); err != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Skip error/not-found responses
|
||||
if resource.Error != "" {
|
||||
return nil
|
||||
}
|
||||
if resource.Name == "" && resource.ID == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
resType := resource.Type
|
||||
if resType == "" {
|
||||
resType = strFromMap(input, "resource_type")
|
||||
}
|
||||
node := resource.Node
|
||||
if node == "" {
|
||||
node = resource.Host
|
||||
}
|
||||
id := resource.ID
|
||||
if id == "" && resource.VMID > 0 {
|
||||
id = fmt.Sprintf("%d", resource.VMID)
|
||||
}
|
||||
if id == "" {
|
||||
id = resource.Name
|
||||
}
|
||||
|
||||
key := fmt.Sprintf("%s:%s:%s:status", resType, node, id)
|
||||
|
||||
var parts []string
|
||||
if resource.Status != "" {
|
||||
parts = append(parts, resource.Status)
|
||||
}
|
||||
if resource.Name != "" {
|
||||
parts = append(parts, resource.Name)
|
||||
}
|
||||
if resource.CPU.Percent > 0 {
|
||||
parts = append(parts, fmt.Sprintf("CPU=%.1f%%", resource.CPU.Percent))
|
||||
}
|
||||
if resource.Memory.Percent > 0 {
|
||||
parts = append(parts, fmt.Sprintf("Mem=%.1f%%", resource.Memory.Percent))
|
||||
}
|
||||
|
||||
if len(parts) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
return []FactEntry{{
|
||||
Category: FactCategoryResource,
|
||||
Key: key,
|
||||
Value: strings.Join(parts, ", "),
|
||||
}}
|
||||
}
|
||||
|
||||
func extractQuerySearchFacts(input map[string]interface{}, resultText string) []FactEntry {
|
||||
query := strFromMap(input, "query")
|
||||
if query == "" {
|
||||
query = strFromMap(input, "search")
|
||||
}
|
||||
|
||||
// ResourceSearchResponse — direct JSON, no wrapper
|
||||
var resp struct {
|
||||
Matches []struct {
|
||||
Name string `json:"name"`
|
||||
Status string `json:"status"`
|
||||
Type string `json:"type"`
|
||||
} `json:"matches"`
|
||||
Total int `json:"total"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(resultText), &resp); err != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
total := resp.Total
|
||||
if total == 0 && len(resp.Matches) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Summarize first 5 matches
|
||||
var summaryParts []string
|
||||
limit := 5
|
||||
if len(resp.Matches) < limit {
|
||||
limit = len(resp.Matches)
|
||||
}
|
||||
for _, m := range resp.Matches[:limit] {
|
||||
entry := m.Name
|
||||
if m.Status != "" {
|
||||
entry += " (" + m.Status + ")"
|
||||
}
|
||||
summaryParts = append(summaryParts, entry)
|
||||
}
|
||||
|
||||
value := fmt.Sprintf("%d results: %s", total, strings.Join(summaryParts, ", "))
|
||||
|
||||
return []FactEntry{{
|
||||
Category: FactCategoryResource,
|
||||
Key: fmt.Sprintf("search:%s:summary", query),
|
||||
Value: truncateValue(value),
|
||||
}}
|
||||
}
|
||||
|
||||
// --- pulse_storage ---
|
||||
|
||||
func extractStorageFacts(input map[string]interface{}, resultText string) []FactEntry {
|
||||
action := strFromMap(input, "action")
|
||||
if action == "" {
|
||||
action = strFromMap(input, "type")
|
||||
}
|
||||
|
||||
switch action {
|
||||
case "pools":
|
||||
return extractStoragePoolFacts(resultText)
|
||||
case "backup_tasks":
|
||||
return extractBackupTaskFacts(resultText)
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func extractStoragePoolFacts(resultText string) []FactEntry {
|
||||
// StorageResponse — direct JSON, no wrapper
|
||||
var resp struct {
|
||||
Pools []struct {
|
||||
Name string `json:"name"`
|
||||
Node string `json:"node"`
|
||||
Nodes []string `json:"nodes"`
|
||||
Type string `json:"type"`
|
||||
Status string `json:"status"`
|
||||
Active bool `json:"active"`
|
||||
UsagePercent float64 `json:"usage_percent"`
|
||||
TotalGB float64 `json:"total_gb"`
|
||||
UsedGB float64 `json:"used_gb"`
|
||||
} `json:"pools"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(resultText), &resp); err != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
var facts []FactEntry
|
||||
for _, pool := range resp.Pools {
|
||||
node := pool.Node
|
||||
if node == "" && len(pool.Nodes) > 0 {
|
||||
node = strings.Join(pool.Nodes, "+")
|
||||
}
|
||||
|
||||
freeGB := pool.TotalGB - pool.UsedGB
|
||||
var parts []string
|
||||
if pool.Type != "" {
|
||||
parts = append(parts, pool.Type)
|
||||
}
|
||||
if pool.Status != "" {
|
||||
parts = append(parts, pool.Status)
|
||||
}
|
||||
if pool.Active {
|
||||
parts = append(parts, fmt.Sprintf("active on %s", node))
|
||||
}
|
||||
parts = append(parts, fmt.Sprintf("%.1f%% used", pool.UsagePercent))
|
||||
if freeGB > 0 {
|
||||
parts = append(parts, fmt.Sprintf("%.0fGB free", freeGB))
|
||||
}
|
||||
|
||||
facts = append(facts, FactEntry{
|
||||
Category: FactCategoryStorage,
|
||||
Key: fmt.Sprintf("storage:%s:%s", node, pool.Name),
|
||||
Value: truncateValue(strings.Join(parts, ", ")),
|
||||
})
|
||||
}
|
||||
return facts
|
||||
}
|
||||
|
||||
func extractBackupTaskFacts(resultText string) []FactEntry {
|
||||
// BackupTasksListResponse — direct JSON, no wrapper
|
||||
var resp struct {
|
||||
Tasks []struct {
|
||||
VMID string `json:"vmid"`
|
||||
Node string `json:"node"`
|
||||
Status string `json:"status"`
|
||||
StartTime string `json:"start_time"`
|
||||
Error string `json:"error"`
|
||||
} `json:"tasks"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(resultText), &resp); err != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
var facts []FactEntry
|
||||
for _, task := range resp.Tasks {
|
||||
// Only record failures
|
||||
if task.Status == "ok" || task.Status == "success" || task.Status == "" {
|
||||
continue
|
||||
}
|
||||
var parts []string
|
||||
parts = append(parts, task.Status)
|
||||
if task.StartTime != "" {
|
||||
parts = append(parts, "at "+task.StartTime)
|
||||
}
|
||||
if task.Error != "" {
|
||||
parts = append(parts, "error="+task.Error)
|
||||
}
|
||||
facts = append(facts, FactEntry{
|
||||
Category: FactCategoryStorage,
|
||||
Key: fmt.Sprintf("backup:%s:%s", task.VMID, task.Node),
|
||||
Value: truncateValue(strings.Join(parts, ", ")),
|
||||
})
|
||||
}
|
||||
return facts
|
||||
}
|
||||
|
||||
// --- pulse_discovery ---
|
||||
|
||||
func extractDiscoveryFacts(input map[string]interface{}, resultText string) []FactEntry {
|
||||
// ResourceDiscoveryInfo — direct JSON, no wrapper
|
||||
var disc struct {
|
||||
ServiceType string `json:"service_type"`
|
||||
Hostname string `json:"hostname"`
|
||||
HostID string `json:"host_id"`
|
||||
ResourceID string `json:"resource_id"`
|
||||
Ports []struct {
|
||||
Port int `json:"port"`
|
||||
} `json:"ports"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(resultText), &disc); err != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
host := disc.HostID
|
||||
if host == "" {
|
||||
host = strFromMap(input, "host")
|
||||
}
|
||||
id := disc.ResourceID
|
||||
if id == "" {
|
||||
id = strFromMap(input, "resource_id")
|
||||
}
|
||||
|
||||
var parts []string
|
||||
if disc.ServiceType != "" {
|
||||
parts = append(parts, "service="+disc.ServiceType)
|
||||
}
|
||||
if disc.Hostname != "" {
|
||||
parts = append(parts, "hostname="+disc.Hostname)
|
||||
}
|
||||
if len(disc.Ports) > 0 {
|
||||
var portStrs []string
|
||||
for _, p := range disc.Ports {
|
||||
portStrs = append(portStrs, fmt.Sprintf("%d", p.Port))
|
||||
}
|
||||
parts = append(parts, "ports=["+strings.Join(portStrs, ",")+"]")
|
||||
}
|
||||
|
||||
if len(parts) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
return []FactEntry{{
|
||||
Category: FactCategoryDiscovery,
|
||||
Key: fmt.Sprintf("discovery:%s:%s", host, id),
|
||||
Value: truncateValue(strings.Join(parts, ", ")),
|
||||
}}
|
||||
}
|
||||
|
||||
// --- pulse_read / pulse_run_command ---
|
||||
|
||||
func extractExecFacts(input map[string]interface{}, resultText string) []FactEntry {
|
||||
host := strFromMap(input, "target_host")
|
||||
if host == "" {
|
||||
host = strFromMap(input, "host")
|
||||
}
|
||||
cmd := strFromMap(input, "command")
|
||||
if cmd == "" {
|
||||
// For pulse_read file/tail/find actions, use action+path to distinguish
|
||||
// different file reads on the same host.
|
||||
action := strFromMap(input, "action")
|
||||
path := strFromMap(input, "path")
|
||||
if action != "" && path != "" {
|
||||
cmd = action + ":" + path
|
||||
} else if action != "" {
|
||||
cmd = action
|
||||
}
|
||||
}
|
||||
if cmd == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Use first 60 chars of command as key prefix (longer to accommodate path)
|
||||
cmdPrefix := cmd
|
||||
if len(cmdPrefix) > 60 {
|
||||
cmdPrefix = cmdPrefix[:60]
|
||||
}
|
||||
|
||||
// Try to parse as CommandResponse (direct JSON, no wrapper)
|
||||
var cmdResp struct {
|
||||
Success bool `json:"success"`
|
||||
ExitCode int `json:"exit_code"`
|
||||
Output string `json:"output"`
|
||||
Stdout string `json:"stdout"`
|
||||
Error string `json:"error"`
|
||||
}
|
||||
|
||||
var value string
|
||||
if err := json.Unmarshal([]byte(resultText), &cmdResp); err == nil && (cmdResp.Output != "" || cmdResp.Stdout != "" || cmdResp.Error != "") {
|
||||
output := cmdResp.Output
|
||||
if output == "" {
|
||||
output = cmdResp.Stdout
|
||||
}
|
||||
if output == "" {
|
||||
output = cmdResp.Error
|
||||
}
|
||||
// Take first 2 lines
|
||||
lines := strings.SplitN(output, "\n", 3)
|
||||
summary := strings.Join(lines[:min(2, len(lines))], "; ")
|
||||
value = fmt.Sprintf("exit=%d, %s", cmdResp.ExitCode, summary)
|
||||
} else {
|
||||
// Fallback: use first 2 lines of raw result text
|
||||
lines := strings.SplitN(resultText, "\n", 3)
|
||||
summary := strings.Join(lines[:min(2, len(lines))], "; ")
|
||||
value = summary
|
||||
}
|
||||
|
||||
return []FactEntry{{
|
||||
Category: FactCategoryExec,
|
||||
Key: fmt.Sprintf("exec:%s:%s", host, cmdPrefix),
|
||||
Value: truncateValue(value),
|
||||
}}
|
||||
}
|
||||
|
||||
// --- pulse_metrics ---
|
||||
|
||||
func extractMetricsFacts(input map[string]interface{}, resultText string) []FactEntry {
|
||||
action := strFromMap(input, "action")
|
||||
if action == "" {
|
||||
action = strFromMap(input, "type")
|
||||
}
|
||||
|
||||
if action != "performance" {
|
||||
return nil
|
||||
}
|
||||
|
||||
resourceID := strFromMap(input, "resource_id")
|
||||
if resourceID == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
// MetricsResponse — direct JSON, no wrapper.
|
||||
// Summary is map[string]ResourceMetricsSummary keyed by resource ID.
|
||||
var resp struct {
|
||||
Summary map[string]struct {
|
||||
AvgCPU float64 `json:"avg_cpu"`
|
||||
MaxCPU float64 `json:"max_cpu"`
|
||||
AvgMemory float64 `json:"avg_memory"`
|
||||
MaxMemory float64 `json:"max_memory"`
|
||||
Trend string `json:"trend"`
|
||||
} `json:"summary"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(resultText), &resp); err != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Look up the summary for this resource (or take the first entry)
|
||||
var avgCPU, maxCPU float64
|
||||
var trend string
|
||||
if s, ok := resp.Summary[resourceID]; ok {
|
||||
avgCPU = s.AvgCPU
|
||||
maxCPU = s.MaxCPU
|
||||
trend = s.Trend
|
||||
} else {
|
||||
// Take first entry if resource ID doesn't match exactly
|
||||
for _, s := range resp.Summary {
|
||||
avgCPU = s.AvgCPU
|
||||
maxCPU = s.MaxCPU
|
||||
trend = s.Trend
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
var parts []string
|
||||
if avgCPU > 0 {
|
||||
parts = append(parts, fmt.Sprintf("avg_cpu=%.1f%%", avgCPU))
|
||||
}
|
||||
if maxCPU > 0 {
|
||||
parts = append(parts, fmt.Sprintf("max=%.1f%%", maxCPU))
|
||||
}
|
||||
if trend != "" {
|
||||
parts = append(parts, "trend="+trend)
|
||||
}
|
||||
|
||||
if len(parts) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
return []FactEntry{{
|
||||
Category: FactCategoryMetrics,
|
||||
Key: fmt.Sprintf("metrics:%s", resourceID),
|
||||
Value: truncateValue(strings.Join(parts, ", ")),
|
||||
}}
|
||||
}
|
||||
|
||||
// --- patrol_report_finding ---
|
||||
|
||||
func extractFindingFacts(input map[string]interface{}, resultText string) []FactEntry {
|
||||
key := strFromMap(input, "key")
|
||||
if key == "" {
|
||||
key = strFromMap(input, "finding_key")
|
||||
}
|
||||
severity := strFromMap(input, "severity")
|
||||
title := strFromMap(input, "title")
|
||||
resourceID := strFromMap(input, "resource_id")
|
||||
|
||||
if key == "" || title == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
var parts []string
|
||||
if severity != "" {
|
||||
parts = append(parts, severity)
|
||||
}
|
||||
parts = append(parts, title)
|
||||
if resourceID != "" {
|
||||
parts = append(parts, "on "+resourceID)
|
||||
}
|
||||
|
||||
return []FactEntry{{
|
||||
Category: FactCategoryFinding,
|
||||
Key: fmt.Sprintf("finding:%s", key),
|
||||
Value: truncateValue(strings.Join(parts, ": ")),
|
||||
}}
|
||||
}
|
||||
|
||||
// PredictFactKeys returns the KA fact keys that this tool call would produce,
|
||||
// based solely on the tool input (without needing the result).
|
||||
// Used by the gate to check if we already have facts for this call.
|
||||
// Returns nil if the key can't be predicted from input alone.
|
||||
func PredictFactKeys(toolName string, toolInput map[string]interface{}) []string {
|
||||
switch toolName {
|
||||
case "pulse_discovery":
|
||||
host := strFromMap(toolInput, "host_id")
|
||||
if host == "" {
|
||||
host = strFromMap(toolInput, "host")
|
||||
}
|
||||
id := strFromMap(toolInput, "resource_id")
|
||||
if host != "" && id != "" {
|
||||
return []string{fmt.Sprintf("discovery:%s:%s", host, id)}
|
||||
}
|
||||
case "pulse_read", "pulse_run_command":
|
||||
host := strFromMap(toolInput, "target_host")
|
||||
if host == "" {
|
||||
host = strFromMap(toolInput, "host")
|
||||
}
|
||||
cmd := strFromMap(toolInput, "command")
|
||||
if cmd == "" {
|
||||
// For file/tail/find actions, include path to distinguish different file reads
|
||||
action := strFromMap(toolInput, "action")
|
||||
path := strFromMap(toolInput, "path")
|
||||
if action != "" && path != "" {
|
||||
cmd = action + ":" + path
|
||||
} else if action != "" {
|
||||
cmd = action
|
||||
}
|
||||
}
|
||||
if host != "" && cmd != "" {
|
||||
cmdPrefix := cmd
|
||||
if len(cmdPrefix) > 60 {
|
||||
cmdPrefix = cmdPrefix[:60]
|
||||
}
|
||||
return []string{fmt.Sprintf("exec:%s:%s", host, cmdPrefix)}
|
||||
}
|
||||
case "pulse_metrics":
|
||||
action := strFromMap(toolInput, "action")
|
||||
if action == "" {
|
||||
action = strFromMap(toolInput, "type")
|
||||
}
|
||||
resourceID := strFromMap(toolInput, "resource_id")
|
||||
if action == "performance" && resourceID != "" {
|
||||
return []string{fmt.Sprintf("metrics:%s", resourceID)}
|
||||
}
|
||||
}
|
||||
// pulse_query get/search and pulse_storage: keys depend on result data,
|
||||
// can't predict from input alone. Return nil — these calls won't be gated.
|
||||
return nil
|
||||
}
|
||||
|
||||
// --- helpers ---
|
||||
|
||||
func strFromMap(m map[string]interface{}, key string) string {
|
||||
if m == nil {
|
||||
return ""
|
||||
}
|
||||
v, ok := m[key]
|
||||
if !ok {
|
||||
return ""
|
||||
}
|
||||
s, ok := v.(string)
|
||||
if !ok {
|
||||
return ""
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func truncateValue(s string) string {
|
||||
if len(s) > maxValueLen {
|
||||
return s[:maxValueLen]
|
||||
}
|
||||
return s
|
||||
}
|
||||
@@ -0,0 +1,279 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestExtractFacts_UnknownTool(t *testing.T) {
|
||||
facts := ExtractFacts("unknown_tool", nil, `{}`)
|
||||
assert.Empty(t, facts)
|
||||
}
|
||||
|
||||
func TestExtractFacts_InvalidJSON(t *testing.T) {
|
||||
facts := ExtractFacts("pulse_query", map[string]interface{}{"action": "get"}, "not json")
|
||||
assert.Empty(t, facts)
|
||||
}
|
||||
|
||||
func TestExtractFacts_QueryGet(t *testing.T) {
|
||||
input := map[string]interface{}{"action": "get", "resource_type": "lxc"}
|
||||
// Actual format from NewJSONResult(ResourceResponse): direct JSON, no wrapper.
|
||||
// CPU/Memory are nested structs.
|
||||
result := `{"type":"lxc","name":"postfix-server","status":"running","node":"delly","id":"lxc/106","vmid":106,"cpu":{"percent":2.5,"cores":4},"memory":{"percent":45.0,"used_gb":1.2,"total_gb":4.0}}`
|
||||
|
||||
facts := ExtractFacts("pulse_query", input, result)
|
||||
require.Len(t, facts, 1)
|
||||
|
||||
f := facts[0]
|
||||
assert.Equal(t, FactCategoryResource, f.Category)
|
||||
assert.Equal(t, "lxc:delly:lxc/106:status", f.Key)
|
||||
assert.Contains(t, f.Value, "running")
|
||||
assert.Contains(t, f.Value, "postfix-server")
|
||||
assert.Contains(t, f.Value, "CPU=2.5%")
|
||||
assert.Contains(t, f.Value, "Mem=45.0%")
|
||||
}
|
||||
|
||||
func TestExtractFacts_QueryGet_NotFound(t *testing.T) {
|
||||
input := map[string]interface{}{"action": "get", "resource_type": "vm"}
|
||||
result := `{"error":"not_found","resource_id":"999","type":"vm"}`
|
||||
|
||||
facts := ExtractFacts("pulse_query", input, result)
|
||||
assert.Empty(t, facts)
|
||||
}
|
||||
|
||||
func TestExtractFacts_QueryGet_NoCPU(t *testing.T) {
|
||||
// Some resources may not have CPU data populated
|
||||
input := map[string]interface{}{"action": "get", "resource_type": "container"}
|
||||
result := `{"type":"container","name":"test-ct","status":"stopped","node":"minipc","id":"lxc/200","cpu":{"percent":0},"memory":{"percent":0}}`
|
||||
|
||||
facts := ExtractFacts("pulse_query", input, result)
|
||||
require.Len(t, facts, 1)
|
||||
assert.Contains(t, facts[0].Value, "stopped")
|
||||
assert.Contains(t, facts[0].Value, "test-ct")
|
||||
assert.NotContains(t, facts[0].Value, "CPU=")
|
||||
}
|
||||
|
||||
func TestExtractFacts_QuerySearch(t *testing.T) {
|
||||
input := map[string]interface{}{"action": "search", "query": "postfix"}
|
||||
result := `{"query":"postfix","matches":[{"name":"postfix-lxc","status":"running","type":"lxc"},{"name":"mail-server","status":"stopped","type":"vm"}],"total":2}`
|
||||
|
||||
facts := ExtractFacts("pulse_query", input, result)
|
||||
require.Len(t, facts, 1)
|
||||
|
||||
f := facts[0]
|
||||
assert.Equal(t, FactCategoryResource, f.Category)
|
||||
assert.Equal(t, "search:postfix:summary", f.Key)
|
||||
assert.Contains(t, f.Value, "2 results")
|
||||
assert.Contains(t, f.Value, "postfix-lxc (running)")
|
||||
assert.Contains(t, f.Value, "mail-server (stopped)")
|
||||
}
|
||||
|
||||
func TestExtractFacts_QuerySearch_Empty(t *testing.T) {
|
||||
input := map[string]interface{}{"action": "search", "query": "nonexistent"}
|
||||
result := `{"query":"nonexistent","matches":[],"total":0}`
|
||||
|
||||
facts := ExtractFacts("pulse_query", input, result)
|
||||
assert.Empty(t, facts)
|
||||
}
|
||||
|
||||
func TestExtractFacts_StoragePools(t *testing.T) {
|
||||
input := map[string]interface{}{"action": "pools"}
|
||||
result := `{"pools":[{"name":"pbs-minipc","node":"","nodes":["delly","minipc"],"type":"PBS","status":"available","active":true,"usage_percent":42.7,"total_gb":1000,"used_gb":427}]}`
|
||||
|
||||
facts := ExtractFacts("pulse_storage", input, result)
|
||||
require.Len(t, facts, 1)
|
||||
|
||||
f := facts[0]
|
||||
assert.Equal(t, FactCategoryStorage, f.Category)
|
||||
assert.Contains(t, f.Key, "storage:")
|
||||
assert.Contains(t, f.Key, "pbs-minipc")
|
||||
assert.Contains(t, f.Value, "PBS")
|
||||
assert.Contains(t, f.Value, "42.7% used")
|
||||
assert.Contains(t, f.Value, "573GB free")
|
||||
}
|
||||
|
||||
func TestExtractFacts_BackupTasks_OnlyFailures(t *testing.T) {
|
||||
input := map[string]interface{}{"action": "backup_tasks"}
|
||||
result := `{"tasks":[{"vmid":"106","node":"delly","status":"ok"},{"vmid":"200","node":"minipc","status":"failed","start_time":"2024-01-15T03:00","error":"snapshot failed"}]}`
|
||||
|
||||
facts := ExtractFacts("pulse_storage", input, result)
|
||||
require.Len(t, facts, 1)
|
||||
|
||||
f := facts[0]
|
||||
assert.Equal(t, FactCategoryStorage, f.Category)
|
||||
assert.Equal(t, "backup:200:minipc", f.Key)
|
||||
assert.Contains(t, f.Value, "failed")
|
||||
assert.Contains(t, f.Value, "snapshot failed")
|
||||
}
|
||||
|
||||
func TestExtractFacts_Discovery(t *testing.T) {
|
||||
input := map[string]interface{}{"host": "delly", "resource_id": "106"}
|
||||
result := `{"service_type":"Postfix","hostname":"patrol-signal-test","host_id":"delly","resource_id":"106","ports":[{"port":25},{"port":22}]}`
|
||||
|
||||
facts := ExtractFacts("pulse_discovery", input, result)
|
||||
require.Len(t, facts, 1)
|
||||
|
||||
f := facts[0]
|
||||
assert.Equal(t, FactCategoryDiscovery, f.Category)
|
||||
assert.Equal(t, "discovery:delly:106", f.Key)
|
||||
assert.Contains(t, f.Value, "service=Postfix")
|
||||
assert.Contains(t, f.Value, "hostname=patrol-signal-test")
|
||||
assert.Contains(t, f.Value, "ports=[25,22]")
|
||||
}
|
||||
|
||||
func TestExtractFacts_Exec_JSON(t *testing.T) {
|
||||
input := map[string]interface{}{"command": "pvesm status | grep pbs-minipc", "target_host": "delly"}
|
||||
result := `{"success":true,"exit_code":0,"output":"pbs-minipc active 42.68%"}`
|
||||
|
||||
facts := ExtractFacts("pulse_read", input, result)
|
||||
require.Len(t, facts, 1)
|
||||
|
||||
f := facts[0]
|
||||
assert.Equal(t, FactCategoryExec, f.Category)
|
||||
assert.Contains(t, f.Key, "exec:delly:")
|
||||
assert.Contains(t, f.Value, "exit=0")
|
||||
assert.Contains(t, f.Value, "pbs-minipc")
|
||||
}
|
||||
|
||||
func TestExtractFacts_Exec_FallbackRaw(t *testing.T) {
|
||||
input := map[string]interface{}{"command": "some-cmd", "target_host": "host1"}
|
||||
result := `not json at all`
|
||||
|
||||
facts := ExtractFacts("pulse_read", input, result)
|
||||
require.Len(t, facts, 1)
|
||||
assert.Contains(t, facts[0].Value, "not json at all")
|
||||
}
|
||||
|
||||
func TestExtractFacts_Metrics(t *testing.T) {
|
||||
input := map[string]interface{}{"action": "performance", "resource_id": "vm101"}
|
||||
// Actual format: summary is map[string]ResourceMetricsSummary keyed by resource ID
|
||||
result := `{"resource_id":"vm101","period":"7d","summary":{"vm101":{"resource_id":"vm101","avg_cpu":12.3,"max_cpu":78.5,"avg_memory":65.0,"max_memory":89.0,"trend":"growing"}}}`
|
||||
|
||||
facts := ExtractFacts("pulse_metrics", input, result)
|
||||
require.Len(t, facts, 1)
|
||||
|
||||
f := facts[0]
|
||||
assert.Equal(t, FactCategoryMetrics, f.Category)
|
||||
assert.Equal(t, "metrics:vm101", f.Key)
|
||||
assert.Contains(t, f.Value, "avg_cpu=12.3%")
|
||||
assert.Contains(t, f.Value, "max=78.5%")
|
||||
assert.Contains(t, f.Value, "trend=growing")
|
||||
}
|
||||
|
||||
func TestExtractFacts_Metrics_EmptySummary(t *testing.T) {
|
||||
input := map[string]interface{}{"action": "performance", "resource_id": "vm101"}
|
||||
result := `{"resource_id":"vm101","period":"7d","summary":{}}`
|
||||
|
||||
facts := ExtractFacts("pulse_metrics", input, result)
|
||||
assert.Empty(t, facts)
|
||||
}
|
||||
|
||||
func TestExtractFacts_Finding(t *testing.T) {
|
||||
input := map[string]interface{}{
|
||||
"key": "high-cpu-vm101",
|
||||
"severity": "warning",
|
||||
"title": "High CPU usage on vm101",
|
||||
"resource_id": "vm101",
|
||||
}
|
||||
result := `{"id":"abc123","status":"created"}`
|
||||
|
||||
facts := ExtractFacts("patrol_report_finding", input, result)
|
||||
require.Len(t, facts, 1)
|
||||
|
||||
f := facts[0]
|
||||
assert.Equal(t, FactCategoryFinding, f.Category)
|
||||
assert.Equal(t, "finding:high-cpu-vm101", f.Key)
|
||||
assert.Contains(t, f.Value, "warning")
|
||||
assert.Contains(t, f.Value, "High CPU usage on vm101")
|
||||
assert.Contains(t, f.Value, "on vm101")
|
||||
}
|
||||
|
||||
func TestExtractFacts_Finding_MissingFields(t *testing.T) {
|
||||
// Missing key and title should return empty
|
||||
input := map[string]interface{}{"severity": "warning"}
|
||||
result := `{"id":"abc123"}`
|
||||
|
||||
facts := ExtractFacts("patrol_report_finding", input, result)
|
||||
assert.Empty(t, facts)
|
||||
}
|
||||
|
||||
func TestPredictFactKeys_Discovery(t *testing.T) {
|
||||
keys := PredictFactKeys("pulse_discovery", map[string]interface{}{
|
||||
"host_id": "delly",
|
||||
"resource_id": "106",
|
||||
})
|
||||
require.Len(t, keys, 1)
|
||||
assert.Equal(t, "discovery:delly:106", keys[0])
|
||||
}
|
||||
|
||||
func TestPredictFactKeys_DiscoveryAltFields(t *testing.T) {
|
||||
keys := PredictFactKeys("pulse_discovery", map[string]interface{}{
|
||||
"host": "minipc",
|
||||
"resource_id": "pbs-minipc",
|
||||
})
|
||||
require.Len(t, keys, 1)
|
||||
assert.Equal(t, "discovery:minipc:pbs-minipc", keys[0])
|
||||
}
|
||||
|
||||
func TestPredictFactKeys_Exec(t *testing.T) {
|
||||
keys := PredictFactKeys("pulse_read", map[string]interface{}{
|
||||
"target_host": "delly",
|
||||
"command": "pvesm status",
|
||||
})
|
||||
require.Len(t, keys, 1)
|
||||
assert.Equal(t, "exec:delly:pvesm status", keys[0])
|
||||
}
|
||||
|
||||
func TestPredictFactKeys_Metrics(t *testing.T) {
|
||||
keys := PredictFactKeys("pulse_metrics", map[string]interface{}{
|
||||
"action": "performance",
|
||||
"resource_id": "vm101",
|
||||
})
|
||||
require.Len(t, keys, 1)
|
||||
assert.Equal(t, "metrics:vm101", keys[0])
|
||||
}
|
||||
|
||||
func TestPredictFactKeys_UnpredictableTools(t *testing.T) {
|
||||
// pulse_query and pulse_storage keys depend on result data
|
||||
assert.Nil(t, PredictFactKeys("pulse_query", map[string]interface{}{"action": "get"}))
|
||||
assert.Nil(t, PredictFactKeys("pulse_storage", map[string]interface{}{"action": "pools"}))
|
||||
assert.Nil(t, PredictFactKeys("unknown_tool", nil))
|
||||
}
|
||||
|
||||
func TestPredictFactKeys_MissingFields(t *testing.T) {
|
||||
// Discovery without host_id should return nil
|
||||
assert.Nil(t, PredictFactKeys("pulse_discovery", map[string]interface{}{"resource_id": "106"}))
|
||||
// Exec without command should return nil
|
||||
assert.Nil(t, PredictFactKeys("pulse_read", map[string]interface{}{"target_host": "delly"}))
|
||||
}
|
||||
|
||||
func TestPredictFactKeys_FileReadDistinctPaths(t *testing.T) {
|
||||
// Different file paths on the same host should produce different keys
|
||||
keys1 := PredictFactKeys("pulse_read", map[string]interface{}{
|
||||
"target_host": "delly",
|
||||
"action": "file",
|
||||
"path": "/etc/pve/storage.cfg",
|
||||
})
|
||||
keys2 := PredictFactKeys("pulse_read", map[string]interface{}{
|
||||
"target_host": "delly",
|
||||
"action": "file",
|
||||
"path": "/var/log/pve/tasks/some-task-log",
|
||||
})
|
||||
require.Len(t, keys1, 1)
|
||||
require.Len(t, keys2, 1)
|
||||
assert.NotEqual(t, keys1[0], keys2[0], "different file paths must produce different keys")
|
||||
assert.Contains(t, keys1[0], "storage.cfg")
|
||||
assert.Contains(t, keys2[0], "tasks/some-task-log")
|
||||
}
|
||||
|
||||
func TestExtractFacts_ValueTruncation(t *testing.T) {
|
||||
input := map[string]interface{}{"command": "long-output-cmd", "target_host": "host1"}
|
||||
// Create a result with very long output
|
||||
longOutput := `{"success":true,"exit_code":0,"output":"` + bigContent(500) + `"}`
|
||||
|
||||
facts := ExtractFacts("pulse_read", input, longOutput)
|
||||
require.Len(t, facts, 1)
|
||||
assert.LessOrEqual(t, len(facts[0].Value), maxValueLen)
|
||||
}
|
||||
@@ -440,6 +440,11 @@ func (s *Service) ExecuteStream(ctx context.Context, req ExecuteRequest, callbac
|
||||
sessionFSM := sessions.GetSessionFSM(session.ID)
|
||||
loop.SetSessionFSM(sessionFSM)
|
||||
|
||||
// Set session-scoped knowledge accumulator for fact extraction across turns.
|
||||
// For user chat, this persists across messages so facts accumulate during a conversation.
|
||||
ka := sessions.GetKnowledgeAccumulator(session.ID)
|
||||
loop.SetKnowledgeAccumulator(ka)
|
||||
|
||||
// If the prefetcher resolved mentions, advance FSM past RESOLVING.
|
||||
// The prefetched context already contains the resource details (type, VMID, node, host)
|
||||
// so forcing the AI to redundantly call a read tool would be wasteful.
|
||||
@@ -469,6 +474,13 @@ func (s *Service) ExecuteStream(ctx context.Context, req ExecuteRequest, callbac
|
||||
}
|
||||
}
|
||||
|
||||
// Override max turns if the caller specified one (e.g. investigations use 15).
|
||||
// Reset after the call to avoid affecting concurrent sessions on the shared loop.
|
||||
if req.MaxTurns > 0 {
|
||||
loop.SetMaxTurns(req.MaxTurns)
|
||||
defer loop.SetMaxTurns(MaxAgenticTurns)
|
||||
}
|
||||
|
||||
resultMessages, err := loop.ExecuteWithTools(ctx, session.ID, messages, filteredTools, callback)
|
||||
|
||||
log.Debug().
|
||||
@@ -599,6 +611,13 @@ func (s *Service) ExecutePatrolStream(ctx context.Context, req PatrolRequest, ca
|
||||
sessionFSM := sessions.GetSessionFSM(session.ID)
|
||||
tempLoop.SetSessionFSM(sessionFSM)
|
||||
|
||||
// Create a fresh knowledge accumulator for this patrol run.
|
||||
// Unlike user chat (which reuses session-scoped KA across messages),
|
||||
// patrol runs need a clean slate to avoid stale facts from prior runs
|
||||
// (patrol-main reuses the same session ID across scheduled runs).
|
||||
ka := sessions.NewKnowledgeAccumulatorForRun(session.ID)
|
||||
tempLoop.SetKnowledgeAccumulator(ka)
|
||||
|
||||
// Set mid-run budget checker if configured
|
||||
if s.budgetChecker != nil {
|
||||
tempLoop.SetBudgetChecker(s.budgetChecker)
|
||||
|
||||
@@ -400,9 +400,8 @@ func TestService_SettersAndUpdateControlSettings(t *testing.T) {
|
||||
assert.True(t, hasTool(service.executor.ListTools(), "pulse_control"))
|
||||
}
|
||||
|
||||
func TestService_FilterToolsForPrompt_AllToolsReturned(t *testing.T) {
|
||||
// All tools should always be returned regardless of prompt content.
|
||||
// The approval mechanism handles control, not tool filtering.
|
||||
func TestService_FilterToolsForPrompt_ReadOnlyFiltersWriteTools(t *testing.T) {
|
||||
// Read-only prompts should not include write/control tools.
|
||||
service := NewService(Config{
|
||||
AIConfig: &config.AIConfig{ControlLevel: config.ControlLevelControlled},
|
||||
StateProvider: &mockStateProvider{},
|
||||
@@ -414,8 +413,24 @@ func TestService_FilterToolsForPrompt_AllToolsReturned(t *testing.T) {
|
||||
require.True(t, hasTool(service.executor.ListTools(), "pulse_docker"))
|
||||
require.True(t, hasTool(service.executor.ListTools(), "pulse_query"))
|
||||
|
||||
// All tools should be returned for any prompt
|
||||
// Read-only prompts should exclude write tools
|
||||
filtered := service.filterToolsForPrompt(context.Background(), "run uptime")
|
||||
assert.False(t, hasProviderTool(filtered, "pulse_control"))
|
||||
assert.False(t, hasProviderTool(filtered, "pulse_docker"))
|
||||
assert.True(t, hasProviderTool(filtered, "pulse_query"))
|
||||
}
|
||||
|
||||
func TestService_FilterToolsForPrompt_WriteIntentIncludesWriteTools(t *testing.T) {
|
||||
service := NewService(Config{
|
||||
AIConfig: &config.AIConfig{ControlLevel: config.ControlLevelControlled},
|
||||
StateProvider: &mockStateProvider{},
|
||||
AgentServer: &mockAgentServer{},
|
||||
})
|
||||
require.True(t, hasTool(service.executor.ListTools(), "pulse_control"))
|
||||
require.True(t, hasTool(service.executor.ListTools(), "pulse_docker"))
|
||||
require.True(t, hasTool(service.executor.ListTools(), "pulse_query"))
|
||||
|
||||
filtered := service.filterToolsForPrompt(context.Background(), "restart vm 101")
|
||||
assert.True(t, hasProviderTool(filtered, "pulse_control"))
|
||||
assert.True(t, hasProviderTool(filtered, "pulse_docker"))
|
||||
assert.True(t, hasProviderTool(filtered, "pulse_query"))
|
||||
|
||||
@@ -32,6 +32,11 @@ type SessionStore struct {
|
||||
// sessionToolSets holds per-session tool allowlists (in-memory only).
|
||||
// These keep tool availability stable across turns while allowing additive expansion.
|
||||
sessionToolSets map[string]map[string]bool
|
||||
|
||||
// knowledgeAccumulators holds per-session knowledge accumulators (in-memory only).
|
||||
// These extract and preserve key facts from tool results to prevent amnesia
|
||||
// when old tool results are compacted from the conversation context.
|
||||
knowledgeAccumulators map[string]*KnowledgeAccumulator
|
||||
}
|
||||
|
||||
// sessionData is the on-disk format for a session
|
||||
@@ -51,10 +56,11 @@ func NewSessionStore(dataDir string) (*SessionStore, error) {
|
||||
}
|
||||
|
||||
return &SessionStore{
|
||||
dataDir: sessionsDir,
|
||||
resolvedContexts: make(map[string]*ResolvedContext),
|
||||
sessionFSMs: make(map[string]*SessionFSM),
|
||||
sessionToolSets: make(map[string]map[string]bool),
|
||||
dataDir: sessionsDir,
|
||||
resolvedContexts: make(map[string]*ResolvedContext),
|
||||
sessionFSMs: make(map[string]*SessionFSM),
|
||||
sessionToolSets: make(map[string]map[string]bool),
|
||||
knowledgeAccumulators: make(map[string]*KnowledgeAccumulator),
|
||||
}, nil
|
||||
}
|
||||
|
||||
@@ -160,10 +166,11 @@ func (s *SessionStore) Delete(id string) error {
|
||||
return fmt.Errorf("failed to delete session: %w", err)
|
||||
}
|
||||
|
||||
// Also clean up resolved context and FSM
|
||||
// Also clean up resolved context, FSM, and knowledge accumulator
|
||||
delete(s.resolvedContexts, id)
|
||||
delete(s.sessionFSMs, id)
|
||||
delete(s.sessionToolSets, id)
|
||||
delete(s.knowledgeAccumulators, id)
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -355,6 +362,32 @@ func (s *SessionStore) GetSessionFSM(sessionID string) *SessionFSM {
|
||||
return fsm
|
||||
}
|
||||
|
||||
// GetKnowledgeAccumulator returns the knowledge accumulator for a session, creating one if needed.
|
||||
// For user chat sessions, this persists across messages (facts accumulate during a conversation).
|
||||
func (s *SessionStore) GetKnowledgeAccumulator(sessionID string) *KnowledgeAccumulator {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
ka, ok := s.knowledgeAccumulators[sessionID]
|
||||
if !ok {
|
||||
ka = NewKnowledgeAccumulator()
|
||||
s.knowledgeAccumulators[sessionID] = ka
|
||||
}
|
||||
return ka
|
||||
}
|
||||
|
||||
// NewKnowledgeAccumulatorForRun creates a fresh KA for a patrol run.
|
||||
// Unlike GetKnowledgeAccumulator (which reuses a session-scoped KA),
|
||||
// this always returns a new instance to avoid stale facts from prior runs.
|
||||
func (s *SessionStore) NewKnowledgeAccumulatorForRun(sessionID string) *KnowledgeAccumulator {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
ka := NewKnowledgeAccumulator()
|
||||
s.knowledgeAccumulators[sessionID] = ka
|
||||
return ka
|
||||
}
|
||||
|
||||
// ResetSessionFSM resets the FSM for a session (e.g., after context clear)
|
||||
func (s *SessionStore) ResetSessionFSM(sessionID string, keepProgress bool) {
|
||||
s.mu.Lock()
|
||||
@@ -434,6 +467,7 @@ func (s *SessionStore) ClearSessionState(sessionID string, keepPinned bool) {
|
||||
|
||||
if !keepPinned {
|
||||
delete(s.sessionToolSets, sessionID)
|
||||
delete(s.knowledgeAccumulators, sessionID)
|
||||
}
|
||||
|
||||
// Reset FSM coherently with context state
|
||||
|
||||
@@ -70,6 +70,7 @@ type ExecuteRequest struct {
|
||||
SessionID string `json:"session_id,omitempty"`
|
||||
Model string `json:"model,omitempty"`
|
||||
Mentions []StructuredMention `json:"mentions,omitempty"`
|
||||
MaxTurns int `json:"max_turns,omitempty"` // Override max agentic turns (0 = use default)
|
||||
}
|
||||
|
||||
// QuestionAnswer represents a user's answer to a question
|
||||
@@ -90,18 +91,20 @@ type ThinkingData struct {
|
||||
|
||||
// ToolStartData is the data for "tool_start" events
|
||||
type ToolStartData struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Input string `json:"input"` // JSON string of input parameters
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Input string `json:"input"` // JSON string of input parameters
|
||||
RawInput string `json:"raw_input,omitempty"` // Unmodified JSON input
|
||||
}
|
||||
|
||||
// ToolEndData is the data for "tool_end" events
|
||||
type ToolEndData struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Input string `json:"input,omitempty"`
|
||||
Output string `json:"output,omitempty"`
|
||||
Success bool `json:"success"`
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Input string `json:"input,omitempty"`
|
||||
RawInput string `json:"raw_input,omitempty"`
|
||||
Output string `json:"output,omitempty"`
|
||||
Success bool `json:"success"`
|
||||
}
|
||||
|
||||
// ApprovalNeededData is the data for "approval_needed" events
|
||||
|
||||
Reference in New Issue
Block a user