1package typesafe23import (4 "bytes"5 "encoding/json/jsontext"6 jsonv2 "encoding/json/v2"7 "errors"8 "fmt"9 "math"10)1112// ErrAnswerNotFound indicates that a named answer is absent.13var ErrAnswerNotFound = errors.New("answer not found")1415// ErrAnswerType indicates that an answer exists but has another type.16var ErrAnswerType = errors.New("answer has unexpected type")1718// Answer is a discriminated answer returned by System One.19type Answer interface {20 answerType() string21 jsonv2.Marshaler22}2324// NoulAnswer contains the probability of a yes answer.25type NoulAnswer struct {26 Noul float64 `json:"noul"`27}2829func (NoulAnswer) answerType() string { return "noul" }3031// ChoiceAnswer contains the selected label, confidence, and label probabilities.32type ChoiceAnswer struct {33 Choice string `json:"choice"`34 Confidence float64 `json:"confidence"`35 Probabilities map[string]float64 `json:"probabilities"`36}3738func (ChoiceAnswer) answerType() string { return "choice" }3940// ScoreAnswer contains an expected score, rubric legend, confidence, and probabilities.41type ScoreAnswer struct {42 Score float64 `json:"score"`43 Confidence float64 `json:"confidence"`44 Legend map[string]any `json:"legend"`45 Probabilities map[string]float64 `json:"probabilities"`46}4748func (ScoreAnswer) answerType() string { return "score" }4950// UnknownAnswer preserves an answer type this SDK does not know yet.51type UnknownAnswer struct {52 Type string53 Raw jsontext.Value54}5556func (a UnknownAnswer) answerType() string { return a.Type }57func (a UnknownAnswer) MarshalJSON() ([]byte, error) { return append([]byte(nil), a.Raw...), nil }5859func marshalAnswer(kind string, value any) ([]byte, error) {60 payload, err := jsonv2.Marshal(value)61 if err != nil {62 return nil, err63 }64 object, err := rawObject(payload)65 if err != nil {66 return nil, err67 }68 object["type"] = jsontext.Value(fmt.Sprintf("%q", kind))69 return jsonv2.Marshal(object)70}7172// MarshalJSON encodes a NoulAnswer with its discriminator.73func (a NoulAnswer) MarshalJSON() ([]byte, error) {74 return marshalAnswer(a.answerType(), struct {75 Noul float64 `json:"noul"`76 }{a.Noul})77}7879// MarshalJSON encodes a ChoiceAnswer with its discriminator.80func (a ChoiceAnswer) MarshalJSON() ([]byte, error) {81 return marshalAnswer(a.answerType(), struct {82 Choice string `json:"choice"`83 Confidence float64 `json:"confidence"`84 Probabilities map[string]float64 `json:"probabilities"`85 }{a.Choice, a.Confidence, a.Probabilities})86}8788// MarshalJSON encodes a ScoreAnswer with its discriminator.89func (a ScoreAnswer) MarshalJSON() ([]byte, error) {90 return marshalAnswer(a.answerType(), struct {91 Score float64 `json:"score"`92 Confidence float64 `json:"confidence"`93 Legend map[string]any `json:"legend"`94 Probabilities map[string]float64 `json:"probabilities"`95 }{a.Score, a.Confidence, a.Legend, a.Probabilities})96}9798// Usage reports token counts for a System One request.99type Usage struct {100 InputTokens int `json:"input_tokens"`101 OutputTokens int `json:"output_tokens"`102}103104// SystemOneResponse contains all named answers and request usage metadata.105type SystemOneResponse struct {106 Model string `json:"model"`107 Answers map[string]Answer `json:"answers"`108 Usage Usage `json:"usage"`109 RequestID string `json:"-"`110}111112// UnmarshalAnswer decodes a discriminated answer and preserves unknown types.113func UnmarshalAnswer(data []byte) (Answer, error) {114 object, err := rawObject(data)115 if err != nil {116 return nil, err117 }118 tag, ok := object["type"]119 if !ok {120 return nil, fmt.Errorf("answer type is required")121 }122 var kind string123 if err := jsonv2.Unmarshal(tag, &kind); err != nil || kind == "" {124 return nil, fmt.Errorf("answer type must be a non-empty string")125 }126 switch kind {127 case "noul":128 var value NoulAnswer129 if err := decodeKnownAnswer(data, object, &value, kind); err != nil {130 return nil, err131 }132 return value, nil133 case "choice":134 var value ChoiceAnswer135 if err := decodeKnownAnswer(data, object, &value, kind); err != nil {136 return nil, err137 }138 return value, nil139 case "score":140 var value ScoreAnswer141 if err := decodeKnownAnswer(data, object, &value, kind); err != nil {142 return nil, err143 }144 return value, nil145 default:146 return UnknownAnswer{Type: kind, Raw: append(jsontext.Value(nil), data...)}, nil147 }148}149150func decodeKnownAnswer(data []byte, object map[string]jsontext.Value, target any, kind string) error {151 if err := jsonv2.Unmarshal(data, target); err != nil {152 return fmt.Errorf("decode %s answer: %w", kind, err)153 }154 switch value := target.(type) {155 case *NoulAnswer:156 if !hasRaw(object, "noul") || !finiteProbability(value.Noul) {157 return fmt.Errorf("invalid noul answer")158 }159 case *ChoiceAnswer:160 if !validChoiceAnswer(*value, object) {161 return fmt.Errorf("invalid choice answer")162 }163 case *ScoreAnswer:164 if !validScoreAnswer(*value, object) {165 return fmt.Errorf("invalid score answer")166 }167 }168 return nil169}170171func validChoiceAnswer(value ChoiceAnswer, object map[string]jsontext.Value) bool {172 return hasRaw(object, "choice") && hasRaw(object, "confidence") &&173 finiteProbability(value.Confidence) && value.Choice != "" &&174 mapHasKey(value.Probabilities, value.Choice) &&175 validProbabilities(value.Probabilities, object, "probabilities")176}177178func validScoreAnswer(value ScoreAnswer, object map[string]jsontext.Value) bool {179 return hasRaw(object, "score") && hasRaw(object, "confidence") &&180 hasRaw(object, "legend") && finiteProbability(value.Confidence) &&181 sameKeys(value.Legend, value.Probabilities) &&182 validProbabilities(value.Probabilities, object, "probabilities")183}184185func finiteProbability(value float64) bool {186 return !math.IsNaN(value) && !math.IsInf(value, 0) && value >= 0 && value <= 1187}188189func validProbabilities(values map[string]float64, object map[string]jsontext.Value, key string) bool {190 raw, ok := object[key]191 if !ok || !hasRaw(object, key) || len(values) == 0 {192 return false193 }194 var rawValues map[string]jsontext.Value195 if err := jsonv2.Unmarshal(raw, &rawValues); err != nil || len(rawValues) != len(values) {196 return false197 }198 for _, rawValue := range rawValues {199 if bytes.Equal(bytes.TrimSpace(rawValue), []byte("null")) {200 return false201 }202 }203 for _, value := range values {204 if !finiteProbability(value) {205 return false206 }207 }208 return true209}210211func mapHasKey(values map[string]float64, key string) bool { _, ok := values[key]; return ok }212213func sameKeys(legend map[string]any, probabilities map[string]float64) bool {214 if len(legend) != len(probabilities) {215 return false216 }217 for key := range legend {218 if _, ok := probabilities[key]; !ok {219 return false220 }221 }222 return true223}224225// UnmarshalJSON decodes and validates a SystemOneResponse.226func (r *SystemOneResponse) UnmarshalJSON(data []byte) error {227 object, err := rawObject(data)228 if err != nil {229 return err230 }231 if !hasRaw(object, "model") || !hasRaw(object, "answers") || !hasRaw(object, "usage") {232 return fmt.Errorf("response requires model, answers, and usage")233 }234 var usageObject map[string]jsontext.Value235 if err := jsonv2.Unmarshal(object["usage"], &usageObject); err != nil || !hasRaw(usageObject, "input_tokens") || !hasRaw(usageObject, "output_tokens") {236 return fmt.Errorf("usage requires input_tokens and output_tokens")237 }238 var envelope struct {239 Model string `json:"model"`240 Answers map[string]jsontext.Value `json:"answers"`241 Usage *Usage `json:"usage"`242 }243 if err := jsonv2.Unmarshal(data, &envelope); err != nil {244 return err245 }246 if envelope.Model == "" || envelope.Answers == nil || envelope.Usage == nil {247 return fmt.Errorf("response has invalid required fields")248 }249 answers := make(map[string]Answer, len(envelope.Answers))250 for name, raw := range envelope.Answers {251 answer, err := UnmarshalAnswer(raw)252 if err != nil {253 return fmt.Errorf("answer %q: %w", name, err)254 }255 answers[name] = answer256 }257 *r = SystemOneResponse{Model: envelope.Model, Answers: answers, Usage: *envelope.Usage}258 return nil259}260261// Noul returns the named Noul answer.262func (r SystemOneResponse) Noul(name string) (NoulAnswer, error) {263 value, err := answerAs[NoulAnswer](r.Answers, name)264 return value, err265}266267// Choice returns the named Choice answer.268func (r SystemOneResponse) Choice(name string) (ChoiceAnswer, error) {269 value, err := answerAs[ChoiceAnswer](r.Answers, name)270 return value, err271}272273// Score returns the named Score answer.274func (r SystemOneResponse) Score(name string) (ScoreAnswer, error) {275 value, err := answerAs[ScoreAnswer](r.Answers, name)276 return value, err277}278279// Answer returns a named answer as the requested concrete answer type.280func (r SystemOneResponse) Answer[T Answer](name string) (T, error) {281 return answerAs[T](r.Answers, name)282}283284func answerAs[T Answer](answers map[string]Answer, name string) (T, error) {285 var zero T286 answer, ok := answers[name]287 if !ok {288 return zero, fmt.Errorf("%w: %s", ErrAnswerNotFound, name)289 }290 value, ok := answer.(T)291 if !ok {292 return zero, fmt.Errorf("%w: %s", ErrAnswerType, name)293 }294 return value, nil295}296297// ModelCard describes an available model.298type ModelCard struct {299 Name string `json:"name"`300 Description string `json:"description"`301 ReleaseDate string `json:"release_date"`302}303304// ListModelsResponse contains the models available to the account.305type ListModelsResponse struct {306 Models []ModelCard `json:"models"`307 RequestID string `json:"-"`308}309310// UnmarshalJSON decodes and validates a model list response.311func (r *ListModelsResponse) UnmarshalJSON(data []byte) error {312 object, err := rawObject(data)313 if err != nil {314 return err315 }316 if !hasRaw(object, "models") {317 return fmt.Errorf("response requires models")318 }319 var rawModels []jsontext.Value320 if err := jsonv2.Unmarshal(object["models"], &rawModels); err != nil {321 return err322 }323 if rawModels == nil {324 return fmt.Errorf("models must be an array")325 }326 models := make([]ModelCard, len(rawModels))327 for index, rawModel := range rawModels {328 modelObject, err := rawObject(rawModel)329 if err != nil {330 return fmt.Errorf("models[%d]: %w", index, err)331 }332 if !hasRaw(modelObject, "name") || !hasRaw(modelObject, "description") || !hasRaw(modelObject, "release_date") {333 return fmt.Errorf("models[%d]: missing required fields", index)334 }335 if err := jsonv2.Unmarshal(rawModel, &models[index]); err != nil {336 return err337 }338 if models[index].Name == "" {339 return fmt.Errorf("models[%d]: name must not be empty", index)340 }341 }342 *r = ListModelsResponse{Models: models}343 return nil344}345346var _ jsonv2.Marshaler = NoulAnswer{}347var _ jsonv2.Marshaler = ChoiceAnswer{}348var _ jsonv2.Marshaler = ScoreAnswer{}349var _ jsonv2.Marshaler = NoulQuestion{}350var _ jsonv2.Marshaler = ChoiceQuestion{}351var _ jsonv2.Marshaler = ScoreQuestion{}