From 00802f78bec4cbad0e071d53f6fa0fe37c928a69 Mon Sep 17 00:00:00 2001 From: ZanzyTHEbar Date: Sat, 21 Feb 2026 18:59:19 +0000 Subject: [PATCH] feat(tools): add agentic_map and llm_map operators with FlatBuffers persistence pkg/tools/agentic_map.go + agentic_map_test.go - agentic_map: batch operator that dispatches each item to a subagent worker; supports configurable concurrency, retries, and timeout per item - Items are processed in parallel up to MaxWorkers; failed items are retried up to MaxRetries with exponential backoff - Results are collected into a structured MapResult with per-item status pkg/tools/llm_map.go + llm_map_test.go - llm_map: batch operator that applies a prompt template to each item using a direct LLM call (no subagent overhead) - Supports Jinja-style {{item}} template substitution and structured JSON output extraction from LLM responses pkg/tools/map_runtime.go + map_runtime_integration_test.go - MapRuntime: manages the lifecycle of a map operator run (create, update, complete, fail); persists run state via memory delegate - Tracks per-item status transitions and aggregates run-level metrics pkg/tools/map_run_tools.go - map_run_status: returns the current status and progress of a running map operation; enables the agent to monitor long-running batch jobs - map_run_cancel: cancels an in-progress map run pkg/tools/map_boundary.go - MapBoundary: defines the input/output contract for map operators; validates item schemas and enforces output format requirements pkg/tools/map_payloads.fbs - FlatBuffers schema for MapRunSpec, MapItemInput, MapItemOutput; enables zero-copy serialization of map operator payloads pkg/tools/mapopsfb/MapRunSpec.go pkg/tools/mapopsfb/MapItemInput.go pkg/tools/mapopsfb/MapItemOutput.go - Generated FlatBuffers Go bindings for map operator payload types pkg/tools/map_flatbuffer_codec.go + map_flatbuffer_codec_test.go - MapFlatbufferCodec: bidirectional codec between Go map payload structs and FlatBuffers wire format; used by MapRuntime for persistence --- pkg/tools/agentic_map.go | 334 +++++++++ pkg/tools/agentic_map_test.go | 113 +++ pkg/tools/llm_map.go | 327 +++++++++ pkg/tools/llm_map_test.go | 111 +++ pkg/tools/map_boundary.go | 204 ++++++ pkg/tools/map_flatbuffer_codec.go | 230 ++++++ pkg/tools/map_flatbuffer_codec_test.go | 83 +++ pkg/tools/map_payloads.fbs | 33 + pkg/tools/map_run_tools.go | 216 ++++++ pkg/tools/map_runtime.go | 833 ++++++++++++++++++++++ pkg/tools/map_runtime_integration_test.go | 383 ++++++++++ pkg/tools/mapopsfb/MapItemInput.go | 93 +++ pkg/tools/mapopsfb/MapItemOutput.go | 134 ++++ pkg/tools/mapopsfb/MapRunSpec.go | 159 +++++ 14 files changed, 3253 insertions(+) create mode 100644 pkg/tools/agentic_map.go create mode 100644 pkg/tools/agentic_map_test.go create mode 100644 pkg/tools/llm_map.go create mode 100644 pkg/tools/llm_map_test.go create mode 100644 pkg/tools/map_boundary.go create mode 100644 pkg/tools/map_flatbuffer_codec.go create mode 100644 pkg/tools/map_flatbuffer_codec_test.go create mode 100644 pkg/tools/map_payloads.fbs create mode 100644 pkg/tools/map_run_tools.go create mode 100644 pkg/tools/map_runtime.go create mode 100644 pkg/tools/map_runtime_integration_test.go create mode 100644 pkg/tools/mapopsfb/MapItemInput.go create mode 100644 pkg/tools/mapopsfb/MapItemOutput.go create mode 100644 pkg/tools/mapopsfb/MapRunSpec.go diff --git a/pkg/tools/agentic_map.go b/pkg/tools/agentic_map.go new file mode 100644 index 000000000..3a0615485 --- /dev/null +++ b/pkg/tools/agentic_map.go @@ -0,0 +1,334 @@ +package tools + +import ( + "context" + "fmt" + "strconv" + "strings" + + memsqlc "github.com/ZanzyTHEbar/dragonscale/pkg/memory/sqlc" + jsonv2 "github.com/go-json-experiment/json" +) + +// AgenticMapTool maps tasks over items using subagent execution with retries. +type AgenticMapTool struct { + manager *SubagentManager + originChannel string + originChatID string + runtime *MapRuntime +} + +func NewAgenticMapTool(manager *SubagentManager) *AgenticMapTool { + return &AgenticMapTool{ + manager: manager, + originChannel: "cli", + originChatID: "direct", + } +} + +func (t *AgenticMapTool) SetRuntime(runtime *MapRuntime) { + t.runtime = runtime +} + +func (t *AgenticMapTool) Name() string { + return "agentic_map" +} + +func (t *AgenticMapTool) Description() string { + return "Run subagent processing over each item with retries. Supports worker-backed runs and generated run reads." +} + +func (t *AgenticMapTool) Parameters() map[string]interface{} { + return map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "items": map[string]interface{}{ + "type": "array", + "description": "Items to process via subagent execution.", + }, + "input_jsonl": map[string]interface{}{ + "type": "string", + "description": "Boundary-only JSONL input. Exactly one of items or input_jsonl is required.", + }, + "task_template": map[string]interface{}{ + "type": "string", + "description": "Task template. Supports {{item_json}} and {{index}} placeholders.", + }, + "max_retries": map[string]interface{}{ + "type": "integer", + "description": "Retries per item after the first attempt (default 1).", + }, + "delegated_scope": map[string]interface{}{ + "type": "string", + "description": "Delegated scope metadata passed to nested subagent calls.", + }, + "kept_work": map[string]interface{}{ + "type": "string", + "description": "Kept work metadata passed to nested subagent calls.", + }, + "execution_mode": map[string]interface{}{ + "type": "string", + "description": "inline or worker. Defaults to worker for JSONL/large batches, inline otherwise.", + }, + "session_key": map[string]interface{}{ + "type": "string", + "description": "Optional map run session key for persistence (default: default).", + }, + "idempotency_key": map[string]interface{}{ + "type": "string", + "description": "Optional key to deduplicate repeated run creation.", + }, + }, + "required": []string{"task_template"}, + } +} + +func (t *AgenticMapTool) SetContext(channel, chatID string) { + t.originChannel = channel + t.originChatID = chatID +} + +func (t *AgenticMapTool) Execute(ctx context.Context, args map[string]interface{}) *ToolResult { + if t.manager == nil && t.runtime == nil { + return ErrorResult("agentic_map manager is not configured").WithError(fmt.Errorf("agentic_map manager is nil")) + } + + items, usedJSONL, err := parseMapBoundaryItems(args) + if err != nil { + return ErrorResult(err.Error()).WithError(err) + } + + taskTemplate, ok := args["task_template"].(string) + if !ok || strings.TrimSpace(taskTemplate) == "" { + return ErrorResult("task_template is required").WithError(fmt.Errorf("task_template is required")) + } + if !strings.Contains(taskTemplate, "{{item_json}}") || !strings.Contains(taskTemplate, "{{index}}") { + return ErrorResult("task_template must include both {{item_json}} and {{index}} placeholders"). + WithError(fmt.Errorf("task_template missing required placeholders")) + } + + maxRetries, err := parseMapMaxRetries(args, 1) + if err != nil { + return ErrorResult(err.Error()).WithError(err) + } + delegatedScope, err := parseOptionalStringArg(args, "delegated_scope") + if err != nil { + return ErrorResult(err.Error()).WithError(err) + } + keptWork, err := parseOptionalStringArg(args, "kept_work") + if err != nil { + return ErrorResult(err.Error()).WithError(err) + } + mode, err := resolveMapExecutionMode(args, len(items), usedJSONL) + if err != nil { + return ErrorResult(err.Error()).WithError(err) + } + sessionKey, err := parseOptionalStringArg(args, "session_key") + if err != nil { + return ErrorResult(err.Error()).WithError(err) + } + idempotencyKey, err := parseOptionalStringArg(args, "idempotency_key") + if err != nil { + return ErrorResult(err.Error()).WithError(err) + } + if sessionKey == "" { + sessionKey = "default" + } + + if t.runtime != nil { + run, reused, err := t.runtime.EnqueueRun(ctx, sessionKey, MapRunSpec{ + OperatorKind: MapOperatorAgentic, + TaskTemplate: taskTemplate, + MaxRetries: uint16(maxRetries), + DelegatedScope: delegatedScope, + KeptWork: keptWork, + OriginChannel: t.originChannel, + OriginChatID: t.originChatID, + }, items, idempotencyKey) + if err != nil { + return ErrorResult(fmt.Sprintf("failed to enqueue agentic_map run: %v", err)).WithError(err) + } + if mode == "worker" { + payload := map[string]interface{}{ + "run_id": run.ID.String(), + "status": run.Status, + "accepted_count": len(items), + "queued_count": run.QueuedItems, + "execution_mode": "worker", + "idempotent_reuse": reused, + } + data, err := jsonv2.Marshal(payload) + if err != nil { + return ErrorResult("failed to serialize agentic_map enqueue result").WithError(err) + } + return &ToolResult{ + ForLLM: string(data), + ForUser: string(data), + Silent: false, + IsError: false, + Async: false, + } + } + runFinal, err := t.runtime.ProcessRunToTerminal(ctx, run.ID, len(items)*(maxRetries+2)+32) + if err != nil { + return ErrorResult(fmt.Sprintf("failed to complete inline agentic_map run %s: %v", run.ID.String(), err)).WithError(err) + } + itemRows, err := t.runtime.ReadRunItems(ctx, run.ID, 0, len(items)) + if err != nil { + return ErrorResult(fmt.Sprintf("failed to read inline agentic_map results: %v", err)).WithError(err) + } + return buildAgenticMapInlineResult(runFinal, itemRows, maxRetries) + } + + // Legacy synchronous path when runtime persistence is not configured. + if mode == "worker" { + return ErrorResult("worker mode requires map runtime persistence").WithError(fmt.Errorf("map runtime is nil")) + } + if t.manager == nil { + return ErrorResult("agentic_map manager is not configured").WithError(fmt.Errorf("agentic_map manager is nil")) + } + + subTool := NewSubagentTool(t.manager) + subTool.SetContext(t.originChannel, t.originChatID) + + type itemResult struct { + Index int `json:"index"` + Success bool `json:"success"` + Attempts int `json:"attempts"` + Output string `json:"output,omitempty"` + Error string `json:"error,omitempty"` + } + results := make([]itemResult, 0, len(items)) + successCount := 0 + + for i, item := range items { + if err := ctx.Err(); err != nil { + return ErrorResult("agentic_map cancelled before completion").WithError(err) + } + + task := strings.ReplaceAll(taskTemplate, "{{item_json}}", item.ItemJSON) + task = strings.ReplaceAll(task, "{{index}}", strconv.Itoa(i)) + + attempts := 0 + var lastErr string + var output string + okItem := false + for attempts <= maxRetries { + if err := ctx.Err(); err != nil { + lastErr = err.Error() + break + } + attempts++ + callArgs := map[string]interface{}{ + "task": task, + "label": fmt.Sprintf("agentic-map-%d", i), + } + if delegatedScope != "" { + callArgs["delegated_scope"] = delegatedScope + } + if keptWork != "" { + callArgs["kept_work"] = keptWork + } + + r := subTool.Execute(ctx, callArgs) + if r != nil && !r.IsError { + okItem = true + output = r.ForUser + break + } + if r != nil { + lastErr = r.ForLLM + } else { + lastErr = "unknown subagent error" + } + } + + if okItem { + successCount++ + } + results = append(results, itemResult{ + Index: i, + Success: okItem, + Attempts: attempts, + Output: output, + Error: lastErr, + }) + } + + payload := map[string]interface{}{ + "count": len(results), + "summary": map[string]interface{}{ + "success_count": successCount, + "failure_count": len(results) - successCount, + "max_retries": maxRetries, + }, + "results": results, + } + data, err := jsonv2.Marshal(payload) + if err != nil { + return ErrorResult("failed to serialize agentic_map results").WithError(err) + } + return &ToolResult{ + ForLLM: string(data), + ForUser: string(data), + Silent: false, + IsError: false, + Async: false, + } +} + +func buildAgenticMapInlineResult(run memsqlc.MapRun, items []MapItemProjection, maxRetries int) *ToolResult { + type itemResult struct { + Index int64 `json:"index"` + Success bool `json:"success"` + Attempts int64 `json:"attempts"` + Output string `json:"output,omitempty"` + Error string `json:"error,omitempty"` + } + results := make([]itemResult, 0, len(items)) + successCount := 0 + for _, item := range items { + outText := "" + switch out := item.Output.(type) { + case map[string]interface{}: + if forUser, ok := out["for_user"].(string); ok { + outText = forUser + } + case string: + outText = out + } + ok := item.Status == mapItemStatusSucceeded + if ok { + successCount++ + } + results = append(results, itemResult{ + Index: item.Index, + Success: ok, + Attempts: item.Attempts, + Output: outText, + Error: item.Error, + }) + } + payload := map[string]interface{}{ + "run_id": run.ID.String(), + "status": run.Status, + "count": len(results), + "summary": map[string]interface{}{ + "success_count": successCount, + "failure_count": len(results) - successCount, + "max_retries": maxRetries, + }, + "results": results, + } + data, err := jsonv2.Marshal(payload) + if err != nil { + return ErrorResult("failed to serialize agentic_map inline results").WithError(err) + } + return &ToolResult{ + ForLLM: string(data), + ForUser: string(data), + Silent: false, + IsError: false, + Async: false, + } +} diff --git a/pkg/tools/agentic_map_test.go b/pkg/tools/agentic_map_test.go new file mode 100644 index 000000000..55bd66c8e --- /dev/null +++ b/pkg/tools/agentic_map_test.go @@ -0,0 +1,113 @@ +package tools + +import ( + "context" + "fmt" + "testing" + + jsonv2 "github.com/go-json-experiment/json" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/ZanzyTHEbar/dragonscale/pkg/bus" +) + +func TestAgenticMapTool_Execute_Success(t *testing.T) { + provider := &MockLanguageModel{} + manager := NewSubagentManager(provider, "test-model", "/tmp/test", bus.NewMessageBus()) + manager.SetRunLoop(func(_ context.Context, _ ToolLoopConfig, _, userPrompt, _, _ string) (*ToolLoopResult, error) { + return &ToolLoopResult{Content: "processed: " + userPrompt, Iterations: 1}, nil + }) + + tool := NewAgenticMapTool(manager) + result := tool.Execute(context.Background(), map[string]interface{}{ + "items": []interface{}{ + map[string]interface{}{"name": "a"}, + map[string]interface{}{"name": "b"}, + }, + "task_template": "Handle item {{index}} => {{item_json}}", + "max_retries": float64(1), + }) + + require.NotNil(t, result) + require.False(t, result.IsError, result.ForLLM) + + var payload struct { + Count int `json:"count"` + Summary struct { + SuccessCount int `json:"success_count"` + FailureCount int `json:"failure_count"` + } `json:"summary"` + } + require.NoError(t, jsonv2.Unmarshal([]byte(result.ForLLM), &payload)) + assert.Equal(t, 2, payload.Count) + assert.Equal(t, 2, payload.Summary.SuccessCount) + assert.Equal(t, 0, payload.Summary.FailureCount) +} + +func TestAgenticMapTool_Execute_RetriesFailedItems(t *testing.T) { + provider := &MockLanguageModel{} + manager := NewSubagentManager(provider, "test-model", "/tmp/test", bus.NewMessageBus()) + callCount := 0 + manager.SetRunLoop(func(_ context.Context, _ ToolLoopConfig, _, userPrompt, _, _ string) (*ToolLoopResult, error) { + callCount++ + if callCount == 1 { + return nil, fmt.Errorf("transient failure for %s", userPrompt) + } + return &ToolLoopResult{Content: "processed on retry", Iterations: 1}, nil + }) + + tool := NewAgenticMapTool(manager) + result := tool.Execute(context.Background(), map[string]interface{}{ + "items": []interface{}{map[string]interface{}{"name": "retry-me"}}, + "task_template": "Retry item {{index}} => {{item_json}}", + "max_retries": float64(2), + }) + + require.NotNil(t, result) + require.False(t, result.IsError, result.ForLLM) + + var payload struct { + Results []struct { + Success bool `json:"success"` + Attempts int `json:"attempts"` + } `json:"results"` + } + require.NoError(t, jsonv2.Unmarshal([]byte(result.ForLLM), &payload)) + require.Len(t, payload.Results, 1) + assert.True(t, payload.Results[0].Success) + assert.Equal(t, 2, payload.Results[0].Attempts) +} + +func TestAgenticMapTool_Execute_RequiresPlaceholders(t *testing.T) { + provider := &MockLanguageModel{} + manager := NewSubagentManager(provider, "test-model", "/tmp/test", bus.NewMessageBus()) + tool := NewAgenticMapTool(manager) + + result := tool.Execute(context.Background(), map[string]interface{}{ + "items": []interface{}{map[string]interface{}{"name": "x"}}, + "task_template": "Handle item without placeholders", + }) + + require.NotNil(t, result) + require.True(t, result.IsError) + assert.Contains(t, result.ForLLM, "placeholders") +} + +func TestAgenticMapTool_Execute_ContextCancelled(t *testing.T) { + provider := &MockLanguageModel{} + manager := NewSubagentManager(provider, "test-model", "/tmp/test", bus.NewMessageBus()) + tool := NewAgenticMapTool(manager) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + result := tool.Execute(ctx, map[string]interface{}{ + "items": []interface{}{map[string]interface{}{"name": "x"}}, + "task_template": "Handle {{index}} => {{item_json}}", + }) + + require.NotNil(t, result) + require.True(t, result.IsError) + assert.Contains(t, result.ForLLM, "cancelled") +} diff --git a/pkg/tools/llm_map.go b/pkg/tools/llm_map.go new file mode 100644 index 000000000..5b5cabb63 --- /dev/null +++ b/pkg/tools/llm_map.go @@ -0,0 +1,327 @@ +package tools + +import ( + "context" + "fmt" + "strings" + + jsonv2 "github.com/go-json-experiment/json" + + fantasy "charm.land/fantasy" + memsqlc "github.com/ZanzyTHEbar/dragonscale/pkg/memory/sqlc" +) + +// LLMMapTool performs schema-validated per-item processing with model calls +// and no side effects (toolless, deterministic structure). +type LLMMapTool struct { + model fantasy.LanguageModel + modelID string + runtime *MapRuntime +} + +func NewLLMMapTool(model fantasy.LanguageModel, modelID string) *LLMMapTool { + return &LLMMapTool{ + model: model, + modelID: modelID, + } +} + +func (t *LLMMapTool) SetRuntime(runtime *MapRuntime) { + t.runtime = runtime +} + +func (t *LLMMapTool) Name() string { + return "llm_map" +} + +func (t *LLMMapTool) Description() string { + return "Apply an instruction to each item with schema-validated JSON output. Supports worker-backed runs with run handles for large batches." +} + +func (t *LLMMapTool) Parameters() map[string]interface{} { + return map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "items": map[string]interface{}{ + "type": "array", + "description": "Array of items to process.", + }, + "input_jsonl": map[string]interface{}{ + "type": "string", + "description": "Boundary-only JSONL input. Exactly one of items or input_jsonl is required.", + }, + "instruction": map[string]interface{}{ + "type": "string", + "description": "Instruction applied to each item. Output must be JSON object only.", + }, + "output_schema": map[string]interface{}{ + "type": "object", + "description": "Optional lightweight JSON schema (properties + required) used to validate each output object.", + }, + "execution_mode": map[string]interface{}{ + "type": "string", + "description": "inline or worker. Defaults to worker for JSONL/large batches, inline otherwise.", + }, + "max_retries": map[string]interface{}{ + "type": "integer", + "description": "Retries per item after first attempt when using worker execution (default 1, max 10).", + }, + "session_key": map[string]interface{}{ + "type": "string", + "description": "Optional map run session key for persistence (default: default).", + }, + "idempotency_key": map[string]interface{}{ + "type": "string", + "description": "Optional key to deduplicate repeated run creation.", + }, + }, + "required": []string{"instruction"}, + } +} + +func (t *LLMMapTool) Execute(ctx context.Context, args map[string]interface{}) *ToolResult { + if t.model == nil && t.runtime == nil { + return ErrorResult("llm_map model is not configured").WithError(fmt.Errorf("llm_map model is nil")) + } + + items, usedJSONL, err := parseMapBoundaryItems(args) + if err != nil { + return ErrorResult(err.Error()).WithError(err) + } + instruction, ok := args["instruction"].(string) + if !ok || strings.TrimSpace(instruction) == "" { + return ErrorResult("instruction is required").WithError(fmt.Errorf("instruction is required")) + } + + var schema map[string]interface{} + if rawSchema, ok := args["output_schema"].(map[string]interface{}); ok { + schema = rawSchema + } + schemaJSON := "" + if schema != nil { + b, err := canonicalJSONString(schema) + if err != nil { + return ErrorResult("output_schema is not JSON-serializable").WithError(err) + } + schemaJSON = b + } + maxRetries, err := parseMapMaxRetries(args, 1) + if err != nil { + return ErrorResult(err.Error()).WithError(err) + } + mode, err := resolveMapExecutionMode(args, len(items), usedJSONL) + if err != nil { + return ErrorResult(err.Error()).WithError(err) + } + sessionKey, err := parseOptionalStringArg(args, "session_key") + if err != nil { + return ErrorResult(err.Error()).WithError(err) + } + idempotencyKey, err := parseOptionalStringArg(args, "idempotency_key") + if err != nil { + return ErrorResult(err.Error()).WithError(err) + } + if sessionKey == "" { + sessionKey = "default" + } + + if t.runtime != nil { + run, reused, err := t.runtime.EnqueueRun(ctx, sessionKey, MapRunSpec{ + OperatorKind: MapOperatorLLM, + Instruction: strings.TrimSpace(instruction), + OutputSchemaJSON: schemaJSON, + MaxRetries: uint16(maxRetries), + }, items, idempotencyKey) + if err != nil { + return ErrorResult(fmt.Sprintf("failed to enqueue llm_map run: %v", err)).WithError(err) + } + if mode == "worker" { + payload := map[string]interface{}{ + "run_id": run.ID.String(), + "status": run.Status, + "accepted_count": len(items), + "queued_count": run.QueuedItems, + "execution_mode": "worker", + "idempotent_reuse": reused, + } + data, err := jsonv2.Marshal(payload) + if err != nil { + return ErrorResult("failed to serialize llm_map enqueue result").WithError(err) + } + return &ToolResult{ + ForLLM: string(data), + ForUser: string(data), + Silent: false, + IsError: false, + Async: false, + } + } + + // Inline mode still goes through the worker lifecycle for a single runtime path. + runFinal, err := t.runtime.ProcessRunToTerminal(ctx, run.ID, len(items)*(maxRetries+2)+32) + if err != nil { + return ErrorResult(fmt.Sprintf("failed to complete inline llm_map run %s: %v", run.ID.String(), err)).WithError(err) + } + itemRows, err := t.runtime.ReadRunItems(ctx, run.ID, 0, len(items)) + if err != nil { + return ErrorResult(fmt.Sprintf("failed to read inline llm_map results: %v", err)).WithError(err) + } + return buildLLMMapInlineResult(runFinal, itemRows) + } + + // Legacy synchronous path when runtime persistence is not configured. + if mode == "worker" { + return ErrorResult("worker mode requires map runtime persistence").WithError(fmt.Errorf("map runtime is nil")) + } + + results := make([]map[string]interface{}, 0, len(items)) + for _, item := range items { + prompt := fmt.Sprintf( + "Instruction:\n%s\n\nItem JSON:\n%s\n\nReturn exactly one JSON object and nothing else.", + instruction, + item.ItemJSON, + ) + maxOut := int64(512) + call := fantasy.Call{ + Prompt: fantasy.Prompt{ + fantasy.NewSystemMessage("You are a pure mapper. Return only strict JSON object output."), + fantasy.NewUserMessage(prompt), + }, + MaxOutputTokens: &maxOut, + } + resp, err := t.model.Generate(ctx, call) + if err != nil { + return ErrorResult(fmt.Sprintf("llm_map failed at index %d", item.Index)).WithError(err) + } + + var out map[string]interface{} + if err := jsonv2.Unmarshal([]byte(resp.Content.Text()), &out); err != nil { + return ErrorResult(fmt.Sprintf("llm_map output at index %d is not valid JSON object", item.Index)).WithError(err) + } + if schema != nil { + if err := validateLLMMapOutputSchema(out, schema); err != nil { + return ErrorResult(fmt.Sprintf("llm_map schema validation failed at index %d: %v", item.Index, err)).WithError(err) + } + } + results = append(results, out) + } + + payload := map[string]interface{}{ + "count": len(results), + "results": results, + } + data, err := jsonv2.Marshal(payload) + if err != nil { + return ErrorResult("failed to serialize llm_map results").WithError(err) + } + return &ToolResult{ + ForLLM: string(data), + ForUser: string(data), + Silent: false, + IsError: false, + Async: false, + } +} + +func buildLLMMapInlineResult(run memsqlc.MapRun, items []MapItemProjection) *ToolResult { + results := make([]map[string]interface{}, 0, len(items)) + failures := make([]map[string]interface{}, 0) + for _, item := range items { + if item.Status == mapItemStatusSucceeded { + outMap, ok := item.Output.(map[string]interface{}) + if !ok { + outMap = map[string]interface{}{ + "value": item.Output, + } + } + results = append(results, outMap) + continue + } + failures = append(failures, map[string]interface{}{ + "index": item.Index, + "status": item.Status, + "attempts": item.Attempts, + "error": item.Error, + }) + } + payload := map[string]interface{}{ + "run_id": run.ID.String(), + "status": run.Status, + "count": len(results), + "summary": map[string]interface{}{ + "success_count": len(results), + "failure_count": len(failures), + }, + "results": results, + "failures": failures, + } + data, err := jsonv2.Marshal(payload) + if err != nil { + return ErrorResult("failed to serialize llm_map inline results").WithError(err) + } + return &ToolResult{ + ForLLM: string(data), + ForUser: string(data), + Silent: false, + IsError: false, + Async: false, + } +} + +func validateLLMMapOutputSchema(out map[string]interface{}, schema map[string]interface{}) error { + if requiredRaw, ok := schema["required"].([]interface{}); ok { + for _, r := range requiredRaw { + key, _ := r.(string) + if key == "" { + continue + } + if _, exists := out[key]; !exists { + return fmt.Errorf("missing required field %q", key) + } + } + } + + props, _ := schema["properties"].(map[string]interface{}) + for key, defRaw := range props { + value, exists := out[key] + if !exists { + continue + } + def, _ := defRaw.(map[string]interface{}) + wantType, _ := def["type"].(string) + if wantType == "" { + continue + } + if !matchesJSONType(value, wantType) { + return fmt.Errorf("field %q expected type %q", key, wantType) + } + } + return nil +} + +func matchesJSONType(value interface{}, want string) bool { + switch want { + case "string": + _, ok := value.(string) + return ok + case "number": + _, ok := value.(float64) + return ok + case "integer": + f, ok := value.(float64) + return ok && f == float64(int64(f)) + case "boolean": + _, ok := value.(bool) + return ok + case "object": + _, ok := value.(map[string]interface{}) + return ok + case "array": + _, ok := value.([]interface{}) + return ok + case "null": + return value == nil + default: + return true + } +} diff --git a/pkg/tools/llm_map_test.go b/pkg/tools/llm_map_test.go new file mode 100644 index 000000000..7acbbf260 --- /dev/null +++ b/pkg/tools/llm_map_test.go @@ -0,0 +1,111 @@ +package tools + +import ( + "context" + "fmt" + "testing" + + jsonv2 "github.com/go-json-experiment/json" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + fantasy "charm.land/fantasy" +) + +type llmMapMockModel struct { + responses []string + next int +} + +func (m *llmMapMockModel) Generate(_ context.Context, _ fantasy.Call) (*fantasy.Response, error) { + if len(m.responses) == 0 { + return nil, fmt.Errorf("no responses configured") + } + resp := m.responses[m.next%len(m.responses)] + m.next++ + return &fantasy.Response{ + Content: fantasy.ResponseContent{fantasy.TextContent{Text: resp}}, + FinishReason: fantasy.FinishReasonStop, + }, nil +} + +func (m *llmMapMockModel) Stream(ctx context.Context, call fantasy.Call) (fantasy.StreamResponse, error) { + resp, err := m.Generate(ctx, call) + if err != nil { + return nil, err + } + return func(yield func(fantasy.StreamPart) bool) { + if !yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextDelta, Delta: resp.Content.Text()}) { + return + } + yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop}) + }, nil +} + +func (m *llmMapMockModel) GenerateObject(_ context.Context, _ fantasy.ObjectCall) (*fantasy.ObjectResponse, error) { + return nil, fmt.Errorf("not implemented") +} + +func (m *llmMapMockModel) StreamObject(_ context.Context, _ fantasy.ObjectCall) (fantasy.ObjectStreamResponse, error) { + return nil, fmt.Errorf("not implemented") +} + +func (m *llmMapMockModel) Provider() string { return "mock" } +func (m *llmMapMockModel) Model() string { return "mock-llm-map" } + +func TestLLMMapTool_Execute_Success(t *testing.T) { + model := &llmMapMockModel{ + responses: []string{ + `{"label":"alpha","priority":1}`, + `{"label":"beta","priority":2}`, + }, + } + tool := NewLLMMapTool(model, "mock-llm-map") + + result := tool.Execute(context.Background(), map[string]interface{}{ + "instruction": "Convert item to {label, priority}", + "items": []interface{}{ + map[string]interface{}{"name": "A"}, + map[string]interface{}{"name": "B"}, + }, + "output_schema": map[string]interface{}{ + "required": []interface{}{"label", "priority"}, + "properties": map[string]interface{}{ + "label": map[string]interface{}{"type": "string"}, + "priority": map[string]interface{}{"type": "integer"}, + }, + }, + }) + require.NotNil(t, result) + require.False(t, result.IsError, result.ForLLM) + + var payload struct { + Count int `json:"count"` + Results []map[string]interface{} `json:"results"` + } + require.NoError(t, jsonv2.Unmarshal([]byte(result.ForLLM), &payload)) + assert.Equal(t, 2, payload.Count) + require.Len(t, payload.Results, 2) + assert.Equal(t, "alpha", payload.Results[0]["label"]) +} + +func TestLLMMapTool_Execute_SchemaValidationFailure(t *testing.T) { + model := &llmMapMockModel{ + responses: []string{`{"only":"value"}`}, + } + tool := NewLLMMapTool(model, "mock-llm-map") + + result := tool.Execute(context.Background(), map[string]interface{}{ + "instruction": "Return normalized object", + "items": []interface{}{map[string]interface{}{"name": "x"}}, + "output_schema": map[string]interface{}{ + "required": []interface{}{"label"}, + "properties": map[string]interface{}{ + "label": map[string]interface{}{"type": "string"}, + }, + }, + }) + require.NotNil(t, result) + require.True(t, result.IsError) + assert.Contains(t, result.ForLLM, "schema validation failed") +} diff --git a/pkg/tools/map_boundary.go b/pkg/tools/map_boundary.go new file mode 100644 index 000000000..65cd51377 --- /dev/null +++ b/pkg/tools/map_boundary.go @@ -0,0 +1,204 @@ +package tools + +import ( + "bufio" + "crypto/sha256" + "encoding/hex" + "fmt" + "strings" + + jsonv2 "github.com/go-json-experiment/json" +) + +const ( + defaultMapInlineCutoff = 12 + defaultMapReadLimit = 100 + maxMapReadLimit = 500 + maxMapItems = 5000 +) + +type mapBoundaryItem struct { + Index int + ItemJSON string + InputHash string +} + +func parseMapBoundaryItems(args map[string]interface{}) ([]mapBoundaryItem, bool, error) { + itemsRaw, hasItems := args["items"] + inputJSONLRaw, hasJSONL := args["input_jsonl"] + if hasItems == hasJSONL { + return nil, false, fmt.Errorf("provide exactly one of items or input_jsonl") + } + + items := make([]mapBoundaryItem, 0) + if hasItems { + array, ok := itemsRaw.([]interface{}) + if !ok { + return nil, false, fmt.Errorf("items must be an array") + } + if len(array) == 0 { + return nil, false, fmt.Errorf("items array is empty") + } + if len(array) > maxMapItems { + return nil, false, fmt.Errorf("items array exceeds max size %d", maxMapItems) + } + for idx, item := range array { + itemJSON, err := canonicalJSONString(item) + if err != nil { + return nil, false, fmt.Errorf("serialize item at index %d: %w", idx, err) + } + items = append(items, mapBoundaryItem{ + Index: idx, + ItemJSON: itemJSON, + InputHash: hashString(itemJSON), + }) + } + return items, false, nil + } + + inputJSONL, ok := inputJSONLRaw.(string) + if !ok || strings.TrimSpace(inputJSONL) == "" { + return nil, true, fmt.Errorf("input_jsonl must be a non-empty string") + } + scanner := bufio.NewScanner(strings.NewReader(inputJSONL)) + // Allow large records while still preventing unbounded memory growth. + const maxJSONLLineBytes = 4 * 1024 * 1024 + scanner.Buffer(make([]byte, 64*1024), maxJSONLLineBytes) + + lineNo := 0 + for scanner.Scan() { + lineNo++ + line := strings.TrimSpace(scanner.Text()) + if line == "" { + continue + } + var item interface{} + if err := jsonv2.Unmarshal([]byte(line), &item); err != nil { + return nil, true, fmt.Errorf("invalid JSONL at line %d: %w", lineNo, err) + } + itemJSON, err := canonicalJSONString(item) + if err != nil { + return nil, true, fmt.Errorf("canonicalize JSONL item at line %d: %w", lineNo, err) + } + items = append(items, mapBoundaryItem{ + Index: len(items), + ItemJSON: itemJSON, + InputHash: hashString(itemJSON), + }) + if len(items) > maxMapItems { + return nil, true, fmt.Errorf("input_jsonl exceeds max item count %d", maxMapItems) + } + } + if err := scanner.Err(); err != nil { + return nil, true, fmt.Errorf("read input_jsonl: %w", err) + } + if len(items) == 0 { + return nil, true, fmt.Errorf("input_jsonl contains no records") + } + return items, true, nil +} + +func resolveMapExecutionMode(args map[string]interface{}, itemCount int, usedJSONL bool) (string, error) { + if raw, ok := args["execution_mode"]; ok { + mode, ok := raw.(string) + if !ok { + return "", fmt.Errorf("execution_mode must be a string") + } + switch strings.ToLower(strings.TrimSpace(mode)) { + case "inline": + return "inline", nil + case "worker": + return "worker", nil + default: + return "", fmt.Errorf("execution_mode must be one of inline|worker") + } + } + if usedJSONL || itemCount > defaultMapInlineCutoff { + return "worker", nil + } + return "inline", nil +} + +func parseMapReadFormat(args map[string]interface{}) (string, error) { + raw, ok := args["format"] + if !ok { + return "json", nil + } + format, ok := raw.(string) + if !ok { + return "", fmt.Errorf("format must be a string") + } + switch strings.ToLower(strings.TrimSpace(format)) { + case "json", "": + return "json", nil + case "jsonl": + return "jsonl", nil + default: + return "", fmt.Errorf("format must be json or jsonl") + } +} + +func parseMapLimit(args map[string]interface{}, key string, defaultValue, maxValue int) (int, error) { + value, ok := args[key] + if !ok { + return defaultValue, nil + } + switch v := value.(type) { + case float64: + if v < 0 { + return 0, fmt.Errorf("%s must be >= 0", key) + } + out := int(v) + if out > maxValue { + out = maxValue + } + return out, nil + case int: + if v < 0 { + return 0, fmt.Errorf("%s must be >= 0", key) + } + if v > maxValue { + return maxValue, nil + } + return v, nil + default: + return 0, fmt.Errorf("%s must be numeric", key) + } +} + +func parseMapMaxRetries(args map[string]interface{}, defaultValue int) (int, error) { + raw, ok := args["max_retries"] + if !ok { + return defaultValue, nil + } + value, err := parseMapLimit(map[string]interface{}{"value": raw}, "value", defaultValue, 10) + if err != nil { + return 0, fmt.Errorf("max_retries %w", err) + } + return value, nil +} + +func parseOptionalStringArg(args map[string]interface{}, key string) (string, error) { + raw, ok := args[key] + if !ok { + return "", nil + } + value, ok := raw.(string) + if !ok { + return "", fmt.Errorf("%s must be a string", key) + } + return strings.TrimSpace(value), nil +} + +func canonicalJSONString(v interface{}) (string, error) { + b, err := jsonv2.Marshal(v) + if err != nil { + return "", err + } + return string(b), nil +} + +func hashString(value string) string { + sum := sha256.Sum256([]byte(value)) + return hex.EncodeToString(sum[:]) +} diff --git a/pkg/tools/map_flatbuffer_codec.go b/pkg/tools/map_flatbuffer_codec.go new file mode 100644 index 000000000..fe89230e0 --- /dev/null +++ b/pkg/tools/map_flatbuffer_codec.go @@ -0,0 +1,230 @@ +package tools + +import ( + "fmt" + + "github.com/ZanzyTHEbar/dragonscale/pkg/tools/mapopsfb" + flatbuffers "github.com/google/flatbuffers/go" +) + +type MapOperatorKind string + +const ( + MapOperatorLLM MapOperatorKind = "llm_map" + MapOperatorAgentic MapOperatorKind = "agentic_map" + mapPayloadVersion uint16 = 1 +) + +type MapRunSpec struct { + Version uint16 + OperatorKind MapOperatorKind + Instruction string + TaskTemplate string + OutputSchemaJSON string + MaxRetries uint16 + DelegatedScope string + KeptWork string + OriginChannel string + OriginChatID string +} + +type MapItemInputRecord struct { + Version uint16 + ItemIndex uint32 + ItemJSON string + InputHash string +} + +type MapItemOutputRecord struct { + Version uint16 + ItemIndex uint32 + Success bool + Attempts uint16 + OutputJSON string + ErrorText string + OutputHash string +} + +func EncodeMapRunSpecFlatBuffer(spec MapRunSpec) ([]byte, error) { + if spec.OperatorKind == "" { + return nil, fmt.Errorf("operator kind is required") + } + version := spec.Version + if version == 0 { + version = 1 + } + maxRetries := spec.MaxRetries + if maxRetries == 0 { + maxRetries = 1 + } + + b := flatbuffers.NewBuilder(256) + operatorKind := b.CreateString(string(spec.OperatorKind)) + instruction := createOptionalFBString(b, spec.Instruction) + taskTemplate := createOptionalFBString(b, spec.TaskTemplate) + outputSchemaJSON := createOptionalFBString(b, spec.OutputSchemaJSON) + delegatedScope := createOptionalFBString(b, spec.DelegatedScope) + keptWork := createOptionalFBString(b, spec.KeptWork) + originChannel := createOptionalFBString(b, spec.OriginChannel) + originChatID := createOptionalFBString(b, spec.OriginChatID) + + mapopsfb.MapRunSpecStart(b) + mapopsfb.MapRunSpecAddVersion(b, version) + mapopsfb.MapRunSpecAddOperatorKind(b, operatorKind) + if instruction != 0 { + mapopsfb.MapRunSpecAddInstruction(b, instruction) + } + if taskTemplate != 0 { + mapopsfb.MapRunSpecAddTaskTemplate(b, taskTemplate) + } + if outputSchemaJSON != 0 { + mapopsfb.MapRunSpecAddOutputSchemaJson(b, outputSchemaJSON) + } + mapopsfb.MapRunSpecAddMaxRetries(b, maxRetries) + if delegatedScope != 0 { + mapopsfb.MapRunSpecAddDelegatedScope(b, delegatedScope) + } + if keptWork != 0 { + mapopsfb.MapRunSpecAddKeptWork(b, keptWork) + } + if originChannel != 0 { + mapopsfb.MapRunSpecAddOriginChannel(b, originChannel) + } + if originChatID != 0 { + mapopsfb.MapRunSpecAddOriginChatId(b, originChatID) + } + obj := mapopsfb.MapRunSpecEnd(b) + b.Finish(obj) + + return b.FinishedBytes(), nil +} + +func DecodeMapRunSpecFlatBuffer(buf []byte) (MapRunSpec, error) { + if len(buf) < flatbuffers.SizeUint32 { + return MapRunSpec{}, fmt.Errorf("flatbuffer payload too short") + } + specFB := mapopsfb.GetRootAsMapRunSpec(buf, 0) + version := specFB.Version() + if version != mapPayloadVersion { + return MapRunSpec{}, fmt.Errorf("unsupported map run spec version %d", version) + } + spec := MapRunSpec{ + Version: version, + OperatorKind: MapOperatorKind(string(specFB.OperatorKind())), + Instruction: string(specFB.Instruction()), + TaskTemplate: string(specFB.TaskTemplate()), + OutputSchemaJSON: string(specFB.OutputSchemaJson()), + MaxRetries: specFB.MaxRetries(), + DelegatedScope: string(specFB.DelegatedScope()), + KeptWork: string(specFB.KeptWork()), + OriginChannel: string(specFB.OriginChannel()), + OriginChatID: string(specFB.OriginChatId()), + } + if spec.OperatorKind == "" { + return MapRunSpec{}, fmt.Errorf("invalid map run spec: operator kind is empty") + } + return spec, nil +} + +func EncodeMapItemInputFlatBuffer(rec MapItemInputRecord) ([]byte, error) { + if rec.ItemJSON == "" { + return nil, fmt.Errorf("item JSON is required") + } + version := rec.Version + if version == 0 { + version = 1 + } + + b := flatbuffers.NewBuilder(128) + itemJSON := b.CreateString(rec.ItemJSON) + inputHash := createOptionalFBString(b, rec.InputHash) + + mapopsfb.MapItemInputStart(b) + mapopsfb.MapItemInputAddVersion(b, version) + mapopsfb.MapItemInputAddItemIndex(b, rec.ItemIndex) + mapopsfb.MapItemInputAddItemJson(b, itemJSON) + if inputHash != 0 { + mapopsfb.MapItemInputAddInputHash(b, inputHash) + } + obj := mapopsfb.MapItemInputEnd(b) + b.Finish(obj) + return b.FinishedBytes(), nil +} + +func DecodeMapItemInputFlatBuffer(buf []byte) (MapItemInputRecord, error) { + if len(buf) < flatbuffers.SizeUint32 { + return MapItemInputRecord{}, fmt.Errorf("flatbuffer payload too short") + } + recFB := mapopsfb.GetRootAsMapItemInput(buf, 0) + version := recFB.Version() + if version != mapPayloadVersion { + return MapItemInputRecord{}, fmt.Errorf("unsupported map item input version %d", version) + } + rec := MapItemInputRecord{ + Version: version, + ItemIndex: recFB.ItemIndex(), + ItemJSON: string(recFB.ItemJson()), + InputHash: string(recFB.InputHash()), + } + if rec.ItemJSON == "" { + return MapItemInputRecord{}, fmt.Errorf("invalid map item input: item JSON is empty") + } + return rec, nil +} + +func EncodeMapItemOutputFlatBuffer(rec MapItemOutputRecord) ([]byte, error) { + version := rec.Version + if version == 0 { + version = 1 + } + + b := flatbuffers.NewBuilder(128) + outputJSON := createOptionalFBString(b, rec.OutputJSON) + errorText := createOptionalFBString(b, rec.ErrorText) + outputHash := createOptionalFBString(b, rec.OutputHash) + + mapopsfb.MapItemOutputStart(b) + mapopsfb.MapItemOutputAddVersion(b, version) + mapopsfb.MapItemOutputAddItemIndex(b, rec.ItemIndex) + mapopsfb.MapItemOutputAddSuccess(b, rec.Success) + mapopsfb.MapItemOutputAddAttempts(b, rec.Attempts) + if outputJSON != 0 { + mapopsfb.MapItemOutputAddOutputJson(b, outputJSON) + } + if errorText != 0 { + mapopsfb.MapItemOutputAddErrorText(b, errorText) + } + if outputHash != 0 { + mapopsfb.MapItemOutputAddOutputHash(b, outputHash) + } + obj := mapopsfb.MapItemOutputEnd(b) + b.Finish(obj) + return b.FinishedBytes(), nil +} + +func DecodeMapItemOutputFlatBuffer(buf []byte) (MapItemOutputRecord, error) { + if len(buf) < flatbuffers.SizeUint32 { + return MapItemOutputRecord{}, fmt.Errorf("flatbuffer payload too short") + } + recFB := mapopsfb.GetRootAsMapItemOutput(buf, 0) + version := recFB.Version() + if version != mapPayloadVersion { + return MapItemOutputRecord{}, fmt.Errorf("unsupported map item output version %d", version) + } + return MapItemOutputRecord{ + Version: version, + ItemIndex: recFB.ItemIndex(), + Success: recFB.Success(), + Attempts: recFB.Attempts(), + OutputJSON: string(recFB.OutputJson()), + ErrorText: string(recFB.ErrorText()), + OutputHash: string(recFB.OutputHash()), + }, nil +} + +func createOptionalFBString(b *flatbuffers.Builder, value string) flatbuffers.UOffsetT { + if value == "" { + return 0 + } + return b.CreateString(value) +} diff --git a/pkg/tools/map_flatbuffer_codec_test.go b/pkg/tools/map_flatbuffer_codec_test.go new file mode 100644 index 000000000..da334d369 --- /dev/null +++ b/pkg/tools/map_flatbuffer_codec_test.go @@ -0,0 +1,83 @@ +package tools + +import ( + "testing" +) + +func TestMapRunSpecFlatBuffer_RoundTrip(t *testing.T) { + original := MapRunSpec{ + Version: 1, + OperatorKind: MapOperatorLLM, + Instruction: "extract label and priority", + TaskTemplate: "", + OutputSchemaJSON: `{"type":"object"}`, + MaxRetries: 3, + DelegatedScope: "", + KeptWork: "", + OriginChannel: "cli", + OriginChatID: "direct", + } + + data, err := EncodeMapRunSpecFlatBuffer(original) + if err != nil { + t.Fatalf("encode run spec: %v", err) + } + decoded, err := DecodeMapRunSpecFlatBuffer(data) + if err != nil { + t.Fatalf("decode run spec: %v", err) + } + + if decoded.OperatorKind != original.OperatorKind { + t.Fatalf("operator kind mismatch: got %q want %q", decoded.OperatorKind, original.OperatorKind) + } + if decoded.Instruction != original.Instruction { + t.Fatalf("instruction mismatch: got %q want %q", decoded.Instruction, original.Instruction) + } + if decoded.OutputSchemaJSON != original.OutputSchemaJSON { + t.Fatalf("schema mismatch: got %q want %q", decoded.OutputSchemaJSON, original.OutputSchemaJSON) + } + if decoded.MaxRetries != original.MaxRetries { + t.Fatalf("max retries mismatch: got %d want %d", decoded.MaxRetries, original.MaxRetries) + } +} + +func TestMapItemFlatBuffer_RoundTrip(t *testing.T) { + input := MapItemInputRecord{ + Version: 1, + ItemIndex: 7, + ItemJSON: `{"name":"alpha","priority":1}`, + InputHash: "abc123", + } + inBuf, err := EncodeMapItemInputFlatBuffer(input) + if err != nil { + t.Fatalf("encode input: %v", err) + } + inDecoded, err := DecodeMapItemInputFlatBuffer(inBuf) + if err != nil { + t.Fatalf("decode input: %v", err) + } + if inDecoded.ItemIndex != input.ItemIndex || inDecoded.ItemJSON != input.ItemJSON || inDecoded.InputHash != input.InputHash { + t.Fatalf("input mismatch: got %+v want %+v", inDecoded, input) + } + + output := MapItemOutputRecord{ + Version: 1, + ItemIndex: 7, + Success: true, + Attempts: 2, + OutputJSON: `{"label":"alpha"}`, + ErrorText: "", + OutputHash: "def456", + } + outBuf, err := EncodeMapItemOutputFlatBuffer(output) + if err != nil { + t.Fatalf("encode output: %v", err) + } + outDecoded, err := DecodeMapItemOutputFlatBuffer(outBuf) + if err != nil { + t.Fatalf("decode output: %v", err) + } + if outDecoded.ItemIndex != output.ItemIndex || outDecoded.Success != output.Success || outDecoded.Attempts != output.Attempts || outDecoded.OutputJSON != output.OutputJSON { + t.Fatalf("output mismatch: got %+v want %+v", outDecoded, output) + } +} diff --git a/pkg/tools/map_payloads.fbs b/pkg/tools/map_payloads.fbs new file mode 100644 index 000000000..aca69a5fc --- /dev/null +++ b/pkg/tools/map_payloads.fbs @@ -0,0 +1,33 @@ +// FlatBuffers contract for map operator persistence payloads. +// Boundary JSON/JSONL is converted to these binary records before DB persistence. +namespace mapopsfb; + +table MapRunSpec { + version: uint16 = 1; + operator_kind: string (required); // llm_map | agentic_map + instruction: string; + task_template: string; + output_schema_json: string; + max_retries: uint16 = 1; + delegated_scope: string; + kept_work: string; + origin_channel: string; + origin_chat_id: string; +} + +table MapItemInput { + version: uint16 = 1; + item_index: uint32; + item_json: string (required); // canonical boundary JSON object + input_hash: string; +} + +table MapItemOutput { + version: uint16 = 1; + item_index: uint32; + success: bool = false; + attempts: uint16 = 0; + output_json: string; + error_text: string; + output_hash: string; +} diff --git a/pkg/tools/map_run_tools.go b/pkg/tools/map_run_tools.go new file mode 100644 index 000000000..c7c84759a --- /dev/null +++ b/pkg/tools/map_run_tools.go @@ -0,0 +1,216 @@ +package tools + +import ( + "context" + "fmt" + "strings" + + "github.com/ZanzyTHEbar/dragonscale/pkg/ids" + jsonv2 "github.com/go-json-experiment/json" +) + +type MapRunStatusTool struct { + runtime *MapRuntime +} + +func NewMapRunStatusTool(runtime *MapRuntime) *MapRunStatusTool { + return &MapRunStatusTool{runtime: runtime} +} + +func (t *MapRunStatusTool) Name() string { + return "map_run_status" +} + +func (t *MapRunStatusTool) Description() string { + return "Return current status/progress for an llm_map or agentic_map run." +} + +func (t *MapRunStatusTool) Parameters() map[string]interface{} { + return map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "run_id": map[string]interface{}{ + "type": "string", + "description": "Map run ID returned by llm_map or agentic_map.", + }, + "process_steps": map[string]interface{}{ + "type": "integer", + "description": "Optional number of worker steps to execute before reading status (default 8).", + }, + }, + "required": []string{"run_id"}, + } +} + +func (t *MapRunStatusTool) Execute(ctx context.Context, args map[string]interface{}) *ToolResult { + if t.runtime == nil { + return ErrorResult("map runtime is not configured").WithError(fmt.Errorf("map runtime is nil")) + } + runIDRaw, ok := args["run_id"].(string) + if !ok || strings.TrimSpace(runIDRaw) == "" { + return ErrorResult("run_id is required").WithError(fmt.Errorf("run_id is required")) + } + runID, err := ids.Parse(strings.TrimSpace(runIDRaw)) + if err != nil { + return ErrorResult("run_id must be a valid UUID").WithError(err) + } + steps, err := parseMapLimit(args, "process_steps", 8, 500) + if err != nil { + return ErrorResult(err.Error()).WithError(err) + } + if steps > 0 { + if _, err := t.runtime.ProcessPending(ctx, steps); err != nil { + return ErrorResult(fmt.Sprintf("failed to process map jobs: %v", err)).WithError(err) + } + } + run, err := t.runtime.GetRun(ctx, runID) + if err != nil { + return ErrorResult(fmt.Sprintf("failed to load map run: %v", err)).WithError(err) + } + payload := toMapRunProjection(run) + data, err := jsonv2.Marshal(payload) + if err != nil { + return ErrorResult("failed to serialize map run status").WithError(err) + } + return &ToolResult{ + ForLLM: string(data), + ForUser: string(data), + Silent: false, + IsError: false, + Async: false, + } +} + +type MapRunReadTool struct { + runtime *MapRuntime +} + +func NewMapRunReadTool(runtime *MapRuntime) *MapRunReadTool { + return &MapRunReadTool{runtime: runtime} +} + +func (t *MapRunReadTool) Name() string { + return "map_run_read" +} + +func (t *MapRunReadTool) Description() string { + return "Read paged item results for a map run; output as JSON or generated JSONL." +} + +func (t *MapRunReadTool) Parameters() map[string]interface{} { + return map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "run_id": map[string]interface{}{ + "type": "string", + "description": "Map run ID returned by llm_map or agentic_map.", + }, + "offset": map[string]interface{}{ + "type": "integer", + "description": "Pagination offset (default 0).", + }, + "limit": map[string]interface{}{ + "type": "integer", + "description": "Pagination limit (default 100, max 500).", + }, + "format": map[string]interface{}{ + "type": "string", + "description": "Projection format: json (default) or jsonl.", + }, + "process_steps": map[string]interface{}{ + "type": "integer", + "description": "Optional worker steps to process before reading (default 8).", + }, + }, + "required": []string{"run_id"}, + } +} + +func (t *MapRunReadTool) Execute(ctx context.Context, args map[string]interface{}) *ToolResult { + if t.runtime == nil { + return ErrorResult("map runtime is not configured").WithError(fmt.Errorf("map runtime is nil")) + } + runIDRaw, ok := args["run_id"].(string) + if !ok || strings.TrimSpace(runIDRaw) == "" { + return ErrorResult("run_id is required").WithError(fmt.Errorf("run_id is required")) + } + runID, err := ids.Parse(strings.TrimSpace(runIDRaw)) + if err != nil { + return ErrorResult("run_id must be a valid UUID").WithError(err) + } + + offset, err := parseMapLimit(args, "offset", 0, 1_000_000) + if err != nil { + return ErrorResult(err.Error()).WithError(err) + } + limit, err := parseMapLimit(args, "limit", defaultMapReadLimit, maxMapReadLimit) + if err != nil { + return ErrorResult(err.Error()).WithError(err) + } + format, err := parseMapReadFormat(args) + if err != nil { + return ErrorResult(err.Error()).WithError(err) + } + steps, err := parseMapLimit(args, "process_steps", 8, 500) + if err != nil { + return ErrorResult(err.Error()).WithError(err) + } + if steps > 0 { + if _, err := t.runtime.ProcessPending(ctx, steps); err != nil { + return ErrorResult(fmt.Sprintf("failed to process map jobs: %v", err)).WithError(err) + } + } + + run, err := t.runtime.GetRun(ctx, runID) + if err != nil { + return ErrorResult(fmt.Sprintf("failed to load map run: %v", err)).WithError(err) + } + items, err := t.runtime.ReadRunItems(ctx, runID, offset, limit) + if err != nil { + return ErrorResult(fmt.Sprintf("failed to load map run items: %v", err)).WithError(err) + } + + payload := map[string]interface{}{ + "run": toMapRunProjection(run), + "offset": offset, + "limit": limit, + "format": format, + } + if format == "jsonl" { + jsonlText, err := encodeMapItemsJSONL(items) + if err != nil { + return ErrorResult("failed to encode JSONL projection").WithError(err) + } + payload["items_jsonl"] = jsonlText + payload["item_count"] = len(items) + } else { + payload["items"] = items + } + + data, err := jsonv2.Marshal(payload) + if err != nil { + return ErrorResult("failed to serialize map run page").WithError(err) + } + return &ToolResult{ + ForLLM: string(data), + ForUser: string(data), + Silent: false, + IsError: false, + Async: false, + } +} + +func encodeMapItemsJSONL(items []MapItemProjection) (string, error) { + if len(items) == 0 { + return "", nil + } + lines := make([]string, 0, len(items)) + for _, item := range items { + b, err := jsonv2.Marshal(item) + if err != nil { + return "", err + } + lines = append(lines, string(b)) + } + return strings.Join(lines, "\n"), nil +} diff --git a/pkg/tools/map_runtime.go b/pkg/tools/map_runtime.go new file mode 100644 index 000000000..42cc0cc27 --- /dev/null +++ b/pkg/tools/map_runtime.go @@ -0,0 +1,833 @@ +package tools + +import ( + "context" + "database/sql" + "errors" + "fmt" + "strconv" + "strings" + "sync" + "time" + + fantasy "charm.land/fantasy" + "github.com/ZanzyTHEbar/dragonscale/pkg/ids" + "github.com/ZanzyTHEbar/dragonscale/pkg/logger" + memsqlc "github.com/ZanzyTHEbar/dragonscale/pkg/memory/sqlc" + "github.com/ZanzyTHEbar/dragonscale/pkg/worker" + jsonv2 "github.com/go-json-experiment/json" +) + +const ( + mapRunStatusQueued = "queued" + mapRunStatusRunning = "running" + mapRunStatusSucceeded = "succeeded" + mapRunStatusFailed = "failed" + mapRunStatusCancelled = "cancelled" + + mapItemStatusQueued = "queued" + mapItemStatusRunning = "running" + mapItemStatusSucceeded = "succeeded" + mapItemStatusFailed = "failed" + + mapJobKindLLMItem = "map_item_llm" + mapJobKindAgenticItem = "map_item_agentic" +) + +type mapJobPayload struct { + RunID string `json:"run_id"` + ItemID string `json:"item_id"` +} + +type MapRunProjection struct { + RunID string `json:"run_id"` + OperatorKind string `json:"operator_kind"` + Status string `json:"status"` + TotalItems int64 `json:"total_items"` + QueuedItems int64 `json:"queued_items"` + RunningItems int64 `json:"running_items"` + SucceededItems int64 `json:"succeeded_items"` + FailedItems int64 `json:"failed_items"` + LastError string `json:"last_error,omitempty"` + CreatedAt string `json:"created_at"` + UpdatedAt string `json:"updated_at"` + CompletedAt string `json:"completed_at,omitempty"` +} + +type MapItemProjection struct { + ItemID string `json:"item_id"` + Index int64 `json:"index"` + Status string `json:"status"` + Attempts int64 `json:"attempts"` + Input interface{} `json:"input,omitempty"` + Output interface{} `json:"output,omitempty"` + Error string `json:"error,omitempty"` + InputHash string `json:"input_hash,omitempty"` + OutputHash string `json:"output_hash,omitempty"` +} + +type MapRuntime struct { + queries *memsqlc.Queries + agentID string + llmModel fantasy.LanguageModel + llmModelID string + subagentManager *SubagentManager + workerLockedBy string + processMu sync.Mutex +} + +func NewMapRuntime( + queries *memsqlc.Queries, + agentID string, + llmModel fantasy.LanguageModel, + llmModelID string, + subagentManager *SubagentManager, +) *MapRuntime { + if queries == nil { + return nil + } + return &MapRuntime{ + queries: queries, + agentID: agentID, + llmModel: llmModel, + llmModelID: llmModelID, + subagentManager: subagentManager, + workerLockedBy: "map-runtime", + } +} + +func (rt *MapRuntime) EnqueueRun( + ctx context.Context, + sessionKey string, + spec MapRunSpec, + items []mapBoundaryItem, + idempotencyKey string, +) (memsqlc.MapRun, bool, error) { + if rt == nil || rt.queries == nil { + return memsqlc.MapRun{}, false, fmt.Errorf("map runtime is not configured") + } + if len(items) == 0 { + return memsqlc.MapRun{}, false, fmt.Errorf("at least one item is required") + } + if sessionKey == "" { + sessionKey = "default" + } + if spec.OperatorKind != MapOperatorLLM && spec.OperatorKind != MapOperatorAgentic { + return memsqlc.MapRun{}, false, fmt.Errorf("unsupported operator kind %q", spec.OperatorKind) + } + seenIndices := make(map[int]struct{}, len(items)) + for _, item := range items { + if item.Index < 0 || uint64(item.Index) > uint64(^uint32(0)) { + return memsqlc.MapRun{}, false, fmt.Errorf("item index %d out of uint32 range", item.Index) + } + if _, exists := seenIndices[item.Index]; exists { + return memsqlc.MapRun{}, false, fmt.Errorf("duplicate item index %d", item.Index) + } + seenIndices[item.Index] = struct{}{} + } + + specFB, err := EncodeMapRunSpecFlatBuffer(spec) + if err != nil { + return memsqlc.MapRun{}, false, fmt.Errorf("encode run spec: %w", err) + } + idemPtr := optionalStringPtr(strings.TrimSpace(idempotencyKey)) + if idemPtr != nil { + existing, err := rt.queries.GetMapRunByIdempotencyKey(ctx, memsqlc.GetMapRunByIdempotencyKeyParams{ + AgentID: rt.agentID, + SessionKey: sessionKey, + OperatorKind: string(spec.OperatorKind), + IdempotencyKey: idemPtr, + }) + if err == nil { + return existing, true, nil + } + if !errors.Is(err, sql.ErrNoRows) { + return memsqlc.MapRun{}, false, fmt.Errorf("lookup idempotent run: %w", err) + } + } + + runID := ids.New() + run, err := rt.queries.InsertMapRun(ctx, memsqlc.InsertMapRunParams{ + ID: runID, + AgentID: rt.agentID, + SessionKey: sessionKey, + OperatorKind: string(spec.OperatorKind), + IdempotencyKey: idemPtr, + Status: mapRunStatusQueued, + TotalItems: int64(len(items)), + QueuedItems: int64(len(items)), + RunningItems: 0, + SucceededItems: 0, + FailedItems: 0, + SpecFb: specFB, + LastError: nil, + }) + if err != nil { + if idemPtr != nil && isUniqueConstraintError(err) { + existing, lookupErr := rt.queries.GetMapRunByIdempotencyKey(ctx, memsqlc.GetMapRunByIdempotencyKeyParams{ + AgentID: rt.agentID, + SessionKey: sessionKey, + OperatorKind: string(spec.OperatorKind), + IdempotencyKey: idemPtr, + }) + if lookupErr == nil { + return existing, true, nil + } + } + return memsqlc.MapRun{}, false, fmt.Errorf("insert map run: %w", err) + } + + jobKind := mapJobKindLLMItem + if spec.OperatorKind == MapOperatorAgentic { + jobKind = mapJobKindAgenticItem + } + maxAttempts := int64(spec.MaxRetries) + 1 + if maxAttempts < 1 { + maxAttempts = 1 + } + + for _, item := range items { + inputFB, encErr := EncodeMapItemInputFlatBuffer(MapItemInputRecord{ + ItemIndex: uint32(item.Index), + ItemJSON: item.ItemJSON, + InputHash: item.InputHash, + }) + if encErr != nil { + msg := fmt.Sprintf("encode item %d: %v", item.Index, encErr) + _, _ = rt.failRun(ctx, run.ID, &msg) + return memsqlc.MapRun{}, false, errors.New(msg) + } + itemID := ids.New() + if _, err := rt.queries.InsertMapItem(ctx, memsqlc.InsertMapItemParams{ + ID: itemID, + RunID: run.ID, + ItemIndex: int64(item.Index), + Status: mapItemStatusQueued, + Attempts: 0, + LastError: nil, + InputFb: inputFB, + OutputFb: nil, + InputHash: optionalStringPtr(item.InputHash), + OutputHash: nil, + }); err != nil { + msg := fmt.Sprintf("insert map item %d: %v", item.Index, err) + _, _ = rt.failRun(ctx, run.ID, &msg) + return memsqlc.MapRun{}, false, errors.New(msg) + } + + dedupeKey := fmt.Sprintf("map:%s:%d", run.ID.String(), item.Index) + if _, err := rt.queries.EnqueueJob(ctx, memsqlc.EnqueueJobParams{ + ID: ids.New(), + Kind: jobKind, + DedupeKey: &dedupeKey, + MaxAttempts: maxAttempts, + RunAt: time.Now().UTC(), + PayloadJson: []byte(`{}`), + }); err != nil { + msg := fmt.Sprintf("enqueue map item %d: %v", item.Index, err) + _, _ = rt.failRun(ctx, run.ID, &msg) + return memsqlc.MapRun{}, false, errors.New(msg) + } + } + + logger.InfoCF("map_runtime", "map run enqueued", map[string]interface{}{ + "run_id": run.ID.String(), + "operator_kind": run.OperatorKind, + "items": len(items), + "idempotency": idempotencyKey, + }) + return run, false, nil +} + +func (rt *MapRuntime) ProcessPending(ctx context.Context, maxSteps int) (int, error) { + if rt == nil || rt.queries == nil { + return 0, nil + } + if maxSteps <= 0 { + maxSteps = 1 + } + opts := &worker.Options{ + LockedBy: rt.workerLockedBy, + Handlers: map[string]worker.HandlerFunc{ + mapJobKindLLMItem: rt.handleLLMMapJob, + mapJobKindAgenticItem: rt.handleAgenticMapJob, + }, + Backoff: func(_ int) time.Duration { + return 10 * time.Millisecond + }, + Now: func() time.Time { + return time.Now().UTC() + }, + } + rt.processMu.Lock() + defer rt.processMu.Unlock() + + processed := 0 + for processed < maxSteps { + if _, err := rt.queries.FindNextRunnableJob(ctx); err != nil { + if errors.Is(err, sql.ErrNoRows) { + break + } + return processed, err + } + if err := worker.RunOnce(ctx, rt.queries, opts); err != nil { + return processed, err + } + processed++ + } + return processed, nil +} + +func (rt *MapRuntime) ProcessRunToTerminal(ctx context.Context, runID ids.UUID, maxSteps int) (memsqlc.MapRun, error) { + if maxSteps <= 0 { + maxSteps = 2000 + } + for i := 0; i < maxSteps; i++ { + run, err := rt.queries.GetMapRunByID(ctx, memsqlc.GetMapRunByIDParams{ID: runID}) + if err != nil { + return memsqlc.MapRun{}, err + } + if isTerminalMapRunStatus(run.Status) { + return run, nil + } + processed, err := rt.ProcessPending(ctx, 1) + if err != nil { + return memsqlc.MapRun{}, err + } + if processed == 0 { + select { + case <-ctx.Done(): + return memsqlc.MapRun{}, ctx.Err() + case <-time.After(10 * time.Millisecond): + } + } + } + return memsqlc.MapRun{}, fmt.Errorf("map run %s did not reach terminal status within %d steps", runID.String(), maxSteps) +} + +func (rt *MapRuntime) GetRun(ctx context.Context, runID ids.UUID) (memsqlc.MapRun, error) { + if rt == nil || rt.queries == nil { + return memsqlc.MapRun{}, fmt.Errorf("map runtime is not configured") + } + return rt.queries.GetMapRunByID(ctx, memsqlc.GetMapRunByIDParams{ID: runID}) +} + +func (rt *MapRuntime) ReadRunItems(ctx context.Context, runID ids.UUID, offset, limit int) ([]MapItemProjection, error) { + if rt == nil || rt.queries == nil { + return nil, fmt.Errorf("map runtime is not configured") + } + rows, err := rt.queries.ListMapItemsByRunPaged(ctx, memsqlc.ListMapItemsByRunPagedParams{ + RunID: runID, + Lim: int64(limit), + Off: int64(offset), + }) + if err != nil { + return nil, err + } + items := make([]MapItemProjection, 0, len(rows)) + for _, row := range rows { + inputRec, inErr := DecodeMapItemInputFlatBuffer(row.InputFb) + if inErr != nil { + return nil, fmt.Errorf("decode item input %s: %w", row.ID.String(), inErr) + } + var inputAny interface{} + if err := jsonv2.Unmarshal([]byte(inputRec.ItemJSON), &inputAny); err != nil { + return nil, fmt.Errorf("decode item JSON %s: %w", row.ID.String(), err) + } + + proj := MapItemProjection{ + ItemID: row.ID.String(), + Index: row.ItemIndex, + Status: row.Status, + Attempts: row.Attempts, + Input: inputAny, + } + if row.InputHash != nil { + proj.InputHash = *row.InputHash + } + if row.LastError != nil { + proj.Error = *row.LastError + } + + if len(row.OutputFb) > 0 { + outputRec, outErr := DecodeMapItemOutputFlatBuffer(row.OutputFb) + if outErr != nil { + return nil, fmt.Errorf("decode item output %s: %w", row.ID.String(), outErr) + } + if row.OutputHash != nil && outputRec.OutputHash != "" && *row.OutputHash != outputRec.OutputHash { + return nil, fmt.Errorf("output hash mismatch for item %s", row.ID.String()) + } + if outputRec.ErrorText != "" && proj.Error == "" { + proj.Error = outputRec.ErrorText + } + if outputRec.OutputHash != "" { + proj.OutputHash = outputRec.OutputHash + } else if row.OutputHash != nil { + proj.OutputHash = *row.OutputHash + } + if strings.TrimSpace(outputRec.OutputJSON) != "" { + var outputAny interface{} + if err := jsonv2.Unmarshal([]byte(outputRec.OutputJSON), &outputAny); err == nil { + proj.Output = outputAny + } else { + // Keep malformed model output visible to callers for debugging. + proj.Output = outputRec.OutputJSON + } + } + } + items = append(items, proj) + } + return items, nil +} + +func (rt *MapRuntime) handleLLMMapJob(ctx context.Context, q *memsqlc.Queries, job memsqlc.Job) error { + if rt.llmModel == nil { + return fmt.Errorf("llm_map runtime model is not configured") + } + item, run, spec, inputRec, err := rt.loadMapJobState(ctx, q, job) + if err != nil { + return err + } + if spec.OperatorKind != MapOperatorLLM { + errMsg := fmt.Sprintf("run %s operator kind mismatch: expected llm_map got %s", run.ID.String(), spec.OperatorKind) + _, _ = q.MarkMapItemFailed(ctx, memsqlc.MarkMapItemFailedParams{LastError: &errMsg, ID: item.ID}) + if _, refreshErr := rt.refreshRunProgress(ctx, q, run.ID, &errMsg); refreshErr != nil { + logger.WarnCF("map_runtime", "failed to refresh run progress", map[string]interface{}{ + "run_id": run.ID.String(), + "error": refreshErr.Error(), + }) + } + return nil + } + + outJSON, callErr := rt.executeLLMMapItem(ctx, spec, inputRec.ItemJSON) + if callErr != nil { + if job.Attempts >= job.MaxAttempts { + msg := callErr.Error() + _, _ = q.MarkMapItemFailed(ctx, memsqlc.MarkMapItemFailedParams{LastError: &msg, ID: item.ID}) + if _, refreshErr := rt.refreshRunProgress(ctx, q, run.ID, &msg); refreshErr != nil { + logger.WarnCF("map_runtime", "failed to refresh run progress", map[string]interface{}{ + "run_id": run.ID.String(), + "error": refreshErr.Error(), + }) + } + return nil + } + return callErr + } + + outHash := hashString(outJSON) + outputFB, err := EncodeMapItemOutputFlatBuffer(MapItemOutputRecord{ + ItemIndex: uint32(item.ItemIndex), + Success: true, + Attempts: uint16(item.Attempts), + OutputJSON: outJSON, + OutputHash: outHash, + }) + if err != nil { + msg := fmt.Sprintf("encode output for item %s: %v", item.ID.String(), err) + _, _ = q.MarkMapItemFailed(ctx, memsqlc.MarkMapItemFailedParams{LastError: &msg, ID: item.ID}) + if _, refreshErr := rt.refreshRunProgress(ctx, q, run.ID, &msg); refreshErr != nil { + logger.WarnCF("map_runtime", "failed to refresh run progress", map[string]interface{}{ + "run_id": run.ID.String(), + "error": refreshErr.Error(), + }) + } + return nil + } + if _, err := q.MarkMapItemSucceeded(ctx, memsqlc.MarkMapItemSucceededParams{ + OutputFb: outputFB, + OutputHash: &outHash, + ID: item.ID, + }); err != nil { + return err + } + if _, refreshErr := rt.refreshRunProgress(ctx, q, run.ID, nil); refreshErr != nil { + logger.WarnCF("map_runtime", "failed to refresh run progress", map[string]interface{}{ + "run_id": run.ID.String(), + "error": refreshErr.Error(), + }) + } + return nil +} + +func (rt *MapRuntime) handleAgenticMapJob(ctx context.Context, q *memsqlc.Queries, job memsqlc.Job) error { + if rt.subagentManager == nil { + return fmt.Errorf("agentic_map runtime manager is not configured") + } + item, run, spec, inputRec, err := rt.loadMapJobState(ctx, q, job) + if err != nil { + return err + } + if spec.OperatorKind != MapOperatorAgentic { + errMsg := fmt.Sprintf("run %s operator kind mismatch: expected agentic_map got %s", run.ID.String(), spec.OperatorKind) + _, _ = q.MarkMapItemFailed(ctx, memsqlc.MarkMapItemFailedParams{LastError: &errMsg, ID: item.ID}) + if _, refreshErr := rt.refreshRunProgress(ctx, q, run.ID, &errMsg); refreshErr != nil { + logger.WarnCF("map_runtime", "failed to refresh run progress", map[string]interface{}{ + "run_id": run.ID.String(), + "error": refreshErr.Error(), + }) + } + return nil + } + outJSON, callErr := rt.executeAgenticMapItem(ctx, spec, int(item.ItemIndex), inputRec.ItemJSON) + if callErr != nil { + if job.Attempts >= job.MaxAttempts { + msg := callErr.Error() + _, _ = q.MarkMapItemFailed(ctx, memsqlc.MarkMapItemFailedParams{LastError: &msg, ID: item.ID}) + if _, refreshErr := rt.refreshRunProgress(ctx, q, run.ID, &msg); refreshErr != nil { + logger.WarnCF("map_runtime", "failed to refresh run progress", map[string]interface{}{ + "run_id": run.ID.String(), + "error": refreshErr.Error(), + }) + } + return nil + } + return callErr + } + + outHash := hashString(outJSON) + outputFB, err := EncodeMapItemOutputFlatBuffer(MapItemOutputRecord{ + ItemIndex: uint32(item.ItemIndex), + Success: true, + Attempts: uint16(item.Attempts), + OutputJSON: outJSON, + OutputHash: outHash, + }) + if err != nil { + msg := fmt.Sprintf("encode output for item %s: %v", item.ID.String(), err) + _, _ = q.MarkMapItemFailed(ctx, memsqlc.MarkMapItemFailedParams{LastError: &msg, ID: item.ID}) + if _, refreshErr := rt.refreshRunProgress(ctx, q, run.ID, &msg); refreshErr != nil { + logger.WarnCF("map_runtime", "failed to refresh run progress", map[string]interface{}{ + "run_id": run.ID.String(), + "error": refreshErr.Error(), + }) + } + return nil + } + if _, err := q.MarkMapItemSucceeded(ctx, memsqlc.MarkMapItemSucceededParams{ + OutputFb: outputFB, + OutputHash: &outHash, + ID: item.ID, + }); err != nil { + return err + } + if _, refreshErr := rt.refreshRunProgress(ctx, q, run.ID, nil); refreshErr != nil { + logger.WarnCF("map_runtime", "failed to refresh run progress", map[string]interface{}{ + "run_id": run.ID.String(), + "error": refreshErr.Error(), + }) + } + return nil +} + +func (rt *MapRuntime) loadMapJobState( + ctx context.Context, + q *memsqlc.Queries, + job memsqlc.Job, +) (memsqlc.MapItem, memsqlc.MapRun, MapRunSpec, MapItemInputRecord, error) { + itemID, dedupeErr := mapJobItemIDFromDedupeKey(ctx, q, job.DedupeKey) + if dedupeErr != nil { + payload, payloadErr := parseMapJobPayload(job.PayloadJson) + if payloadErr != nil { + return memsqlc.MapItem{}, memsqlc.MapRun{}, MapRunSpec{}, MapItemInputRecord{}, fmt.Errorf( + "invalid map job identity: dedupe_key_error=%v payload_error=%v", + dedupeErr, + payloadErr, + ) + } + parsedID, parseErr := ids.Parse(payload.ItemID) + if parseErr != nil { + return memsqlc.MapItem{}, memsqlc.MapRun{}, MapRunSpec{}, MapItemInputRecord{}, fmt.Errorf("invalid item id in job payload: %w", parseErr) + } + itemID = parsedID + } + item, err := q.MarkMapItemRunning(ctx, memsqlc.MarkMapItemRunningParams{ID: itemID}) + if err != nil { + return memsqlc.MapItem{}, memsqlc.MapRun{}, MapRunSpec{}, MapItemInputRecord{}, err + } + run, err := q.GetMapRunByID(ctx, memsqlc.GetMapRunByIDParams{ID: item.RunID}) + if err != nil { + return memsqlc.MapItem{}, memsqlc.MapRun{}, MapRunSpec{}, MapItemInputRecord{}, err + } + spec, err := DecodeMapRunSpecFlatBuffer(run.SpecFb) + if err != nil { + return memsqlc.MapItem{}, memsqlc.MapRun{}, MapRunSpec{}, MapItemInputRecord{}, fmt.Errorf("decode run spec: %w", err) + } + inputRec, err := DecodeMapItemInputFlatBuffer(item.InputFb) + if err != nil { + return memsqlc.MapItem{}, memsqlc.MapRun{}, MapRunSpec{}, MapItemInputRecord{}, fmt.Errorf("decode item input: %w", err) + } + return item, run, spec, inputRec, nil +} + +func (rt *MapRuntime) executeLLMMapItem(ctx context.Context, spec MapRunSpec, itemJSON string) (string, error) { + maxOut := int64(512) + prompt := fmt.Sprintf( + "Instruction:\n%s\n\nItem JSON:\n%s\n\nReturn exactly one JSON object and nothing else.", + spec.Instruction, + itemJSON, + ) + call := fantasy.Call{ + Prompt: fantasy.Prompt{ + fantasy.NewSystemMessage("You are a pure mapper. Return only strict JSON object output."), + fantasy.NewUserMessage(prompt), + }, + MaxOutputTokens: &maxOut, + } + resp, err := rt.llmModel.Generate(ctx, call) + if err != nil { + return "", err + } + outText := strings.TrimSpace(resp.Content.Text()) + var out map[string]interface{} + if err := jsonv2.Unmarshal([]byte(outText), &out); err != nil { + return "", fmt.Errorf("llm_map output is not valid JSON object: %w", err) + } + if strings.TrimSpace(spec.OutputSchemaJSON) != "" { + var schema map[string]interface{} + if err := jsonv2.Unmarshal([]byte(spec.OutputSchemaJSON), &schema); err != nil { + return "", fmt.Errorf("invalid output schema in run spec: %w", err) + } + if err := validateLLMMapOutputSchema(out, schema); err != nil { + return "", fmt.Errorf("schema validation failed: %w", err) + } + } + outJSON, err := canonicalJSONString(out) + if err != nil { + return "", fmt.Errorf("serialize model output: %w", err) + } + return outJSON, nil +} + +func (rt *MapRuntime) executeAgenticMapItem(ctx context.Context, spec MapRunSpec, index int, itemJSON string) (string, error) { + if strings.TrimSpace(spec.TaskTemplate) == "" { + return "", fmt.Errorf("task_template is required for agentic_map") + } + if !strings.Contains(spec.TaskTemplate, "{{item_json}}") || !strings.Contains(spec.TaskTemplate, "{{index}}") { + return "", fmt.Errorf("task_template must include both {{item_json}} and {{index}} placeholders") + } + task := strings.ReplaceAll(spec.TaskTemplate, "{{item_json}}", itemJSON) + task = strings.ReplaceAll(task, "{{index}}", fmt.Sprintf("%d", index)) + + subTool := NewSubagentTool(rt.subagentManager) + originChannel := spec.OriginChannel + if originChannel == "" { + originChannel = "cli" + } + originChatID := spec.OriginChatID + if originChatID == "" { + originChatID = "direct" + } + subTool.SetContext(originChannel, originChatID) + + callArgs := map[string]interface{}{ + "task": task, + "label": fmt.Sprintf("agentic-map-%d", index), + } + if strings.TrimSpace(spec.DelegatedScope) != "" { + callArgs["delegated_scope"] = spec.DelegatedScope + } + if strings.TrimSpace(spec.KeptWork) != "" { + callArgs["kept_work"] = spec.KeptWork + } + + result := subTool.Execute(ctx, callArgs) + if result == nil { + return "", fmt.Errorf("subagent returned nil result") + } + if result.IsError { + if result.Err != nil { + return "", result.Err + } + return "", errors.New(result.ForLLM) + } + + outputPayload := map[string]interface{}{ + "for_user": result.ForUser, + "for_llm": result.ForLLM, + } + outJSON, err := canonicalJSONString(outputPayload) + if err != nil { + return "", fmt.Errorf("serialize agentic output: %w", err) + } + return outJSON, nil +} + +func (rt *MapRuntime) refreshRunProgress( + ctx context.Context, + q *memsqlc.Queries, + runID ids.UUID, + lastErr *string, +) (memsqlc.MapRun, error) { + total, err := q.CountMapItemsByRun(ctx, memsqlc.CountMapItemsByRunParams{RunID: runID}) + if err != nil { + return memsqlc.MapRun{}, err + } + statusCounts, err := q.CountMapItemsByRunAndStatus(ctx, memsqlc.CountMapItemsByRunAndStatusParams{RunID: runID}) + if err != nil { + return memsqlc.MapRun{}, err + } + + var queued, running, succeeded, failed int64 + for _, row := range statusCounts { + switch row.Status { + case mapItemStatusQueued: + queued = row.Count + case mapItemStatusRunning: + running = row.Count + case mapItemStatusSucceeded: + succeeded = row.Count + case mapItemStatusFailed: + failed = row.Count + } + } + + status := mapRunStatusQueued + if running > 0 { + status = mapRunStatusRunning + } + if queued > 0 && running == 0 { + status = mapRunStatusQueued + } + if queued == 0 && running == 0 { + switch { + case failed > 0: + status = mapRunStatusFailed + case succeeded == total: + status = mapRunStatusSucceeded + default: + status = mapRunStatusFailed + } + } + + var completedAt *time.Time + if isTerminalMapRunStatus(status) { + now := time.Now().UTC() + completedAt = &now + } + return q.UpdateMapRunProgress(ctx, memsqlc.UpdateMapRunProgressParams{ + Status: status, + QueuedItems: queued, + RunningItems: running, + SucceededItems: succeeded, + FailedItems: failed, + LastError: lastErr, + CompletedAt: completedAt, + ID: runID, + }) +} + +func (rt *MapRuntime) failRun(ctx context.Context, runID ids.UUID, lastErr *string) (memsqlc.MapRun, error) { + return rt.queries.UpdateMapRunProgress(ctx, memsqlc.UpdateMapRunProgressParams{ + Status: mapRunStatusFailed, + QueuedItems: 0, + RunningItems: 0, + SucceededItems: 0, + FailedItems: 0, + LastError: lastErr, + CompletedAt: ptrTime(time.Now().UTC()), + ID: runID, + }) +} + +func parseMapJobPayload(raw []byte) (mapJobPayload, error) { + var payload mapJobPayload + if err := jsonv2.Unmarshal(raw, &payload); err != nil { + return mapJobPayload{}, err + } + if payload.RunID == "" || payload.ItemID == "" { + return mapJobPayload{}, fmt.Errorf("run_id and item_id are required") + } + return payload, nil +} + +func mapJobItemIDFromDedupeKey(ctx context.Context, q *memsqlc.Queries, dedupeKey *string) (ids.UUID, error) { + runID, itemIndex, err := parseMapJobDedupeKey(dedupeKey) + if err != nil { + return ids.UUID{}, err + } + item, err := q.GetMapItemByRunAndIndex(ctx, memsqlc.GetMapItemByRunAndIndexParams{ + RunID: runID, + ItemIndex: itemIndex, + }) + if err != nil { + return ids.UUID{}, fmt.Errorf("resolve map item by dedupe key: %w", err) + } + return item.ID, nil +} + +func parseMapJobDedupeKey(dedupeKey *string) (ids.UUID, int64, error) { + if dedupeKey == nil || strings.TrimSpace(*dedupeKey) == "" { + return ids.UUID{}, 0, fmt.Errorf("dedupe key is required") + } + parts := strings.Split(strings.TrimSpace(*dedupeKey), ":") + if len(parts) != 3 || parts[0] != "map" { + return ids.UUID{}, 0, fmt.Errorf("invalid dedupe key format: %q", strings.TrimSpace(*dedupeKey)) + } + runID, err := ids.Parse(parts[1]) + if err != nil { + return ids.UUID{}, 0, fmt.Errorf("invalid run id in dedupe key: %w", err) + } + itemIndex, err := strconv.ParseInt(parts[2], 10, 64) + if err != nil { + return ids.UUID{}, 0, fmt.Errorf("invalid item index in dedupe key: %w", err) + } + if itemIndex < 0 { + return ids.UUID{}, 0, fmt.Errorf("item index in dedupe key must be non-negative") + } + return runID, itemIndex, nil +} + +func toMapRunProjection(run memsqlc.MapRun) MapRunProjection { + p := MapRunProjection{ + RunID: run.ID.String(), + OperatorKind: run.OperatorKind, + Status: run.Status, + TotalItems: run.TotalItems, + QueuedItems: run.QueuedItems, + RunningItems: run.RunningItems, + SucceededItems: run.SucceededItems, + FailedItems: run.FailedItems, + CreatedAt: run.CreatedAt.Format(time.RFC3339Nano), + UpdatedAt: run.UpdatedAt.Format(time.RFC3339Nano), + } + if run.LastError != nil { + p.LastError = *run.LastError + } + if run.CompletedAt != nil { + p.CompletedAt = run.CompletedAt.Format(time.RFC3339Nano) + } + return p +} + +func isTerminalMapRunStatus(status string) bool { + switch status { + case mapRunStatusSucceeded, mapRunStatusFailed, mapRunStatusCancelled: + return true + default: + return false + } +} + +func optionalStringPtr(v string) *string { + if strings.TrimSpace(v) == "" { + return nil + } + out := strings.TrimSpace(v) + return &out +} + +func isUniqueConstraintError(err error) bool { + if err == nil { + return false + } + lower := strings.ToLower(err.Error()) + return strings.Contains(lower, "unique constraint") || + strings.Contains(lower, "constraint failed") || + strings.Contains(lower, "duplicate key") +} + +func ptrTime(v time.Time) *time.Time { + return &v +} diff --git a/pkg/tools/map_runtime_integration_test.go b/pkg/tools/map_runtime_integration_test.go new file mode 100644 index 000000000..a8c4a82ee --- /dev/null +++ b/pkg/tools/map_runtime_integration_test.go @@ -0,0 +1,383 @@ +package tools + +import ( + "context" + "fmt" + "sync" + "testing" + + fantasy "charm.land/fantasy" + jsonv2 "github.com/go-json-experiment/json" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/ZanzyTHEbar/dragonscale/pkg/bus" + "github.com/ZanzyTHEbar/dragonscale/pkg/ids" + "github.com/ZanzyTHEbar/dragonscale/pkg/memory/delegate" + memsqlc "github.com/ZanzyTHEbar/dragonscale/pkg/memory/sqlc" +) + +type scriptedMapLLM struct { + mu sync.Mutex + responses []scriptedMapLLMResponse + index int +} + +type scriptedMapLLMResponse struct { + Text string + Err error +} + +func (m *scriptedMapLLM) Generate(_ context.Context, _ fantasy.Call) (*fantasy.Response, error) { + m.mu.Lock() + defer m.mu.Unlock() + if len(m.responses) == 0 { + return nil, fmt.Errorf("no scripted llm responses") + } + resp := m.responses[m.index%len(m.responses)] + m.index++ + if resp.Err != nil { + return nil, resp.Err + } + return &fantasy.Response{ + Content: fantasy.ResponseContent{fantasy.TextContent{Text: resp.Text}}, + FinishReason: fantasy.FinishReasonStop, + }, nil +} + +func (m *scriptedMapLLM) Stream(ctx context.Context, call fantasy.Call) (fantasy.StreamResponse, error) { + resp, err := m.Generate(ctx, call) + if err != nil { + return nil, err + } + return func(yield func(fantasy.StreamPart) bool) { + if !yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextDelta, Delta: resp.Content.Text()}) { + return + } + yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop}) + }, nil +} + +func (m *scriptedMapLLM) GenerateObject(_ context.Context, _ fantasy.ObjectCall) (*fantasy.ObjectResponse, error) { + return nil, fmt.Errorf("not implemented") +} + +func (m *scriptedMapLLM) StreamObject(_ context.Context, _ fantasy.ObjectCall) (fantasy.ObjectStreamResponse, error) { + return nil, fmt.Errorf("not implemented") +} + +func (m *scriptedMapLLM) Provider() string { return "mock" } +func (m *scriptedMapLLM) Model() string { return "mock-map-llm" } + +func newMapRuntimeForTest(t *testing.T, llm fantasy.LanguageModel, manager *SubagentManager) *MapRuntime { + t.Helper() + d, err := delegate.NewLibSQLInMemory() + require.NoError(t, err) + require.NoError(t, d.Init(context.Background())) + t.Cleanup(func() { _ = d.Close() }) + + return NewMapRuntime(d.Queries(), "dragonscale", llm, "mock-map-llm", manager) +} + +func decodeResultMap(t *testing.T, result *ToolResult) map[string]interface{} { + t.Helper() + require.NotNil(t, result) + require.False(t, result.IsError, result.ForLLM) + var payload map[string]interface{} + require.NoError(t, jsonv2.Unmarshal([]byte(result.ForLLM), &payload)) + return payload +} + +func TestLLMMap_WorkerJSONL_StatusAndRead(t *testing.T) { + model := &scriptedMapLLM{ + responses: []scriptedMapLLMResponse{ + {Text: `{"label":"alpha","priority":1}`}, + {Text: `{"label":"beta","priority":2}`}, + }, + } + runtime := newMapRuntimeForTest(t, model, nil) + mapTool := NewLLMMapTool(model, "mock-map-llm") + mapTool.SetRuntime(runtime) + statusTool := NewMapRunStatusTool(runtime) + readTool := NewMapRunReadTool(runtime) + + enqueue := decodeResultMap(t, mapTool.Execute(context.Background(), map[string]interface{}{ + "instruction": "normalize to {label, priority}", + "input_jsonl": "{\"name\":\"a\"}\n{\"name\":\"b\"}", + "execution_mode": "worker", + })) + runID, _ := enqueue["run_id"].(string) + require.NotEmpty(t, runID) + assert.Equal(t, "worker", enqueue["execution_mode"]) + + status := decodeResultMap(t, statusTool.Execute(context.Background(), map[string]interface{}{ + "run_id": runID, + "process_steps": float64(20), + })) + assert.Equal(t, mapRunStatusSucceeded, status["status"]) + assert.EqualValues(t, 2, status["succeeded_items"]) + + readJSON := decodeResultMap(t, readTool.Execute(context.Background(), map[string]interface{}{ + "run_id": runID, + "limit": float64(10), + "format": "json", + })) + itemsAny, ok := readJSON["items"].([]interface{}) + require.True(t, ok) + require.Len(t, itemsAny, 2) + + readJSONL := decodeResultMap(t, readTool.Execute(context.Background(), map[string]interface{}{ + "run_id": runID, + "format": "jsonl", + })) + jsonlText, _ := readJSONL["items_jsonl"].(string) + assert.Contains(t, jsonlText, `"status":"succeeded"`) + assert.Contains(t, jsonlText, "\n") +} + +func TestLLMMap_IdempotencyReuse(t *testing.T) { + model := &scriptedMapLLM{ + responses: []scriptedMapLLMResponse{ + {Text: `{"label":"alpha"}`}, + }, + } + runtime := newMapRuntimeForTest(t, model, nil) + mapTool := NewLLMMapTool(model, "mock-map-llm") + mapTool.SetRuntime(runtime) + + args := map[string]interface{}{ + "instruction": "normalize", + "items": []interface{}{map[string]interface{}{"name": "a"}}, + "execution_mode": "worker", + "idempotency_key": "run-1", + "session_key": "sess-a", + } + first := decodeResultMap(t, mapTool.Execute(context.Background(), args)) + second := decodeResultMap(t, mapTool.Execute(context.Background(), args)) + + assert.Equal(t, first["run_id"], second["run_id"]) + assert.Equal(t, true, second["idempotent_reuse"]) +} + +func TestLLMMap_IdempotencyReuse_Concurrent(t *testing.T) { + model := &scriptedMapLLM{ + responses: []scriptedMapLLMResponse{ + {Text: `{"label":"alpha"}`}, + }, + } + runtime := newMapRuntimeForTest(t, model, nil) + mapTool := NewLLMMapTool(model, "mock-map-llm") + mapTool.SetRuntime(runtime) + + const workers = 12 + start := make(chan struct{}) + var wg sync.WaitGroup + runIDs := make(chan string, workers) + errs := make(chan error, workers) + + for i := 0; i < workers; i++ { + wg.Add(1) + go func() { + defer wg.Done() + <-start + result := mapTool.Execute(context.Background(), map[string]interface{}{ + "instruction": "normalize", + "items": []interface{}{map[string]interface{}{"name": "a"}}, + "execution_mode": "worker", + "idempotency_key": "run-1", + "session_key": "sess-a", + }) + if result == nil || result.IsError { + if result != nil && result.Err != nil { + errs <- result.Err + return + } + errs <- fmt.Errorf("unexpected map tool error: %v", result) + return + } + var payload map[string]interface{} + if err := jsonv2.Unmarshal([]byte(result.ForLLM), &payload); err != nil { + errs <- fmt.Errorf("decode result: %w", err) + return + } + runID, _ := payload["run_id"].(string) + if runID == "" { + errs <- fmt.Errorf("missing run_id in payload") + return + } + runIDs <- runID + }() + } + + close(start) + wg.Wait() + close(errs) + close(runIDs) + + for err := range errs { + require.NoError(t, err) + } + + unique := make(map[string]struct{}) + for runID := range runIDs { + unique[runID] = struct{}{} + } + require.Len(t, unique, 1, "all concurrent requests should reuse the same run id") +} + +func TestLLMMap_InlineRetriesThenSucceeds(t *testing.T) { + model := &scriptedMapLLM{ + responses: []scriptedMapLLMResponse{ + {Err: fmt.Errorf("transient model outage")}, + {Text: `{"label":"ok"}`}, + }, + } + runtime := newMapRuntimeForTest(t, model, nil) + mapTool := NewLLMMapTool(model, "mock-map-llm") + mapTool.SetRuntime(runtime) + readTool := NewMapRunReadTool(runtime) + + result := decodeResultMap(t, mapTool.Execute(context.Background(), map[string]interface{}{ + "instruction": "normalize", + "items": []interface{}{map[string]interface{}{"name": "retry"}}, + "execution_mode": "inline", + "max_retries": float64(2), + })) + assert.Equal(t, mapRunStatusSucceeded, result["status"]) + runID, _ := result["run_id"].(string) + require.NotEmpty(t, runID) + + read := decodeResultMap(t, readTool.Execute(context.Background(), map[string]interface{}{ + "run_id": runID, + "format": "json", + })) + itemsAny, ok := read["items"].([]interface{}) + require.True(t, ok) + require.Len(t, itemsAny, 1) + item0 := itemsAny[0].(map[string]interface{}) + assert.EqualValues(t, 2, item0["attempts"]) + assert.Equal(t, mapItemStatusSucceeded, item0["status"]) +} + +func TestLLMMap_InlineExhaustedRetriesFailsRun(t *testing.T) { + model := &scriptedMapLLM{ + responses: []scriptedMapLLMResponse{ + {Text: `not-json`}, + {Text: `still-not-json`}, + }, + } + runtime := newMapRuntimeForTest(t, model, nil) + mapTool := NewLLMMapTool(model, "mock-map-llm") + mapTool.SetRuntime(runtime) + + payload := decodeResultMap(t, mapTool.Execute(context.Background(), map[string]interface{}{ + "instruction": "normalize", + "items": []interface{}{map[string]interface{}{"name": "bad"}}, + "execution_mode": "inline", + "max_retries": float64(1), + })) + assert.Equal(t, mapRunStatusFailed, payload["status"]) + summary, ok := payload["summary"].(map[string]interface{}) + require.True(t, ok) + assert.EqualValues(t, 1, summary["failure_count"]) +} + +func TestMapRunRead_DetectsOutputHashMismatch(t *testing.T) { + model := &scriptedMapLLM{ + responses: []scriptedMapLLMResponse{ + {Text: `{"label":"alpha"}`}, + }, + } + runtime := newMapRuntimeForTest(t, model, nil) + mapTool := NewLLMMapTool(model, "mock-map-llm") + mapTool.SetRuntime(runtime) + readTool := NewMapRunReadTool(runtime) + + result := decodeResultMap(t, mapTool.Execute(context.Background(), map[string]interface{}{ + "instruction": "normalize", + "items": []interface{}{map[string]interface{}{"name": "a"}}, + "execution_mode": "inline", + })) + runIDText, _ := result["run_id"].(string) + require.NotEmpty(t, runIDText) + + runID, err := ids.Parse(runIDText) + require.NoError(t, err) + rows, err := runtime.queries.ListMapItemsByRunPaged(context.Background(), memsqlc.ListMapItemsByRunPagedParams{ + RunID: runID, + Lim: 10, + Off: 0, + }) + require.NoError(t, err) + require.Len(t, rows, 1) + + _, err = runtime.queries.MarkMapItemSucceeded(context.Background(), memsqlc.MarkMapItemSucceededParams{ + OutputFb: rows[0].OutputFb, + OutputHash: optionalStringPtr("forced-mismatch"), + ID: rows[0].ID, + }) + require.NoError(t, err) + + readResult := readTool.Execute(context.Background(), map[string]interface{}{ + "run_id": runIDText, + "format": "json", + }) + require.NotNil(t, readResult) + require.True(t, readResult.IsError) + assert.Contains(t, readResult.ForLLM, "output hash mismatch") +} + +func TestAgenticMap_WorkerLifecycle(t *testing.T) { + provider := &MockLanguageModel{} + manager := NewSubagentManager(provider, "test-model", "/tmp/test", bus.NewMessageBus()) + manager.SetRunLoop(func(_ context.Context, _ ToolLoopConfig, _, userPrompt, _, _ string) (*ToolLoopResult, error) { + return &ToolLoopResult{Content: "processed: " + userPrompt, Iterations: 1}, nil + }) + + runtime := newMapRuntimeForTest(t, &scriptedMapLLM{responses: []scriptedMapLLMResponse{{Text: `{"unused":true}`}}}, manager) + tool := NewAgenticMapTool(manager) + tool.SetRuntime(runtime) + statusTool := NewMapRunStatusTool(runtime) + readTool := NewMapRunReadTool(runtime) + + enqueue := decodeResultMap(t, tool.Execute(context.Background(), map[string]interface{}{ + "items": []interface{}{map[string]interface{}{"name": "x"}}, + "task_template": "Handle {{index}} => {{item_json}}", + "execution_mode": "worker", + })) + runID, _ := enqueue["run_id"].(string) + require.NotEmpty(t, runID) + + status := decodeResultMap(t, statusTool.Execute(context.Background(), map[string]interface{}{ + "run_id": runID, + "process_steps": float64(20), + })) + assert.Equal(t, mapRunStatusSucceeded, status["status"]) + + read := decodeResultMap(t, readTool.Execute(context.Background(), map[string]interface{}{ + "run_id": runID, + "format": "json", + })) + itemsAny, ok := read["items"].([]interface{}) + require.True(t, ok) + require.Len(t, itemsAny, 1) + item0 := itemsAny[0].(map[string]interface{}) + assert.Equal(t, mapItemStatusSucceeded, item0["status"]) +} + +func TestLLMMap_InvalidJSONLIngestFails(t *testing.T) { + model := &scriptedMapLLM{ + responses: []scriptedMapLLMResponse{{Text: `{"ok":true}`}}, + } + runtime := newMapRuntimeForTest(t, model, nil) + mapTool := NewLLMMapTool(model, "mock-map-llm") + mapTool.SetRuntime(runtime) + + result := mapTool.Execute(context.Background(), map[string]interface{}{ + "instruction": "normalize", + "input_jsonl": "{\"name\":\"ok\"}\nnot-json", + }) + require.NotNil(t, result) + require.True(t, result.IsError) + assert.Contains(t, result.ForLLM, "invalid JSONL") +} diff --git a/pkg/tools/mapopsfb/MapItemInput.go b/pkg/tools/mapopsfb/MapItemInput.go new file mode 100644 index 000000000..b05dcce31 --- /dev/null +++ b/pkg/tools/mapopsfb/MapItemInput.go @@ -0,0 +1,93 @@ +// Code generated by the FlatBuffers compiler. DO NOT EDIT. + +package mapopsfb + +import ( + flatbuffers "github.com/google/flatbuffers/go" +) + +type MapItemInput struct { + _tab flatbuffers.Table +} + +func GetRootAsMapItemInput(buf []byte, offset flatbuffers.UOffsetT) *MapItemInput { + n := flatbuffers.GetUOffsetT(buf[offset:]) + x := &MapItemInput{} + x.Init(buf, n+offset) + return x +} + +func GetSizePrefixedRootAsMapItemInput(buf []byte, offset flatbuffers.UOffsetT) *MapItemInput { + n := flatbuffers.GetUOffsetT(buf[offset+flatbuffers.SizeUint32:]) + x := &MapItemInput{} + x.Init(buf, n+offset+flatbuffers.SizeUint32) + return x +} + +func (rcv *MapItemInput) Init(buf []byte, i flatbuffers.UOffsetT) { + rcv._tab.Bytes = buf + rcv._tab.Pos = i +} + +func (rcv *MapItemInput) Table() flatbuffers.Table { + return rcv._tab +} + +func (rcv *MapItemInput) Version() uint16 { + o := flatbuffers.UOffsetT(rcv._tab.Offset(4)) + if o != 0 { + return rcv._tab.GetUint16(o + rcv._tab.Pos) + } + return 1 +} + +func (rcv *MapItemInput) MutateVersion(n uint16) bool { + return rcv._tab.MutateUint16Slot(4, n) +} + +func (rcv *MapItemInput) ItemIndex() uint32 { + o := flatbuffers.UOffsetT(rcv._tab.Offset(6)) + if o != 0 { + return rcv._tab.GetUint32(o + rcv._tab.Pos) + } + return 0 +} + +func (rcv *MapItemInput) MutateItemIndex(n uint32) bool { + return rcv._tab.MutateUint32Slot(6, n) +} + +func (rcv *MapItemInput) ItemJson() []byte { + o := flatbuffers.UOffsetT(rcv._tab.Offset(8)) + if o != 0 { + return rcv._tab.ByteVector(o + rcv._tab.Pos) + } + return nil +} + +func (rcv *MapItemInput) InputHash() []byte { + o := flatbuffers.UOffsetT(rcv._tab.Offset(10)) + if o != 0 { + return rcv._tab.ByteVector(o + rcv._tab.Pos) + } + return nil +} + +func MapItemInputStart(builder *flatbuffers.Builder) { + builder.StartObject(4) +} +func MapItemInputAddVersion(builder *flatbuffers.Builder, version uint16) { + builder.PrependUint16Slot(0, version, 1) +} +func MapItemInputAddItemIndex(builder *flatbuffers.Builder, itemIndex uint32) { + builder.PrependUint32Slot(1, itemIndex, 0) +} +func MapItemInputAddItemJson(builder *flatbuffers.Builder, itemJson flatbuffers.UOffsetT) { + builder.PrependUOffsetTSlot(2, flatbuffers.UOffsetT(itemJson), 0) +} +func MapItemInputAddInputHash(builder *flatbuffers.Builder, inputHash flatbuffers.UOffsetT) { + builder.PrependUOffsetTSlot(3, flatbuffers.UOffsetT(inputHash), 0) +} +func MapItemInputEnd(builder *flatbuffers.Builder) flatbuffers.UOffsetT { + return builder.EndObject() +} diff --git a/pkg/tools/mapopsfb/MapItemOutput.go b/pkg/tools/mapopsfb/MapItemOutput.go new file mode 100644 index 000000000..8cdf17bf4 --- /dev/null +++ b/pkg/tools/mapopsfb/MapItemOutput.go @@ -0,0 +1,134 @@ +// Code generated by the FlatBuffers compiler. DO NOT EDIT. + +package mapopsfb + +import ( + flatbuffers "github.com/google/flatbuffers/go" +) + +type MapItemOutput struct { + _tab flatbuffers.Table +} + +func GetRootAsMapItemOutput(buf []byte, offset flatbuffers.UOffsetT) *MapItemOutput { + n := flatbuffers.GetUOffsetT(buf[offset:]) + x := &MapItemOutput{} + x.Init(buf, n+offset) + return x +} + +func GetSizePrefixedRootAsMapItemOutput(buf []byte, offset flatbuffers.UOffsetT) *MapItemOutput { + n := flatbuffers.GetUOffsetT(buf[offset+flatbuffers.SizeUint32:]) + x := &MapItemOutput{} + x.Init(buf, n+offset+flatbuffers.SizeUint32) + return x +} + +func (rcv *MapItemOutput) Init(buf []byte, i flatbuffers.UOffsetT) { + rcv._tab.Bytes = buf + rcv._tab.Pos = i +} + +func (rcv *MapItemOutput) Table() flatbuffers.Table { + return rcv._tab +} + +func (rcv *MapItemOutput) Version() uint16 { + o := flatbuffers.UOffsetT(rcv._tab.Offset(4)) + if o != 0 { + return rcv._tab.GetUint16(o + rcv._tab.Pos) + } + return 1 +} + +func (rcv *MapItemOutput) MutateVersion(n uint16) bool { + return rcv._tab.MutateUint16Slot(4, n) +} + +func (rcv *MapItemOutput) ItemIndex() uint32 { + o := flatbuffers.UOffsetT(rcv._tab.Offset(6)) + if o != 0 { + return rcv._tab.GetUint32(o + rcv._tab.Pos) + } + return 0 +} + +func (rcv *MapItemOutput) MutateItemIndex(n uint32) bool { + return rcv._tab.MutateUint32Slot(6, n) +} + +func (rcv *MapItemOutput) Success() bool { + o := flatbuffers.UOffsetT(rcv._tab.Offset(8)) + if o != 0 { + return rcv._tab.GetBool(o + rcv._tab.Pos) + } + return false +} + +func (rcv *MapItemOutput) MutateSuccess(n bool) bool { + return rcv._tab.MutateBoolSlot(8, n) +} + +func (rcv *MapItemOutput) Attempts() uint16 { + o := flatbuffers.UOffsetT(rcv._tab.Offset(10)) + if o != 0 { + return rcv._tab.GetUint16(o + rcv._tab.Pos) + } + return 0 +} + +func (rcv *MapItemOutput) MutateAttempts(n uint16) bool { + return rcv._tab.MutateUint16Slot(10, n) +} + +func (rcv *MapItemOutput) OutputJson() []byte { + o := flatbuffers.UOffsetT(rcv._tab.Offset(12)) + if o != 0 { + return rcv._tab.ByteVector(o + rcv._tab.Pos) + } + return nil +} + +func (rcv *MapItemOutput) ErrorText() []byte { + o := flatbuffers.UOffsetT(rcv._tab.Offset(14)) + if o != 0 { + return rcv._tab.ByteVector(o + rcv._tab.Pos) + } + return nil +} + +func (rcv *MapItemOutput) OutputHash() []byte { + o := flatbuffers.UOffsetT(rcv._tab.Offset(16)) + if o != 0 { + return rcv._tab.ByteVector(o + rcv._tab.Pos) + } + return nil +} + +func MapItemOutputStart(builder *flatbuffers.Builder) { + builder.StartObject(7) +} +func MapItemOutputAddVersion(builder *flatbuffers.Builder, version uint16) { + builder.PrependUint16Slot(0, version, 1) +} +func MapItemOutputAddItemIndex(builder *flatbuffers.Builder, itemIndex uint32) { + builder.PrependUint32Slot(1, itemIndex, 0) +} +func MapItemOutputAddSuccess(builder *flatbuffers.Builder, success bool) { + builder.PrependBoolSlot(2, success, false) +} +func MapItemOutputAddAttempts(builder *flatbuffers.Builder, attempts uint16) { + builder.PrependUint16Slot(3, attempts, 0) +} +func MapItemOutputAddOutputJson(builder *flatbuffers.Builder, outputJson flatbuffers.UOffsetT) { + builder.PrependUOffsetTSlot(4, flatbuffers.UOffsetT(outputJson), 0) +} +func MapItemOutputAddErrorText(builder *flatbuffers.Builder, errorText flatbuffers.UOffsetT) { + builder.PrependUOffsetTSlot(5, flatbuffers.UOffsetT(errorText), 0) +} +func MapItemOutputAddOutputHash(builder *flatbuffers.Builder, outputHash flatbuffers.UOffsetT) { + builder.PrependUOffsetTSlot(6, flatbuffers.UOffsetT(outputHash), 0) +} +func MapItemOutputEnd(builder *flatbuffers.Builder) flatbuffers.UOffsetT { + return builder.EndObject() +} diff --git a/pkg/tools/mapopsfb/MapRunSpec.go b/pkg/tools/mapopsfb/MapRunSpec.go new file mode 100644 index 000000000..e9a97f258 --- /dev/null +++ b/pkg/tools/mapopsfb/MapRunSpec.go @@ -0,0 +1,159 @@ +// Code generated by the FlatBuffers compiler. DO NOT EDIT. + +package mapopsfb + +import ( + flatbuffers "github.com/google/flatbuffers/go" +) + +type MapRunSpec struct { + _tab flatbuffers.Table +} + +func GetRootAsMapRunSpec(buf []byte, offset flatbuffers.UOffsetT) *MapRunSpec { + n := flatbuffers.GetUOffsetT(buf[offset:]) + x := &MapRunSpec{} + x.Init(buf, n+offset) + return x +} + +func GetSizePrefixedRootAsMapRunSpec(buf []byte, offset flatbuffers.UOffsetT) *MapRunSpec { + n := flatbuffers.GetUOffsetT(buf[offset+flatbuffers.SizeUint32:]) + x := &MapRunSpec{} + x.Init(buf, n+offset+flatbuffers.SizeUint32) + return x +} + +func (rcv *MapRunSpec) Init(buf []byte, i flatbuffers.UOffsetT) { + rcv._tab.Bytes = buf + rcv._tab.Pos = i +} + +func (rcv *MapRunSpec) Table() flatbuffers.Table { + return rcv._tab +} + +func (rcv *MapRunSpec) Version() uint16 { + o := flatbuffers.UOffsetT(rcv._tab.Offset(4)) + if o != 0 { + return rcv._tab.GetUint16(o + rcv._tab.Pos) + } + return 1 +} + +func (rcv *MapRunSpec) MutateVersion(n uint16) bool { + return rcv._tab.MutateUint16Slot(4, n) +} + +func (rcv *MapRunSpec) OperatorKind() []byte { + o := flatbuffers.UOffsetT(rcv._tab.Offset(6)) + if o != 0 { + return rcv._tab.ByteVector(o + rcv._tab.Pos) + } + return nil +} + +func (rcv *MapRunSpec) Instruction() []byte { + o := flatbuffers.UOffsetT(rcv._tab.Offset(8)) + if o != 0 { + return rcv._tab.ByteVector(o + rcv._tab.Pos) + } + return nil +} + +func (rcv *MapRunSpec) TaskTemplate() []byte { + o := flatbuffers.UOffsetT(rcv._tab.Offset(10)) + if o != 0 { + return rcv._tab.ByteVector(o + rcv._tab.Pos) + } + return nil +} + +func (rcv *MapRunSpec) OutputSchemaJson() []byte { + o := flatbuffers.UOffsetT(rcv._tab.Offset(12)) + if o != 0 { + return rcv._tab.ByteVector(o + rcv._tab.Pos) + } + return nil +} + +func (rcv *MapRunSpec) MaxRetries() uint16 { + o := flatbuffers.UOffsetT(rcv._tab.Offset(14)) + if o != 0 { + return rcv._tab.GetUint16(o + rcv._tab.Pos) + } + return 1 +} + +func (rcv *MapRunSpec) MutateMaxRetries(n uint16) bool { + return rcv._tab.MutateUint16Slot(14, n) +} + +func (rcv *MapRunSpec) DelegatedScope() []byte { + o := flatbuffers.UOffsetT(rcv._tab.Offset(16)) + if o != 0 { + return rcv._tab.ByteVector(o + rcv._tab.Pos) + } + return nil +} + +func (rcv *MapRunSpec) KeptWork() []byte { + o := flatbuffers.UOffsetT(rcv._tab.Offset(18)) + if o != 0 { + return rcv._tab.ByteVector(o + rcv._tab.Pos) + } + return nil +} + +func (rcv *MapRunSpec) OriginChannel() []byte { + o := flatbuffers.UOffsetT(rcv._tab.Offset(20)) + if o != 0 { + return rcv._tab.ByteVector(o + rcv._tab.Pos) + } + return nil +} + +func (rcv *MapRunSpec) OriginChatId() []byte { + o := flatbuffers.UOffsetT(rcv._tab.Offset(22)) + if o != 0 { + return rcv._tab.ByteVector(o + rcv._tab.Pos) + } + return nil +} + +func MapRunSpecStart(builder *flatbuffers.Builder) { + builder.StartObject(10) +} +func MapRunSpecAddVersion(builder *flatbuffers.Builder, version uint16) { + builder.PrependUint16Slot(0, version, 1) +} +func MapRunSpecAddOperatorKind(builder *flatbuffers.Builder, operatorKind flatbuffers.UOffsetT) { + builder.PrependUOffsetTSlot(1, flatbuffers.UOffsetT(operatorKind), 0) +} +func MapRunSpecAddInstruction(builder *flatbuffers.Builder, instruction flatbuffers.UOffsetT) { + builder.PrependUOffsetTSlot(2, flatbuffers.UOffsetT(instruction), 0) +} +func MapRunSpecAddTaskTemplate(builder *flatbuffers.Builder, taskTemplate flatbuffers.UOffsetT) { + builder.PrependUOffsetTSlot(3, flatbuffers.UOffsetT(taskTemplate), 0) +} +func MapRunSpecAddOutputSchemaJson(builder *flatbuffers.Builder, outputSchemaJson flatbuffers.UOffsetT) { + builder.PrependUOffsetTSlot(4, flatbuffers.UOffsetT(outputSchemaJson), 0) +} +func MapRunSpecAddMaxRetries(builder *flatbuffers.Builder, maxRetries uint16) { + builder.PrependUint16Slot(5, maxRetries, 1) +} +func MapRunSpecAddDelegatedScope(builder *flatbuffers.Builder, delegatedScope flatbuffers.UOffsetT) { + builder.PrependUOffsetTSlot(6, flatbuffers.UOffsetT(delegatedScope), 0) +} +func MapRunSpecAddKeptWork(builder *flatbuffers.Builder, keptWork flatbuffers.UOffsetT) { + builder.PrependUOffsetTSlot(7, flatbuffers.UOffsetT(keptWork), 0) +} +func MapRunSpecAddOriginChannel(builder *flatbuffers.Builder, originChannel flatbuffers.UOffsetT) { + builder.PrependUOffsetTSlot(8, flatbuffers.UOffsetT(originChannel), 0) +} +func MapRunSpecAddOriginChatId(builder *flatbuffers.Builder, originChatId flatbuffers.UOffsetT) { + builder.PrependUOffsetTSlot(9, flatbuffers.UOffsetT(originChatId), 0) +} +func MapRunSpecEnd(builder *flatbuffers.Builder) flatbuffers.UOffsetT { + return builder.EndObject() +}