Convert json.RawMessage to JSON strings instead of byte arrays (#41)

.gitignore

@@ -71,3 +71,6 @@ clicky
 formatters/.playwright-mcp/
 __debug*
 
+ginkgo.report
+*.bak
+examples/uber_demo/uber_demo

ai/agent.go

@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/flanksource/clicky/ai/cache"
+
 	"github.com/spf13/pflag"
 )
 
@@ -92,6 +93,7 @@ func NewAgentManager(config AgentConfig) *AgentManager {
 }
 
 func GetDefaultAgent() (Agent, error) {
+
 	manager := NewAgentManager(DefaultConfig())
 	return manager.GetDefaultAgent()
 }
@@ -111,6 +113,7 @@ func (am *AgentManager) GetAgent(agentType AgentType) (Agent, error) {
 		agent, err = NewClaudeAgent(am.config)
 	case AgentTypeAider:
 		agent, err = NewAiderAgent(am.config)
+
 	default:
 		return nil, fmt.Errorf("unsupported agent type: %s", agentType)
 	}
@@ -185,12 +188,13 @@ var defaultConfig AgentConfig = AgentConfig{
 	Type:          AgentTypeClaude,
 	Model:         "claude-haiku-4-5",
 	MaxTokens:     10000,
-	MaxConcurrent: 3,
+	MaxConcurrent: 4,
 	Debug:         false,
 	Verbose:       false,
-	Temperature:   0.2,
-	CacheTTL:      24 * time.Hour, // Default 24 hour TTL
-	NoCache:       false,
+	// Temperature:   0.2,
+	StrictMCPConfig: true,
+	CacheTTL:        24 * time.Hour, // Default 24 hour TTL
+	NoCache:         false,
 }
 
 // DefaultConfig returns a default agent configuration

ai/api.go

