From badbad4464c35f51ad034ce3cbada9895f766f84 Mon Sep 17 00:00:00 2001 From: rcourtman Date: Wed, 28 Jan 2026 21:21:41 +0000 Subject: [PATCH] refactor(ai): integrate patrol execution into chat service - Add ExecutePatrolStream method to chat.Service for patrol-specific execution - Create chat_service_adapter.go to bridge chat.Service to ai.ChatServiceProvider - Remove standalone patrol.go and patrol_test.go from chat package - Add PatrolRequest/PatrolResponse types to chat service - Add context injection for recent message context This allows patrol to use an isolated agentic loop with its own system prompt while leveraging the common chat infrastructure. --- internal/ai/chat/agentic.go | 150 +++++++++- internal/ai/chat/service.go | 313 ++++++++++++++++++++ internal/ai/chat/service_additional_test.go | 32 +- internal/ai/chat/types.go | 38 +++ internal/api/chat_service_adapter.go | 80 +++++ 5 files changed, 578 insertions(+), 35 deletions(-) create mode 100644 internal/api/chat_service_adapter.go diff --git a/internal/ai/chat/agentic.go b/internal/ai/chat/agentic.go index c30a516ac..2ff0263b6 100644 --- a/internal/ai/chat/agentic.go +++ b/internal/ai/chat/agentic.go @@ -27,6 +27,10 @@ type AgenticLoop struct { providerName string modelName string + // Token accumulation across all turns + totalInputTokens int + totalOutputTokens int + // State for ongoing executions mu sync.Mutex aborted map[string]bool // sessionID -> aborted @@ -92,6 +96,8 @@ func (a *AgenticLoop) executeWithTools(ctx context.Context, sessionID string, me var resultMessages []Message turn := 0 toolsSucceededThisEpisode := false // Track if any tool executed successfully this episode + preferredToolName := "" + preferredToolRetried := false for turn < a.maxTurns { // Check if aborted @@ -131,21 +137,29 @@ func (a *AgenticLoop) executeWithTools(ctx context.Context, sessionID string, me Tools: tools, } - // Determine tool_choice based on turn and intent + // Determine tool_choice based on turn, intent, and explicit tool requests. // We only force tool use when: // 1. Tools are available // 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 len(tools) > 0 { - if turn == 0 && requiresToolUse(providerMessages) { + if preferredToolName == "" { + preferredToolName = getPreferredTool(providerMessages, tools) + } + if preferredToolName != "" { + req.ToolChoice = &providers.ToolChoice{Type: providers.ToolChoiceTool, Name: preferredToolName} + log.Debug(). + Str("session_id", sessionID). + Str("tool", preferredToolName). + Msg("[AgenticLoop] Explicit tool request - forcing tool") + } else if turn == 0 && requiresToolUse(providerMessages) { // First turn with action intent: force the model to use a tool req.ToolChoice = &providers.ToolChoice{Type: providers.ToolChoiceAny} log.Debug(). Str("session_id", sessionID). Msg("[AgenticLoop] First turn with action intent - forcing tool use") } else { - // Conceptual questions or subsequent turns: let the model decide req.ToolChoice = &providers.ToolChoice{Type: providers.ToolChoiceAuto} if turn == 0 { log.Debug(). @@ -227,6 +241,8 @@ func (a *AgenticLoop) executeWithTools(ctx context.Context, sessionID string, me case "done": if data, ok := event.Data.(providers.DoneEvent); ok { toolCalls = data.ToolCalls + a.totalInputTokens += data.InputTokens + a.totalOutputTokens += data.OutputTokens } case "error": @@ -293,6 +309,19 @@ func (a *AgenticLoop) executeWithTools(ctx context.Context, sessionID string, me // If no tool calls, we're done - but first check FSM and phantom execution if len(toolCalls) == 0 { + // If the user explicitly requested a tool and the model didn't comply, retry once. + if preferredToolName != "" && !preferredToolRetried { + preferredToolRetried = true + + retryPrompt := fmt.Sprintf("Tool required: use %s for this request.", preferredToolName) + if len(resultMessages) > 0 { + resultMessages[len(resultMessages)-1].Content = retryPrompt + } + + turn++ + continue + } + // === FSM ENFORCEMENT GATE 2: Check if final answer is allowed === a.mu.Lock() fsm := a.sessionFSM @@ -386,6 +415,10 @@ func (a *AgenticLoop) executeWithTools(ctx context.Context, sessionID string, me } // Execute tool calls + if len(toolCalls) > 0 && preferredToolName != "" { + // Clear preferred tool once the model has used any tool. + preferredToolName = "" + } for _, tc := range toolCalls { // Check for abort a.mu.Lock() @@ -775,6 +808,22 @@ func (a *AgenticLoop) SetProviderInfo(provider, model string) { a.mu.Unlock() } +// GetTotalInputTokens returns the accumulated input tokens across all turns. +func (a *AgenticLoop) GetTotalInputTokens() int { + return a.totalInputTokens +} + +// GetTotalOutputTokens returns the accumulated output tokens across all turns. +func (a *AgenticLoop) GetTotalOutputTokens() int { + return a.totalOutputTokens +} + +// ResetTokenCounts resets the accumulated token counts (for reuse across executions). +func (a *AgenticLoop) ResetTokenCounts() { + a.totalInputTokens = 0 + a.totalOutputTokens = 0 +} + // hasPhantomExecution detects when the model claims to have executed something // but no actual tool calls were made. This catches models that "hallucinate" // tool execution by writing about it instead of calling tools. @@ -976,8 +1025,30 @@ func requiresToolUse(messages []providers.Message) bool { strings.Contains(lastUserContent, "usage") if hasMyInfra && hasStateQuery { - break // Not conceptual, continue to action detection + return true // Explicit state query about user's infrastructure } + + // Exception: explicit resource references should trigger tools even in "tell me about" queries. + resourceNouns := []string{ + "container", "vm", "lxc", "node", "pod", "deployment", "service", "host", "cluster", + } + hasResourceNoun := false + for _, noun := range resourceNouns { + if strings.Contains(lastUserContent, noun) { + hasResourceNoun = true + break + } + } + explicitIndicator := strings.Contains(lastUserContent, "@") || + strings.Contains(lastUserContent, "\"") || + strings.Contains(lastUserContent, "-") || + strings.Contains(lastUserContent, "_") || + strings.Contains(lastUserContent, "/") + + if hasResourceNoun && explicitIndicator { + return true // Treat as action: specific resource is referenced + } + return false } } @@ -1011,6 +1082,10 @@ func requiresToolUse(messages []providers.Message) bool { // Questions about "my" specific infrastructure "my server", "my container", "my vm", "my host", "my infrastructure", "my node", "my cluster", "my proxmox", "my docker", + // Inventory-style queries + "what nodes do i have", "what proxmox nodes", + "what containers do i have", "what vms do i have", + "what is running on", "what's running on", } for _, pattern := range actionPatterns { @@ -1019,10 +1094,77 @@ func requiresToolUse(messages []providers.Message) bool { } } + // Logs or journal queries should always hit tools. + if strings.Contains(lastUserContent, "logs") || + strings.Contains(lastUserContent, " log") || + strings.Contains(lastUserContent, "journal") || + strings.Contains(lastUserContent, "journald") { + return true + } + // Default: assume conceptual question, don't force tools return false } +// getPreferredTool returns a tool name if the user explicitly requested one. +// Only returns tools that are available for this request. +func getPreferredTool(messages []providers.Message, tools []providers.Tool) string { + var lastUserContent string + for i := len(messages) - 1; i >= 0; i-- { + if messages[i].Role == "user" && messages[i].ToolResult == nil { + lastUserContent = strings.ToLower(messages[i].Content) + break + } + } + if lastUserContent == "" { + return "" + } + + toolSet := make(map[string]bool, len(tools)) + for _, tool := range tools { + if tool.Name != "" { + toolSet[tool.Name] = true + } + } + + // Explicit tool mentions + explicitTools := []string{ + "pulse_read", + "pulse_control", + "pulse_query", + "pulse_discovery", + "pulse_docker", + "pulse_kubernetes", + "pulse_metrics", + "pulse_storage", + } + for _, tool := range explicitTools { + if strings.Contains(lastUserContent, tool) && toolSet[tool] { + return tool + } + } + + // Natural language aliases + if (strings.Contains(lastUserContent, "read-only tool") || strings.Contains(lastUserContent, "read only tool")) && toolSet["pulse_read"] { + return "pulse_read" + } + if strings.Contains(lastUserContent, "control tool") && toolSet["pulse_control"] { + return "pulse_control" + } + if strings.Contains(lastUserContent, "query tool") && toolSet["pulse_query"] { + return "pulse_query" + } + + // Context carryover: if we injected an explicit target and logs are requested, force pulse_read. + if strings.Contains(lastUserContent, "explicit target") && + (strings.Contains(lastUserContent, "log") || strings.Contains(lastUserContent, "journal")) && + toolSet["pulse_read"] { + return "pulse_read" + } + + return "" +} + // getSystemPrompt builds the full system prompt including the current mode context. // This is called at request time so the prompt reflects the current mode. func (a *AgenticLoop) getSystemPrompt() string { diff --git a/internal/ai/chat/service.go b/internal/ai/chat/service.go index 6cd040160..1d5dc5360 100644 --- a/internal/ai/chat/service.go +++ b/internal/ai/chat/service.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "fmt" + "regexp" "sort" "strings" "sync" @@ -326,6 +327,10 @@ func (s *Service) ExecuteStream(ctx context.Context, req ExecuteRequest, callbac if prefetcher != nil { prefetchCtx := prefetcher.Prefetch(ctx, req.Prompt) + mentionsFound := false + if prefetchCtx != nil { + mentionsFound = len(prefetchCtx.Mentions) > 0 + } if prefetchCtx != nil && prefetchCtx.Summary != "" { log.Info(). Int("mentions", len(prefetchCtx.Mentions)). @@ -372,6 +377,12 @@ func (s *Service) ExecuteStream(ctx context.Context, req ExecuteRequest, callbac } } } + + if !mentionsFound { + s.injectRecentContextIfNeeded(req.Prompt, session.ID, messages, sessions) + } + } else { + s.injectRecentContextIfNeeded(req.Prompt, session.ID, messages, sessions) } // Run agentic loop @@ -455,6 +466,204 @@ func (s *Service) ExecuteStream(ctx context.Context, req ExecuteRequest, callbac return nil } +// PatrolRequest represents a patrol execution request within the chat service +type PatrolRequest struct { + Prompt string `json:"prompt"` + SystemPrompt string `json:"system_prompt"` + SessionID string `json:"session_id,omitempty"` + UseCase string `json:"use_case"` +} + +// PatrolResponse contains the results of a patrol execution +type PatrolResponse struct { + Content string `json:"content"` + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` +} + +// ExecutePatrolStream creates a temporary agentic loop for patrol execution. +// Unlike ExecuteStream (which uses the shared agentic loop with the chat system prompt), +// this creates an isolated loop with the patrol's own system prompt and model. +func (s *Service) ExecutePatrolStream(ctx context.Context, req PatrolRequest, callback StreamCallback) (*PatrolResponse, error) { + log.Debug(). + Str("session_id", req.SessionID). + Int("prompt_len", len(req.Prompt)). + Msg("[ChatService] ExecutePatrolStream called") + + s.mu.RLock() + if !s.started { + s.mu.RUnlock() + return nil, fmt.Errorf("service not started") + } + sessions := s.sessions + executor := s.executor + cfg := s.cfg + s.mu.RUnlock() + + // Determine model: use patrol model or fall back to chat model + patrolModel := "" + if cfg != nil { + patrolModel = cfg.GetPatrolModel() + if patrolModel == "" { + patrolModel = cfg.GetChatModel() + } + } + if patrolModel == "" { + return nil, fmt.Errorf("no patrol model configured") + } + + // Create a temporary provider for the patrol model + provider, err := s.createProviderForModel(patrolModel) + if err != nil { + return nil, fmt.Errorf("failed to create patrol provider: %w", err) + } + + // Create a temporary agentic loop with the patrol system prompt + systemPrompt := req.SystemPrompt + if systemPrompt == "" { + systemPrompt = s.buildSystemPrompt() + } + tempLoop := NewAgenticLoop(provider, executor, systemPrompt) + tempLoop.SetAutonomousMode(true) // Patrol runs without approval prompts + + // Set provider info for telemetry + parts := strings.SplitN(patrolModel, ":", 2) + if len(parts) == 2 { + tempLoop.SetProviderInfo(parts[0], parts[1]) + } + + // Ensure patrol session exists + sessionID := req.SessionID + if sessionID == "" { + sessionID = "patrol-main" + } + session, err := sessions.EnsureSession(sessionID) + if err != nil { + return nil, fmt.Errorf("failed to ensure patrol session: %w", err) + } + + // Set resolved context on executor (same as ExecuteStream does) + if executor != nil { + resolvedCtx := sessions.GetResolvedContext(session.ID) + executor.SetResolvedContext(resolvedCtx) + } + + // Set session FSM + sessionFSM := sessions.GetSessionFSM(session.ID) + tempLoop.SetSessionFSM(sessionFSM) + + // Add user message + userMsg := Message{ + ID: uuid.New().String(), + Role: "user", + Content: req.Prompt, + Timestamp: time.Now(), + } + if err := sessions.AddMessage(session.ID, userMsg); err != nil { + log.Warn().Err(err).Msg("Failed to save patrol user message") + } + + // Get messages for context + messages, err := sessions.GetMessages(session.ID) + if err != nil { + return nil, fmt.Errorf("failed to get patrol messages: %w", err) + } + + // Get all tools + filteredTools := s.filterToolsForPrompt(ctx, req.Prompt) + + // Run the agentic loop + resultMessages, err := tempLoop.ExecuteWithTools(ctx, session.ID, messages, filteredTools, callback) + if err != nil { + // Still save any messages we got + for _, msg := range resultMessages { + if saveErr := sessions.AddMessage(session.ID, msg); saveErr != nil { + log.Warn().Err(saveErr).Msg("Failed to save patrol message after error") + } + } + return nil, err + } + + // Save result messages + for _, msg := range resultMessages { + if msg.Role == "user" && msg.ToolResult == nil { + continue + } + if err := sessions.AddMessage(session.ID, msg); err != nil { + log.Warn().Err(err).Msg("Failed to save patrol message") + } + } + + // Collect content from result messages + var contentBuilder strings.Builder + for _, msg := range resultMessages { + if msg.Role == "assistant" && msg.Content != "" { + contentBuilder.WriteString(msg.Content) + } + } + + // Send done event + doneData, _ := json.Marshal(DoneData{ + SessionID: session.ID, + InputTokens: tempLoop.GetTotalInputTokens(), + OutputTokens: tempLoop.GetTotalOutputTokens(), + }) + callback(StreamEvent{Type: "done", Data: doneData}) + + return &PatrolResponse{ + Content: contentBuilder.String(), + InputTokens: tempLoop.GetTotalInputTokens(), + OutputTokens: tempLoop.GetTotalOutputTokens(), + }, nil +} + +// createProviderForModel creates a streaming provider for a specific model string (provider:model format). +func (s *Service) createProviderForModel(modelStr string) (providers.StreamingProvider, error) { + if s.cfg == nil { + return nil, fmt.Errorf("no Pulse Assistant config") + } + + parts := strings.SplitN(modelStr, ":", 2) + if len(parts) != 2 { + return nil, fmt.Errorf("invalid model format: %s (expected provider:model)", modelStr) + } + providerName := parts[0] + modelName := parts[1] + + timeout := 5 * time.Minute + + switch providerName { + case "anthropic": + if s.cfg.AnthropicAPIKey == "" { + return nil, fmt.Errorf("Anthropic API key not configured") + } + return providers.NewAnthropicClient(s.cfg.AnthropicAPIKey, modelName, timeout), nil + case "openai": + if s.cfg.OpenAIAPIKey == "" { + return nil, fmt.Errorf("OpenAI API key not configured") + } + return providers.NewOpenAIClient(s.cfg.OpenAIAPIKey, modelName, "", timeout), nil + case "deepseek": + if s.cfg.DeepSeekAPIKey == "" { + return nil, fmt.Errorf("DeepSeek API key not configured") + } + return providers.NewOpenAIClient(s.cfg.DeepSeekAPIKey, modelName, "https://api.deepseek.com", timeout), nil + case "gemini": + if s.cfg.GeminiAPIKey == "" { + return nil, fmt.Errorf("Gemini API key not configured") + } + return providers.NewGeminiClient(s.cfg.GeminiAPIKey, modelName, "", timeout), nil + case "ollama": + baseURL := s.cfg.OllamaBaseURL + if baseURL == "" { + baseURL = "http://localhost:11434" + } + return providers.NewOllamaClient(modelName, baseURL, timeout), nil + default: + return nil, fmt.Errorf("unsupported provider: %s", providerName) + } +} + // ListAvailableTools returns tool names available for the given prompt. func (s *Service) ListAvailableTools(ctx context.Context, prompt string) []string { s.mu.RLock() @@ -928,6 +1137,110 @@ DOCKER BIND MOUNTS: - Use pulse_discovery to find bind mount mappings` } +var recentContextPronounPattern = regexp.MustCompile(`(?i)\b(it|its|that|those|this|them|previous|earlier|last|same|former|latter)\b`) +var recentContextNounPattern = regexp.MustCompile(`(?i)\b(the (service|container|vm|lxc|node|host|docker|instance|one))\b`) + +func shouldInjectRecentContext(prompt string) bool { + return recentContextPronounPattern.MatchString(prompt) || recentContextNounPattern.MatchString(prompt) +} + +func (s *Service) injectRecentContextIfNeeded(prompt, sessionID string, messages []Message, sessions *SessionStore) { + if !shouldInjectRecentContext(prompt) { + return + } + + if sessions == nil { + return + } + resolvedCtx := sessions.GetResolvedContext(sessionID) + if resolvedCtx == nil { + return + } + + recentIDs := resolvedCtx.GetRecentlyAccessedResourcesSorted(tools.RecentAccessWindow, 3) + if len(recentIDs) == 0 { + return + } + + var lines []string + primaryName := "" + primaryTarget := "" + for _, resourceID := range recentIDs { + res, ok := resolvedCtx.GetResourceByID(resourceID) + if !ok || res == nil { + continue + } + label := res.Name + if label == "" { + label = resourceID + } + kind := res.Kind + if kind == "" { + kind = res.ResourceType + } + location := res.Node + if location == "" { + location = res.Scope.HostName + } + if kind != "" && location != "" { + label = fmt.Sprintf("%s (%s on %s)", label, kind, location) + } else if kind != "" { + label = fmt.Sprintf("%s (%s)", label, kind) + } + lines = append(lines, "- "+label) + if primaryName == "" { + primaryName = res.Name + primaryTarget = res.TargetHost + if primaryName == "" { + primaryName = label + } + } + } + + if len(lines) == 0 { + return + } + + primary := strings.TrimPrefix(lines[0], "- ") + if primaryName == "" { + primaryName = primary + } + targetHint := "" + if primaryTarget != "" { + targetHint = fmt.Sprintf(" Use target_host=\"%s\".", primaryTarget) + } + summary := fmt.Sprintf("Context: The most recently referenced resource is %s. If the user says \"it/its/that\", assume they mean this resource unless they specify otherwise. Do not ask for clarification unless the user names a different resource.%s", primary, targetHint) + if len(lines) > 1 { + others := strings.Join(lines[1:], "\n") + summary += "\nOther recent resources:\n" + others + } + + lowerPrompt := strings.ToLower(prompt) + if strings.Contains(lowerPrompt, "log") || strings.Contains(lowerPrompt, "journal") { + rewrite := fmt.Sprintf("Show logs for %s (last 50 lines).", primaryName) + if primaryTarget != "" { + summary += fmt.Sprintf("\nInstruction: %s Use pulse_read action=logs target_host=\"%s\" lines=50.", rewrite, primaryTarget) + } else { + summary += fmt.Sprintf("\nInstruction: %s Use pulse_read action=logs target_host=\"%s\" lines=50.", rewrite, primaryName) + } + } + + log.Debug(). + Str("session_id", sessionID). + Strs("recent_resource_ids", recentIDs). + Msg("[ChatService] Injecting recent context") + + if len(messages) == 0 { + return + } + lastIdx := len(messages) - 1 + if messages[lastIdx].Role != "user" { + return + } + + messages[lastIdx].Content = summary + "\n\n---\nExplicit target: " + primaryName + "\nUser question (targeted): " + messages[lastIdx].Content +} + func (s *Service) filterToolsForPrompt(ctx context.Context, prompt string) []providers.Tool { mcpTools := s.executor.ListTools() providerTools := ConvertMCPToolsToProvider(mcpTools) diff --git a/internal/ai/chat/service_additional_test.go b/internal/ai/chat/service_additional_test.go index 10de6be44..9790aeb96 100644 --- a/internal/ai/chat/service_additional_test.go +++ b/internal/ai/chat/service_additional_test.go @@ -37,34 +37,4 @@ func TestServiceExecuteCommand_NoExecutor(t *testing.T) { } } -func TestPatrolServiceSessionLifecycle(t *testing.T) { - store, err := NewSessionStore(t.TempDir()) - if err != nil { - t.Fatalf("NewSessionStore error: %v", err) - } - - service := &Service{ - sessions: store, - started: true, - } - - patrol := NewPatrolService(service) - if err := patrol.CreatePatrolSession(context.Background()); err != nil { - t.Fatalf("CreatePatrolSession error: %v", err) - } - if patrol.GetSessionID() == "" { - t.Fatalf("expected session ID to be set") - } - - patrol.mu.Lock() - patrol.running = true - patrol.mu.Unlock() - if !patrol.IsRunning() { - t.Fatalf("expected patrol to be running") - } - - service.started = false - if err := patrol.CreatePatrolSession(context.Background()); err == nil { - t.Fatalf("expected error when service not running") - } -} +// TestPatrolServiceSessionLifecycle was removed: it tested the deleted chat/patrol.go bridge. diff --git a/internal/ai/chat/types.go b/internal/ai/chat/types.go index 817a2e4cb..4129c995c 100644 --- a/internal/ai/chat/types.go +++ b/internal/ai/chat/types.go @@ -3,6 +3,7 @@ package chat import ( "encoding/json" + "sort" "strings" "time" @@ -603,6 +604,43 @@ func (rc *ResolvedContext) GetRecentlyAccessedResources(window time.Duration) [] return recent } +// GetRecentlyAccessedResourcesSorted returns recently accessed resources ordered by most recent access. +// If max <= 0, all recent resources are returned. +func (rc *ResolvedContext) GetRecentlyAccessedResourcesSorted(window time.Duration, max int) []string { + if rc.explicitlyAccessed == nil { + return nil + } + + cutoff := time.Now().Add(-window) + type recentResource struct { + id string + ts time.Time + } + var recent []recentResource + + for resourceID, explicitAccess := range rc.explicitlyAccessed { + if explicitAccess.After(cutoff) { + if _, ok := rc.ResourcesByID[resourceID]; ok { + recent = append(recent, recentResource{id: resourceID, ts: explicitAccess}) + } + } + } + + sort.Slice(recent, func(i, j int) bool { + return recent[i].ts.After(recent[j].ts) + }) + + if max > 0 && len(recent) > max { + recent = recent[:max] + } + + ids := make([]string, len(recent)) + for i, res := range recent { + ids[i] = res.id + } + return ids +} + // AddResource adds a resolved resource to the context (internal use). // NOTE: This does NOT mark the resource as "recently accessed" for routing validation. // Use MarkExplicitAccess() when the user explicitly selects/queries a single resource. diff --git a/internal/api/chat_service_adapter.go b/internal/api/chat_service_adapter.go new file mode 100644 index 000000000..da15211e1 --- /dev/null +++ b/internal/api/chat_service_adapter.go @@ -0,0 +1,80 @@ +package api + +import ( + "context" + "encoding/json" + + "github.com/rcourtman/pulse-go-rewrite/internal/ai" + "github.com/rcourtman/pulse-go-rewrite/internal/ai/chat" +) + +// chatServiceAdapter wraps chat.Service to implement ai.ChatServiceProvider. +// This bridges the chat package (concrete) to the ai package (interface) without +// creating an import cycle. +type chatServiceAdapter struct { + svc *chat.Service +} + +func (a *chatServiceAdapter) CreateSession(ctx context.Context) (*ai.ChatSession, error) { + session, err := a.svc.CreateSession(ctx) + if err != nil { + return nil, err + } + return &ai.ChatSession{ID: session.ID}, nil +} + +func (a *chatServiceAdapter) ExecuteStream(ctx context.Context, req ai.ChatExecuteRequest, callback ai.ChatStreamCallback) error { + return a.svc.ExecuteStream(ctx, chat.ExecuteRequest{ + Prompt: req.Prompt, + SessionID: req.SessionID, + }, adaptCallback(callback)) +} + +func (a *chatServiceAdapter) ExecutePatrolStream(ctx context.Context, req ai.PatrolExecuteRequest, callback ai.ChatStreamCallback) (*ai.PatrolStreamResponse, error) { + resp, err := a.svc.ExecutePatrolStream(ctx, chat.PatrolRequest{ + Prompt: req.Prompt, + SystemPrompt: req.SystemPrompt, + SessionID: req.SessionID, + UseCase: req.UseCase, + }, adaptCallback(callback)) + if err != nil { + return nil, err + } + return &ai.PatrolStreamResponse{ + Content: resp.Content, + InputTokens: resp.InputTokens, + OutputTokens: resp.OutputTokens, + }, nil +} + +func (a *chatServiceAdapter) GetMessages(ctx context.Context, sessionID string) ([]ai.ChatMessage, error) { + messages, err := a.svc.GetMessages(ctx, sessionID) + if err != nil { + return nil, err + } + result := make([]ai.ChatMessage, len(messages)) + for i, m := range messages { + result[i] = ai.ChatMessage{ + ID: m.ID, + Role: m.Role, + Content: m.Content, + Timestamp: m.Timestamp, + } + } + return result, nil +} + +func (a *chatServiceAdapter) DeleteSession(ctx context.Context, sessionID string) error { + return a.svc.DeleteSession(ctx, sessionID) +} + +// adaptCallback converts an ai.ChatStreamCallback to a chat.StreamCallback. +// The ai package uses []byte for event data, while the chat package uses json.RawMessage. +func adaptCallback(callback ai.ChatStreamCallback) chat.StreamCallback { + return func(event chat.StreamEvent) { + callback(ai.ChatStreamEvent{ + Type: event.Type, + Data: json.RawMessage(event.Data), + }) + } +}