git.giacomolli.uk - onechat/blob - chain.go

   1 package main
   2
   3 import (
   4         "context"
   5         "encoding/json"
   6         "errors"
   7         "fmt"
   8         "strings"
   9 )
  10
  11 // errChainLogWrite marks an error of a log write, which ends the chain.
  12 var errChainLogWrite = errors.New("log write failed")
  13
  14 // ChainClient sends model calls.
  15 type ChainClient interface {
  16         Create(ctx context.Context, request ModelRequest) (ModelResponse, error)
  17 }
  18
  19 // Chain is one request chain of a turn or of a subagent. It sends request with
  20 // the input of the first request, runs the tool calls of each response, and
  21 // sends their results back, until a response has no tool call: the final
  22 // answer. The replies, the tool calls, and the results go to the log of tools
  23 // as they happen. The thoughts never do.
  24 //
  25 // Each later request begins with the whole input of the request before it,
  26 // then its output items, then the tool results. The output items hold the
  27 // reasoning with its encrypted content, because the calls do not store it. So
  28 // each request begins with all of the previous one and reads it from the prompt
  29 // cache. The messages of take join each later request after the tool results.
  30 type Chain struct {
  31         client  ChainClient
  32         tools   *Toolbox
  33         request ModelRequest
  34         take    func() []Message
  35 }
  36
  37 // Run runs the chain from input, the input of its first request, and returns
  38 // the text of the final answer. An error of a log write wraps
  39 // errChainLogWrite.
  40 func (c *Chain) Run(ctx context.Context, input []any) (string, error) {
  41         request := c.request
  42         for request.Turn = "first"; ; request.Turn = "next" {
  43                 request.Input = input
  44                 response, err := c.client.Create(ctx, request)
  45                 if err != nil {
  46                         return "", err
  47                 }
  48
  49                 var text string
  50                 input, text, err = c.handle(ctx, input, response)
  51                 if err != nil {
  52                         return "", err
  53                 }
  54
  55                 if len(input) == 0 {
  56                         return text, nil
  57                 }
  58         }
  59 }
  60
  61 // handle logs the replies of response and runs its tool calls. It returns the
  62 // input of the next request: input, the output items, the tool results, and
  63 // the messages of take. After the final answer it returns no input, and the
  64 // replies of the answer as its text.
  65 func (c *Chain) handle(ctx context.Context, input []any, response ModelResponse) ([]any, string, error) {
  66         var outputs []any
  67         var replies []string
  68         for _, raw := range response.Output {
  69                 item, err := chainInputItem(raw)
  70                 if err != nil {
  71                         return nil, "", err
  72                 }
  73
  74                 input = append(input, item)
  75
  76                 var call ChainFunctionCall
  77                 err = json.Unmarshal(raw, &call)
  78                 if err != nil {
  79                         return nil, "", fmt.Errorf("decode output item: %w", err)
  80                 }
  81
  82                 switch call.Type {
  83                 case "message":
  84                         text, err := modelOutputText([]json.RawMessage{raw})
  85                         if err != nil {
  86                                 return nil, "", err
  87                         }
  88
  89                         text = strings.TrimSpace(text)
  90                         if text == "" {
  91                                 continue
  92                         }
  93
  94                         _, err = c.tools.log.Append("reply", text)
  95                         if err != nil {
  96                                 return nil, "", fmt.Errorf("%w: %w", errChainLogWrite, err)
  97                         }
  98
  99                         replies = append(replies, text)
 100                 case "function_call":
 101                         output, err := c.tools.Call(ctx, call.Name, call.Arguments)
 102                         if err != nil {
 103                                 return nil, "", fmt.Errorf("%w: %w", errChainLogWrite, err)
 104                         }
 105
 106                         outputs = append(outputs, ChainFunctionCallOutput{
 107                                 Type:   "function_call_output",
 108                                 CallID: call.CallID,
 109                                 Output: output,
 110                         })
 111                 }
 112         }
 113
 114         if len(outputs) == 0 {
 115                 return nil, strings.Join(replies, "\n\n"), nil
 116         }
 117
 118         input = append(input, outputs...)
 119         for _, message := range c.take() {
 120                 input = append(input, modelUserMessage(message.Text))
 121         }
 122
 123         return input, "", nil
 124 }
 125
 126 // ChainFunctionCall holds the fields of a function call output item.
 127 type ChainFunctionCall struct {
 128         Type      string `json:"type"`
 129         Name      string `json:"name"`
 130         Arguments string `json:"arguments"`
 131         CallID    string `json:"call_id"`
 132 }
 133
 134 // ChainFunctionCallOutput is the input item with the result of a function
 135 // call.
 136 type ChainFunctionCallOutput struct {
 137         Type   string `json:"type"`
 138         CallID string `json:"call_id"`
 139         Output string `json:"output"`
 140 }
 141
 142 // chainInputItem returns an output item as an input item for the next request.
 143 // The calls do not store their items, so the item goes back whole, without its
 144 // id.
 145 func chainInputItem(raw json.RawMessage) (map[string]any, error) {
 146         var item map[string]any
 147         err := json.Unmarshal(raw, &item)
 148         if err != nil {
 149                 return nil, fmt.Errorf("decode output item: %w", err)
 150         }
 151
 152         delete(item, "id")
 153
 154         return item, nil
 155 }