More cleanup

This commit is contained in:
ParthSareen 2024-12-11 18:11:00 -08:00
parent 97abd7bfea
commit 2536ffe0ab

View File

@ -61,21 +61,6 @@ type Usage struct {
TotalTokens int `json:"total_tokens"` TotalTokens int `json:"total_tokens"`
} }
// ChunkUsage is an alias for Usage with the ability to marshal a marker
// value as null. This is to allow omitting the field in chunks when usage
// isn't requested, and otherwise return null on non-final chunks when it
// is requested to follow OpenAI's behavior.
type ChunkUsage = Usage
// var nullChunkUsage = ChunkUsage{}
// func (u *ChunkUsage) MarshalJSON() ([]byte, error) {
// if u == &nullChunkUsage {
// return []byte("null"), nil
// }
// return json.Marshal(*u)
// }
type ResponseFormat struct { type ResponseFormat struct {
Type string `json:"type"` Type string `json:"type"`
JsonSchema *JsonSchema `json:"json_schema,omitempty"` JsonSchema *JsonSchema `json:"json_schema,omitempty"`
@ -127,7 +112,7 @@ type ChatCompletionChunk struct {
Model string `json:"model"` Model string `json:"model"`
SystemFingerprint string `json:"system_fingerprint"` SystemFingerprint string `json:"system_fingerprint"`
Choices []ChunkChoice `json:"choices"` Choices []ChunkChoice `json:"choices"`
Usage *ChunkUsage `json:"usage,omitempty"` Usage *Usage `json:"usage,omitempty"`
} }
// TODO (https://github.com/ollama/ollama/issues/5259): support []string, []int and [][]int // TODO (https://github.com/ollama/ollama/issues/5259): support []string, []int and [][]int
@ -163,7 +148,7 @@ type CompletionChunk struct {
Choices []CompleteChunkChoice `json:"choices"` Choices []CompleteChunkChoice `json:"choices"`
Model string `json:"model"` Model string `json:"model"`
SystemFingerprint string `json:"system_fingerprint"` SystemFingerprint string `json:"system_fingerprint"`
Usage *ChunkUsage `json:"usage,omitempty"` Usage *Usage `json:"usage,omitempty"`
} }
type ToolCall struct { type ToolCall struct {
@ -655,7 +640,7 @@ func (w *ChatWriter) writeResponse(data []byte) (int, error) {
if w.stream { if w.stream {
c := toChunk(w.id, chatResponse) c := toChunk(w.id, chatResponse)
if w.streamOptions != nil && w.streamOptions.IncludeUsage { if w.streamOptions != nil && w.streamOptions.IncludeUsage {
c.Usage = &ChunkUsage{} c.Usage = &Usage{}
} }
d, err := json.Marshal(c) d, err := json.Marshal(c)
if err != nil { if err != nil {
@ -719,7 +704,7 @@ func (w *CompleteWriter) writeResponse(data []byte) (int, error) {
if w.stream { if w.stream {
c := toCompleteChunk(w.id, generateResponse) c := toCompleteChunk(w.id, generateResponse)
if w.streamOptions != nil && w.streamOptions.IncludeUsage { if w.streamOptions != nil && w.streamOptions.IncludeUsage {
c.Usage = &ChunkUsage{} c.Usage = &Usage{}
} }
d, err := json.Marshal(c) d, err := json.Marshal(c)
if err != nil { if err != nil {