summaryrefslogtreecommitdiff
path: root/json.go
blob: 818bbf94e555a4365f51935bf90ccdad456fcc12 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
package main

import (
	"encoding/json"
	"os"

	"go.wit.com/lib/protobuf/chatpb"
	"go.wit.com/log"
)

// GeminiRequest matches the overall structure of the gemini-cli JSON output.
type GeminiRequest struct {
	Model    string    `json:"model"`
	Contents []Content `json:"contents"`
	// Config is left as a raw message because its structure is complex and not needed for now.
	Config json.RawMessage `json:"config"`
}

// Content matches the 'contents' array elements.
type Content struct {
	Role  string `json:"role"`
	Parts []Part `json:"parts"`
}

// Part matches the 'parts' array elements.
// It can contain one of several types of data.
type Part struct {
	Text             string            `json:"text,omitempty"`
	ThoughtSignature string            `json:"thoughtSignature,omitempty"`
	FunctionCall     *FunctionCall     `json:"functionCall,omitempty"`
	FunctionResponse *FunctionResponse `json:"functionResponse,omitempty"`
}

// FunctionCall matches the 'functionCall' object.
type FunctionCall struct {
	Name string            `json:"name"`
	Args map[string]string `json:"args"`
}

// FunctionResponse matches the 'functionResponse' object.
type FunctionResponse struct {
	ID       string                 `json:"id"`
	Name     string                 `json:"name"`
	Response map[string]interface{} `json:"response"`
}

func parsePB(filename string) (*chatpb.GeminiRequest, error) {
	// Read the entire file
	data, err := os.ReadFile(filename)
	if err != nil {
		return nil, log.Errorf("failed to read file %s: %w", filename, err)
	}
	pb := new(chatpb.GeminiRequest)
	if err := pb.UnmarshalJSON(data); err != nil {
		return nil, err
	}
	return pb, nil
}

// parseJSON opens the given file, reads it, and unmarshals it into our structs.
func parseJSON(filename string) (*GeminiRequest, error) {
	log.Infof("Attempting to parse file: %s\n", filename)

	// Read the entire file
	data, err := os.ReadFile(filename)
	if err != nil {
		return nil, log.Errorf("failed to read file %s: %w", filename, err)
	}

	// Unmarshal the JSON data
	var req *GeminiRequest
	req = new(GeminiRequest)
	if err := json.Unmarshal(data, &req); err != nil {
		return nil, log.Errorf("failed to unmarshal JSON from %s: %w", filename, err)
	}

	dumpSummaryJSON(req)
	return req, nil
}

func dumpSummaryJSON(req *GeminiRequest) {
	var totalFC, totalTexts, totalFR int

	// Log the parsed data to confirm it worked

	// Example of accessing deeper data
	for _, content := range req.Contents {
		// log.Infof("Content[%d] Role: %s", i, content.Role)
		for _, part := range content.Parts {
			if part.Text != "" {
				// log.Infof("  Part[%d] Text: %.60s...", j, part.Text) // Print snippet
				totalTexts += 1
			}
			if part.FunctionCall != nil {
				// log.Infof("  Part[%d] FunctionCall: %s", j, part.FunctionCall.Name)
				totalFC += 1
			}
			if part.FunctionResponse != nil {
				// log.Infof("  Part[%d] FunctionCall: %s", j, part.FunctionCall.Name)
				totalFR += 1
			}
		}
	}
	log.Printf("Parsed JSON (Model: %s) (# of content blocks %d) (Text #=%d) (FC=%d) (FR=%d)\n", req.Model, len(req.Contents), totalTexts, totalFC, totalFR)
}

func dumpFullJSON(req *GeminiRequest) {
	// Log the parsed data to confirm it worked
	log.Info("Successfully parsed JSON file.")
	log.Infof("Model: %s", req.Model)
	log.Infof("Number of content blocks: %d", len(req.Contents))

	// Example of accessing deeper data
	for i, content := range req.Contents {
		log.Infof("Content[%d] Role: %s", i, content.Role)
		for j, part := range content.Parts {
			if part.Text != "" {
				log.Infof("  Part[%d] Text: %.60s...", j, part.Text) // Print snippet
			}
			if part.FunctionCall != nil {
				log.Infof("  Part[%d] FunctionCall: %s", j, part.FunctionCall.Name)
			}
		}
	}
}