From ec83831eb7bd9e2a3a1fc43a15b37407c7572873 Mon Sep 17 00:00:00 2001 From: Leo Date: Thu, 18 Jun 2026 12:24:46 +0200 Subject: [PATCH] Unwrap structured scrape JSON output --- internal/cmd/output_helpers.go | 77 ++++++++++++++++++++++------------ internal/cmd/sessions_test.go | 27 +++++++++++- 2 files changed, 75 insertions(+), 29 deletions(-) diff --git a/internal/cmd/output_helpers.go b/internal/cmd/output_helpers.go index 9e272a1..b361ded 100644 --- a/internal/cmd/output_helpers.go +++ b/internal/cmd/output_helpers.go @@ -80,12 +80,18 @@ func PrintListOrEmpty(items any, emptyMsg string) (bool, error) { } // PrintScrapeResponse formats scrape output consistently across all scrape commands. -// In JSON mode, returns the full response. In text mode without instructions, -// returns just the markdown. With instructions, checks data.success and returns -// the extracted data or an error message. +// In JSON mode without instructions, returns the full response. With instructions, +// returns the extracted structured data directly. In text mode without instructions, +// returns just the markdown. func PrintScrapeResponse(resp *api.DataSpace, hasInstructions bool) error { - // JSON mode: return full response if IsJSONOutput() { + if hasInstructions { + data, err := extractScrapeStructuredData(resp) + if err != nil { + return err + } + return GetFormatter().Print(data) + } return GetFormatter().Print(resp) } @@ -96,31 +102,21 @@ func PrintScrapeResponse(resp *api.DataSpace, hasInstructions bool) error { } // Structured mode: check data.success - if resp.Structured != nil { - // Check success field - if resp.Structured.Success != nil && !*resp.Structured.Success { - if resp.Structured.Error != nil { - return fmt.Errorf("%s", *resp.Structured.Error) - } - return fmt.Errorf("scrape failed") - } - // Return data if present - if resp.Structured.Data != nil { - // Extract actual data from the union type wrapper - if data, err := resp.Structured.Data.AsBaseModel(); err == nil && data != nil { - // Print with a nice header for scrape results - fmt.Println("Scraped content from the current page:") - fmt.Println() - // Print as indented JSON for better readability - jsonBytes, err := json.MarshalIndent(data, "", " ") - if err != nil { - return GetFormatter().Print(data) - } - fmt.Println(string(jsonBytes)) - return nil - } + data, err := extractScrapeStructuredData(resp) + if err == nil { + fmt.Println("Scraped content from the current page:") + fmt.Println() + jsonBytes, err := json.MarshalIndent(data, "", " ") + if err != nil { + return GetFormatter().Print(data) } + fmt.Println(string(jsonBytes)) + return nil + } + if resp != nil && resp.Structured != nil && resp.Structured.Success != nil && !*resp.Structured.Success { + return err } + fmt.Println("Scraped content from the current page:") fmt.Println() // Print as indented JSON for better readability @@ -132,6 +128,33 @@ func PrintScrapeResponse(resp *api.DataSpace, hasInstructions bool) error { return nil } +// extractScrapeStructuredData returns the direct payload from structured scrape +// responses, matching the SDK default for instruction-based scrapes. +func extractScrapeStructuredData(resp *api.DataSpace) (any, error) { + if resp == nil || resp.Structured == nil { + return nil, fmt.Errorf("scrape did not return structured data") + } + if resp.Structured.Success != nil && !*resp.Structured.Success { + if resp.Structured.Error != nil && *resp.Structured.Error != "" { + return nil, fmt.Errorf("%s", *resp.Structured.Error) + } + return nil, fmt.Errorf("scrape failed") + } + if resp.Structured.Data == nil { + return nil, fmt.Errorf("scrape did not return structured data") + } + + raw, err := json.Marshal(resp.Structured.Data) + if err != nil { + return nil, err + } + var data any + if err := json.Unmarshal(raw, &data); err != nil { + return nil, err + } + return data, nil +} + // printSessionStatus formats session status output with simplified Steps display. // In JSON mode, returns the full response. In text mode, formats Steps as a simple list. func printSessionStatus(resp *api.SessionResponse) error { diff --git a/internal/cmd/sessions_test.go b/internal/cmd/sessions_test.go index d009e72..b39a389 100644 --- a/internal/cmd/sessions_test.go +++ b/internal/cmd/sessions_test.go @@ -2,6 +2,7 @@ package cmd import ( "context" + "encoding/json" "fmt" "os" "path/filepath" @@ -391,7 +392,7 @@ func TestRunSessionExecute_InvalidJSON(t *testing.T) { func TestRunSessionScrape(t *testing.T) { server := setupSessionTest(t) - scrapeResp := fmt.Sprintf(`{"markdown":"hi","structured":{},"session":%s}`, sessionJSON()) + scrapeResp := fmt.Sprintf(`{"markdown":"hi","structured":{"data":{"result":"hi"},"success":true},"session":%s}`, sessionJSON()) server.AddResponse("/sessions/"+sessionIDTest+"/page/scrape", 200, scrapeResp) origInstructions := sessionScrapeInstructions @@ -457,7 +458,7 @@ func TestRunSessionScrape_Defaults(t *testing.T) { func TestRunSessionScrape_InstructionsOnly(t *testing.T) { server := setupSessionTest(t) - scrapeResp := fmt.Sprintf(`{"markdown":"content","structured":{"data":"value"},"session":%s}`, sessionJSON()) + scrapeResp := fmt.Sprintf(`{"markdown":"content","structured":{"data":{"title":"Extracted Title","count":2},"success":true},"session":%s}`, sessionJSON()) server.AddResponse("/sessions/"+sessionIDTest+"/page/scrape", 200, scrapeResp) origInstructions := sessionScrapeInstructions @@ -486,6 +487,20 @@ func TestRunSessionScrape_InstructionsOnly(t *testing.T) { if stdout == "" { t.Error("expected output, got empty string") } + + var parsed map[string]any + if err := json.Unmarshal([]byte(stdout), &parsed); err != nil { + t.Fatalf("expected JSON output, got %q: %v", stdout, err) + } + if _, ok := parsed["markdown"]; ok { + t.Fatalf("expected structured data only, got markdown key in %v", parsed) + } + if parsed["title"] != "Extracted Title" { + t.Fatalf("expected structured title, got %v", parsed) + } + if parsed["count"] != float64(2) { + t.Fatalf("expected structured count, got %v", parsed) + } } func TestRunSessionScrape_OnlyMainContentOnly(t *testing.T) { @@ -519,6 +534,14 @@ func TestRunSessionScrape_OnlyMainContentOnly(t *testing.T) { if stdout == "" { t.Error("expected output, got empty string") } + + var parsed map[string]any + if err := json.Unmarshal([]byte(stdout), &parsed); err != nil { + t.Fatalf("expected JSON output, got %q: %v", stdout, err) + } + if parsed["markdown"] != "main content only" { + t.Fatalf("expected full scrape response with markdown, got %v", parsed) + } } func TestRunSessionScrape_TextOutput(t *testing.T) {