X-Git-Url: http://git.giacomolli.uk/onechat/blobdiff_plain/5da3ea038d1e9f9a2634d5ab137aef6989a3a789..e0ab6220b6fd35f7e92a04e4bf149940185713ab:/responses.go

diff --git a/responses.go b/responses.go
index 3055e70..19ab6a1 100644
--- a/responses.go
+++ b/responses.go
@@ -73,8 +73,10 @@ type Client struct {
 
 // Request is one model call. Input holds the input items in order, such as
 // userMessage values or output items of an earlier response. Node names the
-// node that a compaction builds, or is nil. It goes only to the usage file,
-// never to the endpoint.
+// node that a compaction builds, or is nil. Turn names the request of a turn:
+// first, next, or warmup. Node and Turn go only to the usage file, never to the
+// endpoint. A call with StopAtOutput ends at the first output item, after the
+// endpoint has read the whole input, and gives no output and no usage.
 type Request struct {
 	Model        string
 	Effort       string
@@ -82,6 +84,8 @@ type Request struct {
 	Tools        []Tool
 	Input        []any
 	Node         *nodeKey
+	Turn         string
+	StopAtOutput bool
 }
 
 // Tool is the definition of a function tool.
@@ -258,7 +262,8 @@ func (c *Client) Create(ctx context.Context, request Request) (Response, error)
 		record.Node = &[2]int{request.Node.level, request.Node.index}
 	}
 
-	response, err := c.send(httpRequest, &record)
+	record.Turn = request.Turn
+	response, err := c.send(httpRequest, &record, request.StopAtOutput)
 	if err != nil {
 		record.Error = err.Error()
 	}
@@ -274,7 +279,7 @@ func (c *Client) Create(ctx context.Context, request Request) (Response, error)
 
 // send makes the HTTP call and reads its stream. It puts the status and the
 // quota headers into record.
-func (c *Client) send(httpRequest *http.Request, record *usageRecord) (response Response, err error) {
+func (c *Client) send(httpRequest *http.Request, record *usageRecord, stopAtOutput bool) (response Response, err error) {
 	httpResponse, err := c.http.Do(httpRequest)
 	if err != nil {
 		return Response{}, fmt.Errorf("send model call: %w", err)
@@ -293,7 +298,7 @@ func (c *Client) send(httpRequest *http.Request, record *usageRecord) (response
 		return Response{}, errors.Join(err, readErr)
 	}
 
-	return readStream(httpResponse.Body)
+	return readStream(httpResponse.Body, stopAtOutput)
 }
 
 // usageLog appends one JSON line for each model call to usage.jsonl. Calls can
@@ -306,13 +311,14 @@ type usageLog struct {
 // usageRecord is one line of usage.jsonl. Date is the start of the call, and
 // Seconds its length. Node is the node [l, i] of a compaction. Status is the
 // HTTP status, or 0 when no response came. TextBytes is the size of the text of
-// the reply. Quota holds the x-codex-* headers of the response, without the
+// the reply. Turn names the request of a turn. Quota holds the x-codex-* headers of the response, without the
 // turn state.
 type usageRecord struct {
 	Date         string            `json:"date"`
 	Seconds      float64           `json:"seconds"`
 	Model        string            `json:"model"`
 	Node         *[2]int           `json:"node,omitempty"`
+	Turn         string            `json:"turn,omitempty"`
 	Status       int               `json:"status"`
 	InputTokens  int               `json:"input_tokens"`
 	CachedTokens int               `json:"cached_tokens"`
@@ -400,8 +406,9 @@ func quotaHeaders(header http.Header) map[string]string {
 
 // readStream reads the server-sent events of a call until the response
 // completes. The output items come from the output_item.done events, or from
-// the completed response when the stream sent none.
-func readStream(stream io.Reader) (Response, error) {
+// the completed response when the stream sent none. With stopAtOutput it ends
+// at the first output item.
+func readStream(stream io.Reader, stopAtOutput bool) (Response, error) {
 	var response Response
 	reader := bufio.NewReader(stream)
 	var data bytes.Buffer
@@ -419,7 +426,7 @@ func readStream(stream io.Reader) (Response, error) {
 				data.WriteByte('\n')
 			}
 		} else if data.Len() > 0 {
-			done, eventErr := readEvent(&response, data.Bytes())
+			done, eventErr := readEvent(&response, data.Bytes(), stopAtOutput)
 			data.Reset()
 			if eventErr != nil || done {
 				return response, eventErr
@@ -434,7 +441,7 @@ func readStream(stream io.Reader) (Response, error) {
 
 // readEvent adds one stream event to response, and tells if it ends the
 // stream.
-func readEvent(response *Response, data []byte) (bool, error) {
+func readEvent(response *Response, data []byte, stopAtOutput bool) (bool, error) {
 	var event streamEvent
 	err := json.Unmarshal(data, &event)
 	if err != nil {
@@ -442,6 +449,8 @@ func readEvent(response *Response, data []byte) (bool, error) {
 	}
 
 	switch event.Type {
+	case "response.output_item.added":
+		return stopAtOutput, nil
 	case "response.output_item.done":
 		response.Output = append(response.Output, event.Item)
 	case "response.completed":