180 lines · 6.5 KB · Go
RawDownload
package executor import (	"bufio"	"bytes"	"context"	"io"	"net/http"	"strings" 	"github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps"	cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"	cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"	log "github.com/sirupsen/logrus"	"github.com/tidwall/gjson") func (e *XAIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (_ *cliproxyexecutor.StreamResult, err error) {	ctx = helps.EnsureSessionContext(ctx, opts, req.Payload)	if opts.Alt == "responses/compact" {		return nil, statusErr{code: http.StatusBadRequest, msg: "streaming not supported for /responses/compact"}	}	if xaiInputHasItemType(req.Payload, "compaction_trigger") {		return e.executeCompactionTriggerStream(ctx, auth, req, opts)	} 	token, _ := xaiCreds(auth)	baseURL := xaiChatBaseURL(auth)	logXAIResolvedBaseURL(ctx, baseURL) 	prepared, err := e.prepareResponsesRequest(ctx, req, opts, true)	if err != nil {		return nil, err	} 	reporter := helps.NewExecutorUsageReporter(ctx, e, prepared.baseModel, auth)	defer reporter.TrackFailure(ctx, &err)	reporter.SetTranslatedReasoningEffort(prepared.body, e.Identifier()) 	url := strings.TrimSuffix(baseURL, "/") + "/responses"	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(prepared.body))	if err != nil {		return nil, err	}	applyXAIChatHeaders(httpReq, auth, token, true, prepared.sessionID, opts.Headers)	e.recordXAIRequest(ctx, auth, url, httpReq.Header.Clone(), prepared.body) 	httpClient := helps.NewProxyAwareHTTPClient(ctx, e.cfg, auth, 0)	httpClient = reporter.TrackHTTPClient(httpClient)	httpResp, err := httpClient.Do(httpReq)	if err != nil {		helps.RecordAPIResponseError(ctx, e.cfg, err)		return nil, err	}	helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())	if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {		data, errRead := io.ReadAll(httpResp.Body)		if errClose := httpResp.Body.Close(); errClose != nil {			log.Errorf("xai executor: close response body error: %v", errClose)		}		if errRead != nil {			helps.RecordAPIResponseError(ctx, e.cfg, errRead)			return nil, errRead		}		helps.AppendAPIResponseChunk(ctx, e.cfg, data)		helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data))		return nil, xaiStatusErr(httpResp.StatusCode, data)	} 	out := make(chan cliproxyexecutor.StreamChunk)	go func() {		defer close(out)		defer func() {			if errClose := httpResp.Body.Close(); errClose != nil {				log.Errorf("xai executor: close response body error: %v", errClose)			}		}()		scanner := bufio.NewScanner(httpResp.Body)		scanner.Buffer(nil, 52_428_800)		claudeInputTokens := helps.NewClaudeInputTokenState(prepared.from, prepared.to, prepared.responseFormat, prepared.originalPayload)		var param any		outputItemsByIndex := make(map[int64][]byte)		var outputItemsFallback [][]byte		responseFilter := newXAIInternalXSearchResponseFilter(prepared.filterInternalXSearch, prepared.clientDeclaredTools)		namespaceRestorer := newXAINamespaceRestorer(prepared.namespaceTools)		var pendingEventLine []byte		emitTranslatedLine := func(translatedLine []byte) bool {			chunks := helps.TranslateStreamWithClaudeInputTokens(ctx, prepared.to, prepared.responseFormat, req.Model, prepared.originalPayload, prepared.body, translatedLine, &param, claudeInputTokens)			for i := range chunks {				select {				case out <- cliproxyexecutor.StreamChunk{Payload: chunks[i]}:				case <-ctx.Done():					return false				}			}			return true		}		for scanner.Scan() {			line := scanner.Bytes()			helps.AppendAPIResponseChunk(ctx, e.cfg, line) 			if bytes.HasPrefix(line, xaiEventTag) {				if pendingEventLine != nil && !emitTranslatedLine(xaiNormalizeReasoningSummaryEventLine(pendingEventLine, "")) {					return				}				pendingEventLine = bytes.Clone(line)				continue			} 			if bytes.HasPrefix(line, xaiDataTag) {				eventDataList := xaiNormalizeReasoningSummaryDataEvents(bytes.TrimSpace(line[len(xaiDataTag):]))				hasPendingEventLine := pendingEventLine != nil				for i, eventData := range eventDataList {					eventData = namespaceRestorer.restore(eventData)					eventData = responseFilter.apply(eventData)					if len(eventData) == 0 {						if hasPendingEventLine && i == 0 {							pendingEventLine = nil						}						continue					}					normalizedEventName := gjson.GetBytes(eventData, "type").String()					switch normalizedEventName {					case "response.output_item.done":						xaiCollectOutputItemDone(eventData, outputItemsByIndex, &outputItemsFallback)					case "response.completed", "response.incomplete":						if detail, ok := helps.ParseCodexUsage(eventData); ok {							reporter.Publish(ctx, detail)						}						eventData = xaiPatchCompletedOutput(eventData, outputItemsByIndex, outputItemsFallback)						eventData = xaiNormalizeReasoningSummaryData(eventData)						if normalizedEventName == "response.completed" {							// A truncated turn carries no replayable terminal state, so only a							// completed response may refresh the reasoning replay cache.							cacheXAIReasoningReplayFromCompleted(ctx, prepared.replayScope, eventData)						}						normalizedEventName = gjson.GetBytes(eventData, "type").String()					} 					if hasPendingEventLine {						eventLine := []byte("event: " + normalizedEventName)						if i == 0 {							eventLine = xaiNormalizeReasoningSummaryEventLine(pendingEventLine, normalizedEventName)							pendingEventLine = nil						}						if !emitTranslatedLine(eventLine) {							return						}					}					if !emitTranslatedLine(append([]byte("data: "), eventData...)) {						return					}				}				continue			} 			if pendingEventLine != nil {				if !emitTranslatedLine(xaiNormalizeReasoningSummaryEventLine(pendingEventLine, "")) {					return				}				pendingEventLine = nil			}			if !emitTranslatedLine(bytes.Clone(line)) {				return			}		}		if pendingEventLine != nil {			emitTranslatedLine(xaiNormalizeReasoningSummaryEventLine(pendingEventLine, ""))		}		if errScan := scanner.Err(); errScan != nil {			helps.RecordAPIResponseError(ctx, e.cfg, errScan)			reporter.PublishFailure(ctx, errScan)			select {			case out <- cliproxyexecutor.StreamChunk{Err: errScan}:			case <-ctx.Done():			}		}	}()	return &cliproxyexecutor.StreamResult{Headers: httpResp.Header.Clone(), Chunks: out}, nil}