grail/gf-lt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Enha: save/load message token stats

Browse Source
This commit is contained in:
Grail Finder
2026-02-27 11:23:03 +03:00
parent 915b029d2c
commit c855c30ae2
4 changed files with 30 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
models/models.go
View File

@@ -115,33 +115,39 @@ type RoleMsg struct {
}
// MarshalJSON implements custom JSON marshaling for RoleMsg
func (m *RoleMsg) MarshalJSON() ([]byte, error) {
//
//nolint:gocritic
func (m RoleMsg) MarshalJSON() ([]byte, error) {
if m.hasContentParts {
// Use structured content format
aux := struct {
Role string `json:"role"`
Content []any `json:"content"`
ToolCallID string `json:"tool_call_id,omitempty"`
KnownTo []string `json:"known_to,omitempty"`
Role string `json:"role"`
Content []any `json:"content"`
ToolCallID string `json:"tool_call_id,omitempty"`
KnownTo []string `json:"known_to,omitempty"`
Stats *ResponseStats `json:"stats,omitempty"`
}{
Role: m.Role,
Content: m.ContentParts,
ToolCallID: m.ToolCallID,
KnownTo: m.KnownTo,
Stats: m.Stats,
}
return json.Marshal(aux)
} else {
// Use simple content format
aux := struct {
Role string `json:"role"`
Content string `json:"content"`
ToolCallID string `json:"tool_call_id,omitempty"`
KnownTo []string `json:"known_to,omitempty"`
Role string `json:"role"`
Content string `json:"content"`
ToolCallID string `json:"tool_call_id,omitempty"`
KnownTo []string `json:"known_to,omitempty"`
Stats *ResponseStats `json:"stats,omitempty"`
}{
Role: m.Role,
Content: m.Content,
ToolCallID: m.ToolCallID,
KnownTo: m.KnownTo,
Stats: m.Stats,
}
return json.Marshal(aux)
}
@@ -151,26 +157,29 @@ func (m *RoleMsg) MarshalJSON() ([]byte, error) {
func (m *RoleMsg) UnmarshalJSON(data []byte) error {
// First, try to unmarshal as structured content format
var structured struct {
Role string `json:"role"`
Content []any `json:"content"`
ToolCallID string `json:"tool_call_id,omitempty"`
KnownTo []string `json:"known_to,omitempty"`
Role string `json:"role"`
Content []any `json:"content"`
ToolCallID string `json:"tool_call_id,omitempty"`
KnownTo []string `json:"known_to,omitempty"`
Stats *ResponseStats `json:"stats,omitempty"`
}
if err := json.Unmarshal(data, &structured); err == nil && len(structured.Content) > 0 {
m.Role = structured.Role
m.ContentParts = structured.Content
m.ToolCallID = structured.ToolCallID
m.KnownTo = structured.KnownTo
m.Stats = structured.Stats
m.hasContentParts = true
return nil
}
// Otherwise, unmarshal as simple content format
var simple struct {
Role string `json:"role"`
Content string `json:"content"`
ToolCallID string `json:"tool_call_id,omitempty"`
KnownTo []string `json:"known_to,omitempty"`
Role string `json:"role"`
Content string `json:"content"`
ToolCallID string `json:"tool_call_id,omitempty"`
KnownTo []string `json:"known_to,omitempty"`
Stats *ResponseStats `json:"stats,omitempty"`
}
if err := json.Unmarshal(data, &simple); err != nil {
return err
@@ -179,6 +188,7 @@ func (m *RoleMsg) UnmarshalJSON(data []byte) error {
m.Content = simple.Content
m.ToolCallID = simple.ToolCallID
m.KnownTo = simple.KnownTo
m.Stats = simple.Stats
m.hasContentParts = false
return nil
}