@@ -2,6 +2,7 @@ package ai
 
 import (
 	"context"
+	"encoding/json"
 	"time"
 
 	"github.com/flanksource/clicky/api"
@@ -40,27 +41,63 @@ type AgentConfig struct {
 	NoCache         bool          `json:"no_cache,omitempty"`
 }
 
+func (a AgentConfig) Pretty() api.Text {
+	t := api.Text{}.Append(string(a.Type))
+	if a.Model != "" {
+		t = t.Space().Append(a.Model, "font-mono")
+	}
+	if a.NoCache {
+		t = t.Space().Append(" No Cache", "text-orange-500")
+	}
+	return t
+}
+
 // PromptRequest represents a request to process a prompt
 type PromptRequest struct {
+	Context          map[string]string `json:"context,omitempty"`
+	Name             string            `json:"name"`
+	Prompt           string            `json:"prompt"`
+	StructuredOutput interface{}       `json:"structured_output,omitempty"` // Schema for structured JSON output
+}
+
+// TypedPromptRequest represents a type-safe request with structured output
+type TypedPromptRequest[T any] struct {
 	Context map[string]string `json:"context,omitempty"`
 	Name    string            `json:"name"`
 	Prompt  string            `json:"prompt"`
 }
 
 // PromptResponse represents the response from processing a prompt
 type PromptResponse struct {
-	Request PromptRequest `json:"request,omitempty"`
-	Result  string        `json:"result"`
-	Costs   Costs         `json:"costs,omitempty"`
-	Model   string        `json:"model,omitempty"`
-	Error   string        `json:"error,omitempty"`
+	Request        PromptRequest `json:"request,omitempty"`
+	Result         string        `json:"result"`
+	StructuredData interface{}   `json:"structured_data,omitempty"` // Populated if structured output was requested
+	Costs          Costs         `json:"costs,omitempty"`
+	Model          string        `json:"model,omitempty"`
+	Error          string        `json:"error,omitempty"`
 	// Total wall-clock duration of the request
 	Duration time.Duration `json:"duration,omitempty"`
 	// Duration spent in the model processing as reported by the API
 	DurationModel time.Duration `json:"duration_model,omitempty"`
 	CacheHit      bool          `json:"cache_hit,omitempty"`
 }
 
+// TypedPromptResponse represents a type-safe response with structured output
+type TypedPromptResponse[T any] struct {
+	Request       TypedPromptRequest[T] `json:"request,omitempty"`
+	Data          T                     `json:"data"`
+	Costs         Costs                 `json:"costs,omitempty"`
+	Model         string                `json:"model,omitempty"`
+	Error         string                `json:"error,omitempty"`
+	Duration      time.Duration         `json:"duration,omitempty"`
+	DurationModel time.Duration         `json:"duration_model,omitempty"`
+	CacheHit      bool                  `json:"cache_hit,omitempty"`
+}
+
+func (tr TypedPromptResponse[T]) IsOK() bool {
+	return tr.Error == ""
+}
+
 func (pr PromptResponse) IsOK() bool {
 	return pr.Error == ""
 }
@@ -111,3 +148,56 @@ type Agent interface {
 	// Close cleans up resources
 	Close() error
 }
+
+// ExecutePromptTyped is a generic helper that executes a prompt with type-safe structured output.
+// It automatically cleans up the JSON response before unmarshaling into the target type.
+func ExecutePromptTyped[T any](ctx context.Context, agent Agent, request TypedPromptRequest[T]) (*TypedPromptResponse[T], error) {
+	// Convert to regular PromptRequest
+	var schema T
+	promptReq := PromptRequest{
+		Context:          request.Context,
+		Name:             request.Name,
+		Prompt:           request.Prompt,
+		StructuredOutput: &schema,
+	}
+
+	// Execute the request
+	resp, err := agent.ExecutePrompt(ctx, promptReq)
+	if err != nil {
+		return &TypedPromptResponse[T]{
+			Request: request,
+			Error:   err.Error(),
+		}, err
+	}
+
+	// Convert to typed response
+	typedResp := &TypedPromptResponse[T]{
+		Request:       request,
+		Costs:         resp.Costs,
+		Model:         resp.Model,
+		Duration:      resp.Duration,
+		DurationModel: resp.DurationModel,
+		CacheHit:      resp.CacheHit,
+		Error:         resp.Error,
+	}
+
+	// If there's structured data, convert it
+	if resp.StructuredData != nil {
+		if data, ok := resp.StructuredData.(*T); ok {
+			typedResp.Data = *data
+		} else {
+			// Try to marshal and unmarshal through JSON
+			jsonData, marshalErr := json.Marshal(resp.StructuredData)
+			if marshalErr != nil {
+				typedResp.Error = fmt.Sprintf("failed to marshal structured data: %v", marshalErr)
+				return typedResp, marshalErr
+			}
+			if unmarshalErr := json.Unmarshal(jsonData, &typedResp.Data); unmarshalErr != nil {
+				typedResp.Error = fmt.Sprintf("failed to unmarshal structured data to target type: %v", unmarshalErr)
+				return typedResp, unmarshalErr
+			}
+		}
+	}
+
+	return typedResp, nil
+}

ai/claude_agent.go

@@ -10,6 +10,7 @@ import (
 
 	flanksourcecontext "github.com/flanksource/commons/context"
 	"github.com/flanksource/commons/logger"
+	"github.com/samber/lo"
 
 	"github.com/flanksource/clicky"
 	"github.com/flanksource/clicky/ai/cache"
@@ -134,7 +135,6 @@ func (ca *ClaudeAgent) ExecutePrompt(ctx context.Context, request PromptRequest)
 
 	task := clicky.StartTask(request.Name,
 		func(ctx flanksourcecontext.Context, t *clicky.Task) (interface{}, error) {
-			t.Infof("Starting Claude request")
 
 			resp, execErr := ca.executeClaude(ctx, request, t)
 			if execErr != nil {
@@ -184,7 +184,6 @@ func (ca *ClaudeAgent) ExecuteBatch(ctx context.Context, requests []PromptReques
 
 		task := clicky.StartTask(req.Name,
 			func(ctx flanksourcecontext.Context, t *clicky.Task) (interface{}, error) {
-				t.Infof("Processing request")
 
 				response, err := ca.executeClaude(t.Context(), req, t)
 
@@ -299,7 +298,7 @@ func (ca *ClaudeAgent) executeClaude(ctx context.Context, request PromptRequest,
 	}
 
 	if logger.V(3).Enabled() {
-		task.Tracef("%s", prompt)
+		task.Tracef("%s", lo.Ellipsis(prompt, 200))
 	}
 
 	startTime := time.Now()

ai/cost.go

@@ -17,6 +17,10 @@ type Cost struct {
 	OutputCost   float64 `json:"output_cost"`
 }
 
+func (c Cost) IsEmpty() bool {
+	return c.InputTokens == 0 && c.OutputTokens == 0 && c.InputCost == 0 && c.OutputCost == 0
+}
+
 func (c Cost) Total() float64 {
 	return c.InputCost + c.OutputCost
 }
@@ -31,11 +35,38 @@ func (c Cost) Pretty() api.Text {
 	}
 
 	if c.InputTokens+c.OutputTokens > 0 {
-		t = t.Space().Append(fmt.Sprintf("(%d in, %d out)", api.Human(c.InputTokens), api.Human(c.OutputTokens)), "text-muted")
+		t = t.Space().Append(fmt.Sprintf("(%v in, %v out)", api.Human(c.InputTokens), api.Human(c.OutputTokens)), "text-muted")
 	}
 	return t
 }
 
+func (c Costs) AggregateByModel() Costs {
+	modelCosts := make(map[string]Cost)
+	for _, cost := range c {
+		if cost.IsEmpty() {
+			continue
+		}
+		model := cost.Model
+		if model == "" {
+			model = "unknown"
+		}
+
+		existing, ok := modelCosts[model]
+		if !ok {
+			existing = Cost{Model: model}
+		}
+
+		modelCosts[model] = existing.Add(cost)
+	}
+
+	aggregated := Costs{}
+	for _, cost := range modelCosts {
+		aggregated = append(aggregated, cost)
+	}
+
+	return aggregated
+}
+
 func (c Costs) Sum() Cost {
 	total := Cost{}
 	for _, cost := range c {

ai/json_cleanup.go

@@ -0,0 +1,79 @@
+package ai
+
+import (
+	"encoding/json"
+	"regexp"
+	"strings"
+)
+
+var (
+	// Match markdown code blocks with optional language specifier
+	codeBlockRegex = regexp.MustCompile("(?s)```(?:json)?\\s*(.+?)```")
+
+	// Match JSON objects or arrays
+	jsonObjectRegex = regexp.MustCompile(`(?s)\{.*\}`)
+	jsonArrayRegex  = regexp.MustCompile(`(?s)\[.*\]`)
+)
+
+// CleanupJSONResponse attempts to extract and clean JSON from LLM responses
+// that may contain markdown formatting, explanatory text, or other noise.
+//
+// It tries the following strategies in order:
+// 1. Extract JSON from markdown code blocks (```json or ```)
+// 2. Extract the first JSON object {...}
+// 3. Extract the first JSON array [...]
+// 4. Return the trimmed original string
+//
+// After extraction, it validates that the result is valid JSON.
+func CleanupJSONResponse(response string) string {
+	// Trim whitespace
+	response = strings.TrimSpace(response)
+
+	// If it's already valid JSON, return as-is
+	if isValidJSON(response) {
+		return response
+	}
+
+	// Strategy 1: Extract from markdown code blocks
+	if matches := codeBlockRegex.FindStringSubmatch(response); len(matches) > 1 {
+		extracted := strings.TrimSpace(matches[1])
+		if isValidJSON(extracted) {
+			return extracted
+		}
+	}
+
+	// Strategy 2: Extract first JSON object
+	if match := jsonObjectRegex.FindString(response); match != "" {
+		if isValidJSON(match) {
+			return match
+		}
+	}
+
+	// Strategy 3: Extract first JSON array
+	if match := jsonArrayRegex.FindString(response); match != "" {
+		if isValidJSON(match) {
+			return match
+		}
+	}
+
+	// Strategy 4: Return original trimmed
+	return response
+}
+
+// isValidJSON checks if a string is valid JSON
+func isValidJSON(s string) bool {
+	var js interface{}
+	return json.Unmarshal([]byte(s), &js) == nil
+}
+
+// UnmarshalWithCleanup attempts to unmarshal JSON with automatic cleanup
+func UnmarshalWithCleanup(data string, v interface{}) error {
+	// First try direct unmarshal
+	if err := json.Unmarshal([]byte(data), v); err == nil {
+		return nil
+	}
+
+	// If that fails, try with cleanup
+	cleaned := CleanupJSONResponse(data)
+	return json.Unmarshal([]byte(cleaned), v)
+}

api/filter.go

@@ -201,7 +201,7 @@ func rowToCELMap(row PrettyDataRow) map[string]interface{} {
 		// - bool for booleans
 		// - string for strings
 		// - time.Time for dates
-		result[key] = fieldValue.Primitive()
+		result[key] = fieldValue.String()
 	}
 	return result
 }
@@ -256,7 +256,7 @@ func getVariableDeclarationsFromRow(row PrettyDataRow) []cel.EnvOption {
 	var decls []cel.EnvOption
 
 	for key, fieldValue := range row {
-		celType := inferCELTypeFromValue(fieldValue.Primitive())
+		celType := inferCELTypeFromValue(fieldValue.String())
 		decls = append(decls, cel.Variable(key, celType))
 	}