Refactor llm package
This commit is contained in:
dwrz
+280
-352
@@ -9,119 +9,70 @@ import (
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"slices"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// ErrMaxIterations is returned when the agent loop exceeds the configured
|
||||
// maximum number of iterations.
|
||||
var ErrMaxIterations = fmt.Errorf("max iterations exceeded")
|
||||
|
||||
// CompletionsResult bundles the messages generated during a completions
|
||||
// request along with usage and timings from the final API response.
|
||||
type CompletionsResult struct {
|
||||
Messages []Message
|
||||
Usage *Usage
|
||||
Timings *Timings
|
||||
}
|
||||
|
||||
// CompletionsParams holds parameters for Completions and CompletionsStream.
|
||||
// CompletionsParams contains the inputs to a chat completion.
|
||||
type CompletionsParams struct {
|
||||
Model string
|
||||
SystemMessage string
|
||||
MaxIterations int
|
||||
Messages []Message
|
||||
// Messages is the conversation presented to the model, in order.
|
||||
Messages []Message
|
||||
// Model identifies the model according to the configured endpoint.
|
||||
Model string
|
||||
// ThinkingBudgetTokens is sent as the endpoint's reasoning-token budget.
|
||||
// Nil omits both the budget and enable_thinking; zero sets enable_thinking
|
||||
// to false; any other value sets it to true.
|
||||
ThinkingBudgetTokens *int
|
||||
// Tools describes the functions available for the model to call.
|
||||
Tools []APITool
|
||||
}
|
||||
|
||||
// Completions sends messages to the LLM using the specified model.
|
||||
// If systemMessage is non-empty, it overrides the configured system message.
|
||||
// maxIterations caps the number of agent loop iterations.
|
||||
// A value of 0 means unlimited (no cap).
|
||||
// Returns all messages generated during the completions request, including
|
||||
// tool calls and tool results, along with usage and timings from the final response.
|
||||
// CompletionResult contains one generated Message and its response metadata.
|
||||
type CompletionResult struct {
|
||||
// FinishReason is the endpoint's reason for ending generation.
|
||||
FinishReason string
|
||||
// Message is the generated message.
|
||||
Message Message
|
||||
// Timings contains llama.cpp performance data, when reported.
|
||||
Timings *Timings
|
||||
// Usage contains token counts, when reported.
|
||||
Usage *Usage
|
||||
}
|
||||
|
||||
// Completions requests a non-streaming chat completion and returns the first
|
||||
// choice. It returns an error if the response contains no choices.
|
||||
func (c *Client) Completions(
|
||||
ctx context.Context, p CompletionsParams,
|
||||
) (*CompletionsResult, error) {
|
||||
) (*CompletionResult, error) {
|
||||
ctx, cancel := context.WithTimeout(ctx, c.timeout)
|
||||
defer cancel()
|
||||
|
||||
model, systemMessage, maxIterations, messages, thinkingBudgetTokens :=
|
||||
p.Model, p.SystemMessage, p.MaxIterations, p.Messages, p.ThinkingBudgetTokens
|
||||
|
||||
// Use request system message if provided, config default otherwise.
|
||||
smsg := c.systemMessage
|
||||
if systemMessage != "" {
|
||||
smsg = systemMessage
|
||||
cr, err := c.completions(ctx, p)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// Prepend system message, if configured and not already present.
|
||||
if smsg != "" &&
|
||||
(len(messages) == 0 || messages[0].Role != RoleSystem) {
|
||||
messages = append(
|
||||
[]Message{{
|
||||
Role: RoleSystem,
|
||||
ContentParts: []ContentPart{{
|
||||
Type: ContentTypeText,
|
||||
Text: smsg,
|
||||
}},
|
||||
}},
|
||||
messages...,
|
||||
)
|
||||
}
|
||||
|
||||
// Remember the starting length so we can return only the messages
|
||||
// generated during this call (the suffix of the conversation history).
|
||||
start := len(messages)
|
||||
// Loop for tool calls. maxIterations == 0 means unlimited.
|
||||
var finalUsage *Usage
|
||||
var finalTimings *Timings
|
||||
for i := 0; maxIterations == 0 || i < maxIterations; i++ {
|
||||
cr, err := c.completions(ctx, model, thinkingBudgetTokens, messages)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("completions: %w", err)
|
||||
}
|
||||
// Capture usage and timings from this iteration.
|
||||
if cr.Usage != nil {
|
||||
finalUsage = cr.Usage
|
||||
}
|
||||
if cr.Timings != nil {
|
||||
finalTimings = cr.Timings
|
||||
}
|
||||
// Use first choice for tool call loop.
|
||||
message := &cr.Choices[0].Message
|
||||
messages = append(messages, *message)
|
||||
if len(message.ToolCalls) == 0 {
|
||||
return &CompletionsResult{
|
||||
Messages: messages[start:],
|
||||
Usage: finalUsage,
|
||||
Timings: finalTimings,
|
||||
}, nil
|
||||
}
|
||||
toolResults := c.callTools(ctx, *message)
|
||||
messages = append(messages, toolResults...)
|
||||
}
|
||||
|
||||
return &CompletionsResult{
|
||||
Messages: messages[start:],
|
||||
Usage: finalUsage,
|
||||
Timings: finalTimings,
|
||||
}, ErrMaxIterations
|
||||
choice := cr.Choices[0]
|
||||
return &CompletionResult{
|
||||
Message: choice.Message,
|
||||
FinishReason: choice.FinishReason,
|
||||
Usage: cr.Usage,
|
||||
Timings: cr.Timings,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// completions sends a single non-streaming request and returns the
|
||||
// full API response.
|
||||
func (c *Client) completions(
|
||||
ctx context.Context, model string, thinkingBudgetTokens *int, messages []Message,
|
||||
ctx context.Context, p CompletionsParams,
|
||||
) (*ChatResponse, error) {
|
||||
kwargs := map[string]any{}
|
||||
if thinkingBudgetTokens != nil {
|
||||
kwargs["enable_thinking"] = *thinkingBudgetTokens != 0
|
||||
if p.ThinkingBudgetTokens != nil {
|
||||
kwargs[chatTemplateEnableThinking] = *p.ThinkingBudgetTokens != 0
|
||||
}
|
||||
body, err := json.Marshal(ChatRequest{
|
||||
Model: model,
|
||||
Messages: messages,
|
||||
Tools: c.tools,
|
||||
ThinkingBudgetTokens: thinkingBudgetTokens,
|
||||
Model: p.Model,
|
||||
Messages: p.Messages,
|
||||
Tools: p.Tools,
|
||||
ThinkingBudgetTokens: p.ThinkingBudgetTokens,
|
||||
ChatTemplateKWARGS: kwargs,
|
||||
ParallelToolCalls: new(true),
|
||||
})
|
||||
@@ -133,7 +84,7 @@ func (c *Client) completions(
|
||||
req, err := http.NewRequestWithContext(
|
||||
ctx,
|
||||
http.MethodPost,
|
||||
c.baseURL+"/chat/completions",
|
||||
c.baseURL+chatCompletionsPath,
|
||||
bytes.NewReader(body),
|
||||
)
|
||||
if err != nil {
|
||||
@@ -163,11 +114,7 @@ func (c *Client) completions(
|
||||
)
|
||||
}
|
||||
var apiError struct {
|
||||
Error *struct {
|
||||
Type string `json:"type"`
|
||||
Code string `json:"code"`
|
||||
Message string `json:"message"`
|
||||
} `json:"error"`
|
||||
Error *APIError `json:"error"`
|
||||
}
|
||||
if err := json.Unmarshal(
|
||||
body, &apiError,
|
||||
@@ -182,11 +129,7 @@ func (c *Client) completions(
|
||||
),
|
||||
slog.String("type", apiError.Error.Type),
|
||||
)
|
||||
return nil, fmt.Errorf(
|
||||
"completions: %s: %s",
|
||||
apiError.Error.Code,
|
||||
apiError.Error.Message,
|
||||
)
|
||||
return nil, fmt.Errorf("completions: %w", apiError.Error)
|
||||
}
|
||||
c.log.ErrorContext(
|
||||
ctx,
|
||||
@@ -212,131 +155,140 @@ func (c *Client) completions(
|
||||
return cr, nil
|
||||
}
|
||||
|
||||
// StreamEvent wraps a message or delta produced during streaming.
|
||||
type StreamEvent struct {
|
||||
// Type identifies the event kind: delta, reasoning_delta, done, or message.
|
||||
Type string
|
||||
// Delta is a text fragment from the assistant response.
|
||||
Delta string
|
||||
// ReasoningDelta is a reasoning text fragment from the assistant response.
|
||||
ReasoningDelta string
|
||||
// Message is the complete chat completion message.
|
||||
Message Message
|
||||
// Usage is token usage info from the final chunk, sent on the done event.
|
||||
Usage *Usage
|
||||
// Timings is llama.cpp timing info from the final chunk, sent on the done event.
|
||||
Timings *Timings
|
||||
// StreamAccumulator reconstructs a Message from the ordered deltas for one
|
||||
// streaming choice. The zero value is ready for use.
|
||||
type StreamAccumulator struct {
|
||||
// Content contains the accumulated message text.
|
||||
Content strings.Builder
|
||||
// Reasoning combines fragments received under either supported reasoning
|
||||
// field name.
|
||||
Reasoning strings.Builder
|
||||
// Refusal contains the accumulated refusal text.
|
||||
Refusal strings.Builder
|
||||
// Role is the latest non-empty role received.
|
||||
Role Role
|
||||
// ToolCalls contains calls assembled by tool-call index.
|
||||
ToolCalls []ToolCall
|
||||
}
|
||||
|
||||
const (
|
||||
streamEventDelta = "delta"
|
||||
streamEventReasoningDelta = "reasoning_delta"
|
||||
streamEventDone = "done"
|
||||
streamEventMessage = "message"
|
||||
)
|
||||
// Add incorporates the next delta. Content, refusal, reasoning, and tool-call
|
||||
// fragments are appended in arrival order. Tool-call fragments are grouped by
|
||||
// Index; a missing index denotes the first call. A non-empty Role replaces the
|
||||
// previously accumulated role.
|
||||
func (acc *StreamAccumulator) Add(delta Delta) {
|
||||
if delta.Role != "" {
|
||||
acc.Role = delta.Role
|
||||
}
|
||||
acc.Content.WriteString(delta.Content)
|
||||
acc.Refusal.WriteString(delta.Refusal)
|
||||
acc.Reasoning.WriteString(delta.ReasoningContent)
|
||||
acc.Reasoning.WriteString(delta.Reasoning)
|
||||
|
||||
// CompletionsStream sends messages to the LLM using the specified model and
|
||||
// streams results. Each complete message (assistant reply, tool call,
|
||||
// tool result) is sent to the events channel as it becomes available.
|
||||
// The channel is closed before returning.
|
||||
// If systemMessage is non-empty, it overrides the configured system message.
|
||||
// maxIterations caps the number of agent loop iterations; a value of 0
|
||||
// means unlimited (no cap).
|
||||
// Returns ErrMaxIterations if the iteration limit is exceeded.
|
||||
for _, tc := range delta.ToolCalls {
|
||||
idx := 0
|
||||
if tc.Index != nil {
|
||||
idx = *tc.Index
|
||||
}
|
||||
for len(acc.ToolCalls) <= idx {
|
||||
acc.ToolCalls = append(
|
||||
acc.ToolCalls,
|
||||
ToolCall{Type: ToolTypeFunction},
|
||||
)
|
||||
}
|
||||
if tc.ID != "" {
|
||||
acc.ToolCalls[idx].ID = tc.ID
|
||||
}
|
||||
if tc.Type != "" {
|
||||
acc.ToolCalls[idx].Type = tc.Type
|
||||
}
|
||||
acc.ToolCalls[idx].Function.Name += tc.Function.Name
|
||||
acc.ToolCalls[idx].Function.Arguments += tc.Function.Arguments
|
||||
}
|
||||
}
|
||||
|
||||
// Message returns the accumulated state without resetting the accumulator. If
|
||||
// the stream contains a refusal, the returned content consists of that refusal
|
||||
// rather than the accumulated text.
|
||||
func (acc *StreamAccumulator) Message() Message {
|
||||
var contentParts []ContentPart
|
||||
if refusal := acc.Refusal.String(); refusal != "" {
|
||||
contentParts = []ContentPart{{
|
||||
Type: ContentTypeRefusal,
|
||||
Refusal: refusal,
|
||||
}}
|
||||
} else {
|
||||
contentParts = []ContentPart{{
|
||||
Type: ContentTypeText,
|
||||
Text: acc.Content.String(),
|
||||
}}
|
||||
}
|
||||
|
||||
return Message{
|
||||
ContentParts: contentParts,
|
||||
ReasoningContent: acc.Reasoning.String(),
|
||||
Refusal: acc.Refusal.String(),
|
||||
Role: acc.Role,
|
||||
ToolCalls: acc.ToolCalls,
|
||||
}
|
||||
}
|
||||
|
||||
// CompletionStream is a caller-driven chat-completion stream. A stream owns
|
||||
// its response body and timeout context until it terminates or is closed.
|
||||
type CompletionStream struct {
|
||||
ctx context.Context
|
||||
reader *bufio.Reader
|
||||
body io.ReadCloser
|
||||
cancel context.CancelFunc
|
||||
|
||||
acc StreamAccumulator
|
||||
current StreamChunk
|
||||
finishReason string
|
||||
usage *Usage
|
||||
timings *Timings
|
||||
err error
|
||||
cleanupErr error
|
||||
terminated bool
|
||||
resourcesDone bool
|
||||
}
|
||||
|
||||
// CompletionsStream opens a streaming chat completion. Request setup,
|
||||
// transport, and non-success response errors are returned before a stream is
|
||||
// created.
|
||||
func (c *Client) CompletionsStream(
|
||||
ctx context.Context, p CompletionsParams, events chan<- StreamEvent,
|
||||
) error {
|
||||
defer close(events)
|
||||
ctx context.Context,
|
||||
p CompletionsParams,
|
||||
) (*CompletionStream, error) {
|
||||
ctx, cancel := context.WithTimeout(ctx, c.timeout)
|
||||
defer cancel()
|
||||
|
||||
model, systemMessage, maxIterations, messages, thinkingBudgetTokens :=
|
||||
p.Model, p.SystemMessage, p.MaxIterations, p.Messages, p.ThinkingBudgetTokens
|
||||
|
||||
// Use request system message if provided, config default otherwise.
|
||||
smsg := c.systemMessage
|
||||
if systemMessage != "" {
|
||||
smsg = systemMessage
|
||||
}
|
||||
|
||||
// Prepend system message, if configured and not already present.
|
||||
if smsg != "" &&
|
||||
(len(messages) == 0 || messages[0].Role != RoleSystem) {
|
||||
messages = append(
|
||||
[]Message{{
|
||||
Role: RoleSystem,
|
||||
ContentParts: []ContentPart{
|
||||
{Type: ContentTypeText, Text: smsg},
|
||||
},
|
||||
}},
|
||||
messages...,
|
||||
)
|
||||
}
|
||||
|
||||
// Loop for tool calls. maxIterations == 0 means unlimited.
|
||||
for i := 0; maxIterations == 0 || i < maxIterations; i++ {
|
||||
choices, err := c.completionsStream(ctx, messages, model, thinkingBudgetTokens, events)
|
||||
if err != nil {
|
||||
return fmt.Errorf("completions stream: %w", err)
|
||||
}
|
||||
// Use first choice for tool call loop.
|
||||
message := choices[0]
|
||||
if len(message.ToolCalls) == 0 {
|
||||
return nil
|
||||
}
|
||||
toolResults := c.callTools(ctx, message)
|
||||
for _, tr := range toolResults {
|
||||
events <- StreamEvent{
|
||||
Type: streamEventMessage,
|
||||
Message: tr,
|
||||
}
|
||||
}
|
||||
messages = append(messages, choices...)
|
||||
messages = append(messages, toolResults...)
|
||||
}
|
||||
|
||||
return ErrMaxIterations
|
||||
}
|
||||
|
||||
// streamAccum accumulates deltas for a single choice in a streaming response.
|
||||
type streamAccum struct {
|
||||
content strings.Builder
|
||||
reasoning strings.Builder
|
||||
refusal strings.Builder
|
||||
role string
|
||||
toolCalls []ToolCall
|
||||
}
|
||||
|
||||
// completionsStream sends a single streaming request, accumulates the
|
||||
// response, and returns all assistant messages from the response choices.
|
||||
// Sends each message to the events channel. Usage and timings are included
|
||||
// on the done event for the last message.
|
||||
func (c *Client) completionsStream(ctx context.Context, messages []Message, model string, thinkingBudgetTokens *int, events chan<- StreamEvent) ([]Message, error) {
|
||||
kwargs := map[string]any{}
|
||||
if thinkingBudgetTokens != nil {
|
||||
kwargs["enable_thinking"] = *thinkingBudgetTokens != 0
|
||||
if p.ThinkingBudgetTokens != nil {
|
||||
kwargs[chatTemplateEnableThinking] = *p.ThinkingBudgetTokens != 0
|
||||
}
|
||||
body, err := json.Marshal(ChatRequest{
|
||||
Model: model,
|
||||
Messages: messages,
|
||||
Stream: true,
|
||||
Tools: c.tools,
|
||||
ThinkingBudgetTokens: thinkingBudgetTokens,
|
||||
ChatTemplateKWARGS: kwargs,
|
||||
Messages: p.Messages,
|
||||
Model: p.Model,
|
||||
ParallelToolCalls: new(true),
|
||||
Stream: true,
|
||||
StreamOptions: &StreamOptions{IncludeUsage: true},
|
||||
ThinkingBudgetTokens: p.ThinkingBudgetTokens,
|
||||
Tools: p.Tools,
|
||||
})
|
||||
if err != nil {
|
||||
cancel()
|
||||
return nil, fmt.Errorf(
|
||||
"completions stream: marshal request body: %w", err,
|
||||
)
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(
|
||||
ctx, http.MethodPost,
|
||||
c.baseURL+"/chat/completions",
|
||||
ctx,
|
||||
http.MethodPost,
|
||||
c.baseURL+chatCompletionsPath,
|
||||
bytes.NewReader(body),
|
||||
)
|
||||
if err != nil {
|
||||
cancel()
|
||||
return nil, fmt.Errorf(
|
||||
"completions stream: create request: %w", err,
|
||||
)
|
||||
@@ -349,12 +301,14 @@ func (c *Client) completionsStream(ctx context.Context, messages []Message, mode
|
||||
|
||||
res, err := c.httpc.Do(req)
|
||||
if err != nil {
|
||||
cancel()
|
||||
return nil, fmt.Errorf("completions stream: %w", err)
|
||||
}
|
||||
|
||||
if res.StatusCode != http.StatusOK {
|
||||
defer cancel()
|
||||
defer res.Body.Close()
|
||||
body, err := io.ReadAll(res.Body)
|
||||
res.Body.Close()
|
||||
if err != nil {
|
||||
c.log.ErrorContext(
|
||||
ctx,
|
||||
@@ -367,11 +321,7 @@ func (c *Client) completionsStream(ctx context.Context, messages []Message, mode
|
||||
)
|
||||
}
|
||||
var apiError struct {
|
||||
Error *struct {
|
||||
Type string `json:"type"`
|
||||
Code string `json:"code"`
|
||||
Message string `json:"message"`
|
||||
} `json:"error"`
|
||||
Error *APIError `json:"error"`
|
||||
}
|
||||
if err := json.Unmarshal(
|
||||
body, &apiError,
|
||||
@@ -381,15 +331,11 @@ func (c *Client) completionsStream(ctx context.Context, messages []Message, mode
|
||||
"completions stream: request error",
|
||||
slog.Int("status", res.StatusCode),
|
||||
slog.String("code", apiError.Error.Code),
|
||||
slog.String(
|
||||
"message", apiError.Error.Message,
|
||||
),
|
||||
slog.String("message", apiError.Error.Message),
|
||||
slog.String("type", apiError.Error.Type),
|
||||
)
|
||||
return nil, fmt.Errorf(
|
||||
"completions stream: %s: %s",
|
||||
apiError.Error.Code,
|
||||
apiError.Error.Message,
|
||||
"completions stream: %w", apiError.Error,
|
||||
)
|
||||
}
|
||||
c.log.ErrorContext(
|
||||
@@ -400,161 +346,143 @@ func (c *Client) completionsStream(ctx context.Context, messages []Message, mode
|
||||
)
|
||||
return nil, fmt.Errorf(
|
||||
"completions stream: %d: %s",
|
||||
res.StatusCode, string(body),
|
||||
res.StatusCode,
|
||||
string(body),
|
||||
)
|
||||
}
|
||||
|
||||
// Accumulate the streamed response per choice index.
|
||||
accums := make(map[int]*streamAccum)
|
||||
var usage *Usage
|
||||
var timings *Timings
|
||||
return &CompletionStream{
|
||||
ctx: ctx,
|
||||
reader: bufio.NewReader(res.Body),
|
||||
body: res.Body,
|
||||
cancel: cancel,
|
||||
}, nil
|
||||
}
|
||||
|
||||
scanner := bufio.NewScanner(res.Body)
|
||||
for scanner.Scan() {
|
||||
line := scanner.Text()
|
||||
// Skip empty lines and comments.
|
||||
if line == "" || line == ":" {
|
||||
continue
|
||||
}
|
||||
// Parse data lines.
|
||||
if !strings.HasPrefix(line, "data: ") {
|
||||
continue
|
||||
}
|
||||
data := strings.TrimPrefix(line, "data: ")
|
||||
// [DONE] signals end of stream.
|
||||
if data == "[DONE]" {
|
||||
break
|
||||
// Next advances to the next JSON data event. It returns false after [DONE], an
|
||||
// error, or Close. On an error, Err reports the cause and Result remains
|
||||
// available with all chunks accumulated before the error.
|
||||
func (s *CompletionStream) Next() bool {
|
||||
if s.terminated {
|
||||
return false
|
||||
}
|
||||
|
||||
var data []string
|
||||
for {
|
||||
line, err := s.reader.ReadString('\n')
|
||||
if err != nil && err != io.EOF {
|
||||
s.fail(fmt.Errorf("completions stream: read response: %w", err))
|
||||
return false
|
||||
}
|
||||
var chunk StreamChunk
|
||||
if err := json.Unmarshal([]byte(data), &chunk); err != nil {
|
||||
continue
|
||||
}
|
||||
for _, ch := range chunk.Choices {
|
||||
acc, ok := accums[ch.Index]
|
||||
if !ok {
|
||||
acc = &streamAccum{}
|
||||
accums[ch.Index] = acc
|
||||
}
|
||||
delta := ch.Delta
|
||||
if delta.Role != "" {
|
||||
acc.role = delta.Role
|
||||
}
|
||||
if delta.Content != "" {
|
||||
acc.content.WriteString(delta.Content)
|
||||
events <- StreamEvent{
|
||||
Type: streamEventDelta,
|
||||
Delta: delta.Content,
|
||||
}
|
||||
}
|
||||
if delta.Refusal != "" {
|
||||
acc.refusal.WriteString(delta.Refusal)
|
||||
}
|
||||
if delta.ReasoningContent != "" {
|
||||
acc.reasoning.WriteString(
|
||||
delta.ReasoningContent,
|
||||
)
|
||||
events <- StreamEvent{
|
||||
Type: streamEventReasoningDelta,
|
||||
ReasoningDelta: delta.ReasoningContent,
|
||||
}
|
||||
}
|
||||
if delta.Reasoning != "" {
|
||||
acc.reasoning.WriteString(delta.Reasoning)
|
||||
events <- StreamEvent{
|
||||
Type: streamEventReasoningDelta,
|
||||
ReasoningDelta: delta.Reasoning,
|
||||
}
|
||||
}
|
||||
// Accumulate tool call deltas by index.
|
||||
for _, tc := range delta.ToolCalls {
|
||||
idx := 0
|
||||
if tc.Index != nil {
|
||||
idx = *tc.Index
|
||||
}
|
||||
// Grow the slice as needed.
|
||||
for len(acc.toolCalls) <= idx {
|
||||
acc.toolCalls = append(
|
||||
acc.toolCalls,
|
||||
ToolCall{Type: "function"},
|
||||
)
|
||||
if len(line) > 0 {
|
||||
line = strings.TrimSuffix(line, "\n")
|
||||
line = strings.TrimSuffix(line, "\r")
|
||||
if line == "" {
|
||||
if len(data) > 0 {
|
||||
break
|
||||
}
|
||||
if tc.ID != "" {
|
||||
acc.toolCalls[idx].ID = tc.ID
|
||||
} else if line[0] != ':' {
|
||||
field, value, found := strings.Cut(line, ":")
|
||||
if !found {
|
||||
field = line
|
||||
value = ""
|
||||
}
|
||||
if tc.Function.Name != "" {
|
||||
acc.toolCalls[idx].Function.Name +=
|
||||
tc.Function.Name
|
||||
value = strings.TrimPrefix(value, " ")
|
||||
if field == "data" {
|
||||
data = append(data, value)
|
||||
}
|
||||
if tc.Function.Arguments != "" {
|
||||
acc.toolCalls[idx].Function.Arguments +=
|
||||
tc.Function.Arguments
|
||||
}
|
||||
}
|
||||
}
|
||||
// Capture usage from the final chunk (has null choices).
|
||||
if chunk.Usage != nil {
|
||||
usage = chunk.Usage
|
||||
}
|
||||
if chunk.Timings != nil {
|
||||
timings = chunk.Timings
|
||||
if err == io.EOF {
|
||||
s.fail(io.ErrUnexpectedEOF)
|
||||
return false
|
||||
}
|
||||
}
|
||||
res.Body.Close()
|
||||
|
||||
if err := scanner.Err(); err != nil {
|
||||
return nil, fmt.Errorf(
|
||||
"completions stream: read error: %w", err,
|
||||
)
|
||||
payload := strings.Join(data, "\n")
|
||||
if payload == sseDone {
|
||||
s.terminated = true
|
||||
s.release()
|
||||
return false
|
||||
}
|
||||
|
||||
// Build messages in index order.
|
||||
indices := make([]int, 0, len(accums))
|
||||
for idx := range accums {
|
||||
indices = append(indices, idx)
|
||||
var chunk StreamChunk
|
||||
if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
|
||||
s.fail(fmt.Errorf("completions stream: decode chunk: %w", err))
|
||||
return false
|
||||
}
|
||||
for _, choice := range chunk.Choices {
|
||||
if choice.Index != 0 {
|
||||
s.fail(fmt.Errorf(
|
||||
"completions stream: unexpected choice index %d",
|
||||
choice.Index,
|
||||
))
|
||||
return false
|
||||
}
|
||||
if choice.FinishReason != "" {
|
||||
s.finishReason = choice.FinishReason
|
||||
}
|
||||
s.acc.Add(choice.Delta)
|
||||
}
|
||||
if chunk.Usage != nil {
|
||||
s.usage = chunk.Usage
|
||||
}
|
||||
slices.Sort(indices)
|
||||
if chunk.Timings != nil {
|
||||
s.timings = chunk.Timings
|
||||
}
|
||||
s.current = chunk
|
||||
return true
|
||||
}
|
||||
|
||||
msgs := make([]Message, 0, len(accums))
|
||||
for i, idx := range indices {
|
||||
acc := accums[idx]
|
||||
// Current returns the most recently decoded chunk. The value is replaced by
|
||||
// the next successful call to Next.
|
||||
func (s *CompletionStream) Current() StreamChunk {
|
||||
return s.current
|
||||
}
|
||||
|
||||
// Content and refusal are mutually exclusive in the API.
|
||||
var contentParts []ContentPart
|
||||
if refusalStr := acc.refusal.String(); refusalStr != "" {
|
||||
contentParts = []ContentPart{{
|
||||
Type: ContentTypeRefusal,
|
||||
Refusal: refusalStr,
|
||||
}}
|
||||
} else {
|
||||
contentParts = []ContentPart{{
|
||||
Type: ContentTypeText,
|
||||
Text: acc.content.String(),
|
||||
}}
|
||||
}
|
||||
// Result returns a snapshot of the state accumulated so far. Err reports
|
||||
// whether the snapshot is a complete response or a partial result.
|
||||
func (s *CompletionStream) Result() CompletionResult {
|
||||
message := s.acc.Message()
|
||||
message.ToolCalls = append([]ToolCall(nil), message.ToolCalls...)
|
||||
return CompletionResult{
|
||||
Message: message,
|
||||
FinishReason: s.finishReason,
|
||||
Usage: s.usage,
|
||||
Timings: s.timings,
|
||||
}
|
||||
}
|
||||
|
||||
// Err returns the stream failure, if any. Closing a stream does not create or
|
||||
// replace a stream failure.
|
||||
func (s *CompletionStream) Err() error {
|
||||
return s.err
|
||||
}
|
||||
|
||||
message := Message{
|
||||
Role: acc.role,
|
||||
ContentParts: contentParts,
|
||||
Refusal: acc.refusal.String(),
|
||||
ReasoningContent: acc.reasoning.String(),
|
||||
ToolCalls: acc.toolCalls,
|
||||
}
|
||||
// Close releases the response body and timeout context. It is safe to call
|
||||
// repeatedly and returns the first response-body cleanup error.
|
||||
func (s *CompletionStream) Close() error {
|
||||
s.terminated = true
|
||||
s.release()
|
||||
return s.cleanupErr
|
||||
}
|
||||
|
||||
// Attach usage and timings to the done event of the last message.
|
||||
var evUsage *Usage
|
||||
var evTimings *Timings
|
||||
if i == len(indices)-1 {
|
||||
evUsage = usage
|
||||
evTimings = timings
|
||||
}
|
||||
events <- StreamEvent{
|
||||
Type: streamEventDone,
|
||||
Message: message,
|
||||
Usage: evUsage,
|
||||
Timings: evTimings,
|
||||
}
|
||||
msgs = append(msgs, message)
|
||||
func (s *CompletionStream) fail(err error) {
|
||||
if ctxErr := s.ctx.Err(); ctxErr != nil {
|
||||
err = ctxErr
|
||||
}
|
||||
s.err = err
|
||||
s.terminated = true
|
||||
s.release()
|
||||
}
|
||||
|
||||
return msgs, nil
|
||||
func (s *CompletionStream) release() {
|
||||
if s.resourcesDone {
|
||||
return
|
||||
}
|
||||
s.resourcesDone = true
|
||||
if err := s.body.Close(); err != nil {
|
||||
s.cleanupErr = err
|
||||
}
|
||||
s.cancel()
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user