chore: refactoring

2 files changed, 10 insertions, 6 deletions

diff --git a/internal/message/content.go b/internal/message/content.go
index d7b76f8ce..0e1608753 100644
--- a/internal/message/content.go
+++ b/internal/message/content.go
@@ -109,6 +109,11 @@ type Finish struct {
 	Time   time.Time    `json:"time"`
 }
 
+type DBFinish struct {
+	Reason FinishReason `json:"reason"`
+	Time   int64        `json:"time"`
+}
+
 func (Finish) isPart() {}
 
 func (m *Message) Content() *TextContent {
diff --git a/internal/message/message.go b/internal/message/message.go
index 197c2f0f5..35a414533 100644
--- a/internal/message/message.go
+++ b/internal/message/message.go
@@ -387,7 +387,10 @@ func marshallParts(parts []ContentPart) ([]byte, error) {
 			dataBytes, err = json.Marshal(p)
 		case Finish:
 			typ = finishType
-			dataBytes, err = json.Marshal(p)
+			var dbFinish DBFinish
+			dbFinish.Reason = p.Reason
+			dbFinish.Time = p.Time.UnixMilli()
+			dataBytes, err = json.Marshal(dbFinish)
 		default:
 			return nil, fmt.Errorf("unknown part type for marshalling: %T", part)
 		}
@@ -464,11 +467,7 @@ func unmarshallParts(data []byte) ([]ContentPart, error) {
 			}
 			parts = append(parts, p)
 		case finishType:
-			type dbFinish struct {
-				Reason FinishReason `json:"reason"`
-				Time   int64        `json:"time"`
-			}
-			var p dbFinish
+			var p DBFinish
 			if err := json.Unmarshal(wrapper.Data, &p); err != nil {
 				return nil, fmt.Errorf("unmarshal Finish: %w. Data: %s", err, string(wrapper.Data))
 			}