243 lines
7.6 KiB
Go
243 lines
7.6 KiB
Go
package llm
|
||
|
||
import (
|
||
"bytes"
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
"io"
|
||
"net/http"
|
||
"strings"
|
||
"time"
|
||
)
|
||
|
||
// Client calls a chat completions endpoint with streaming.
|
||
// 默认讲 OpenAI 兼容协议；Protocol 为 ProtocolAnthropic 时讲原生
|
||
// Anthropic Messages（组包与 BuildBody 共用，见 anthropic.go）。
|
||
type Client struct {
|
||
BaseURL string
|
||
APIKey string
|
||
HTTPClient *http.Client
|
||
// Protocol 空 = openai-chat（零值兼容所有直接构造点）。
|
||
Protocol Protocol
|
||
// ReasoningSplit 在 openai-chat 请求体加 "reasoning_split": true：
|
||
// MiniMax 等 interleaved-thinking 提供商据此把思考放进独立的
|
||
// reasoning_content 字段而不是混在 content 的 <think> 块里（分流后
|
||
// SSE 的 reasoning_content 捕获照常生效）。对不认识该字段的服务器
|
||
// 属未知字段，多数按忽略处理；宿主按模型名自行决定是否置位。
|
||
ReasoningSplit bool
|
||
// RouteContract optionally releases tool calls from an openai-chat stream
|
||
// that ended without [DONE]; nil or the zero value is off.
|
||
RouteContract *RouteContractEOF
|
||
}
|
||
|
||
// Stream implements StreamFunc. Transport and provider failures are encoded as
|
||
// KindError events; the returned channel is always non-nil and is closed.
|
||
func (c *Client) Stream(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions) <-chan Event {
|
||
out := make(chan Event)
|
||
go func() {
|
||
defer close(out)
|
||
if c.Protocol == ProtocolAnthropic {
|
||
c.streamAnthropic(ctx, model, messages, tools, opts, out)
|
||
return
|
||
}
|
||
c.run(ctx, model, messages, tools, opts, out)
|
||
}()
|
||
return out
|
||
}
|
||
|
||
func (c *Client) run(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions, out chan<- Event) {
|
||
base := strings.TrimRight(c.BaseURL, "/")
|
||
if opts.APIBase != "" {
|
||
base = strings.TrimRight(opts.APIBase, "/")
|
||
}
|
||
key := c.APIKey
|
||
if opts.APIKey != "" {
|
||
key = opts.APIKey
|
||
}
|
||
if base == "" || key == "" || model == "" {
|
||
emitError(out, "", false, "agentx llm client is missing base url, api key, or model")
|
||
return
|
||
}
|
||
if opts.MessageShape != "" {
|
||
messages = PrepareMessages(messages, opts.MessageShape)
|
||
}
|
||
body := map[string]any{
|
||
"model": model,
|
||
"messages": toProviderMessages(messages),
|
||
"stream": true,
|
||
"stream_options": map[string]any{
|
||
"include_usage": true,
|
||
},
|
||
}
|
||
if opts.MaxTokens > 0 {
|
||
body["max_tokens"] = opts.MaxTokens
|
||
}
|
||
if effort := OpenAIReasoningEffort(opts.ThinkingLevel); effort != "" {
|
||
body["reasoning_effort"] = effort
|
||
}
|
||
if c.ReasoningSplit {
|
||
body["reasoning_split"] = true
|
||
}
|
||
if len(tools) > 0 {
|
||
body["tools"] = toProviderTools(tools)
|
||
body["tool_choice"] = "auto"
|
||
}
|
||
raw, err := json.Marshal(body)
|
||
if err != nil {
|
||
emitError(out, "", false, err.Error())
|
||
return
|
||
}
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, base+"/chat/completions", bytes.NewReader(raw))
|
||
if err != nil {
|
||
emitError(out, "", false, err.Error())
|
||
return
|
||
}
|
||
req.Header.Set("Authorization", "Bearer "+key)
|
||
req.Header.Set("Content-Type", "application/json")
|
||
req.Header.Set("Accept", "text/event-stream")
|
||
hc := c.HTTPClient
|
||
if hc == nil {
|
||
hc = &http.Client{Timeout: 30 * time.Minute}
|
||
}
|
||
resp, err := hc.Do(req)
|
||
if err != nil {
|
||
if ctx.Err() != nil {
|
||
emitAbort(out, "", ctx.Err(), &TerminalObservation{ReadOutcome: "cancelled"})
|
||
return
|
||
}
|
||
emitError(out, "", true, err.Error(), &TerminalObservation{ReadOutcome: "transport_error"})
|
||
return
|
||
}
|
||
defer resp.Body.Close()
|
||
if resp.StatusCode >= 400 {
|
||
slurp, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
|
||
msg := fmt.Sprintf("provider status %d: %s", resp.StatusCode, strings.TrimSpace(string(slurp)))
|
||
ev := Event{
|
||
Kind: KindError,
|
||
HTTPStatus: resp.StatusCode,
|
||
TerminalObservation: &TerminalObservation{ReadOutcome: "provider_error"},
|
||
Retryable: StatusRetryable(resp.StatusCode) && !ContextOverflow(msg),
|
||
ErrorMessage: msg,
|
||
Err: &StreamError{Message: msg, Retryable: StatusRetryable(resp.StatusCode) && !ContextOverflow(msg)},
|
||
Message: Message{Role: RoleAssistant, StopReason: StopError},
|
||
}
|
||
if d, ok := ParseRetryAfter(resp.Header.Get("Retry-After"), time.Now()); ok {
|
||
ev.HasRetryAfter = true
|
||
ev.RetryAfter = d
|
||
}
|
||
out <- ev
|
||
return
|
||
}
|
||
observed := newBodyObserver(resp)
|
||
contract := c.newRouteContract(model, base, req, resp, tools)
|
||
var facts sseFacts
|
||
var hooks sseHooks
|
||
if contract != nil {
|
||
hooks.facts = &facts
|
||
hooks.decideEOF = func(success, failure Event) Event {
|
||
check := contract.check(observed.snapshot(), failure.TerminalObservation, facts, success.Message.ToolCalls)
|
||
if check != "" {
|
||
failure.TerminalObservation.ContractCheck = check
|
||
return failure
|
||
}
|
||
success.CompletionBasis = CompletionRouteContractEOF
|
||
success.TerminalObservation.CompletionBasis = CompletionRouteContractEOF
|
||
return success
|
||
}
|
||
}
|
||
for ev := range parseSSE(ctx, observed, hooks) {
|
||
if ev.Kind == KindMessageEnd || ev.Kind == KindError {
|
||
ev.HTTPStatus = resp.StatusCode
|
||
// ParseSSE sends its terminal event only after its last body read.
|
||
if obs := ev.TerminalObservation; obs != nil {
|
||
obs.BodyBoundary = observed.snapshot()
|
||
if ev.Kind == KindMessageEnd && ev.Err == nil && ev.StreamComplete {
|
||
ev.CompletionBasis = CompletionDoneMarker
|
||
obs.CompletionBasis = CompletionDoneMarker
|
||
}
|
||
if contract != nil && ev.Err != nil && !ev.StreamComplete &&
|
||
ev.Message.StopReason != StopAborted && obs.ContractCheck == "" {
|
||
obs.ContractCheck = contract.check(obs.BodyBoundary, obs, facts, ev.Message.ToolCalls)
|
||
}
|
||
}
|
||
}
|
||
out <- ev
|
||
}
|
||
}
|
||
|
||
// replayThink 拼接一条 assistant 消息捕获到的全部思考块。
|
||
func replayThink(m Message) string {
|
||
var b strings.Builder
|
||
for _, t := range m.Thinking {
|
||
if t.Thinking != "" {
|
||
b.WriteString(t.Thinking)
|
||
b.WriteString("\n")
|
||
}
|
||
}
|
||
return strings.TrimRight(b.String(), "\n")
|
||
}
|
||
|
||
func toProviderMessages(messages []Message) []map[string]any {
|
||
out := make([]map[string]any, 0, len(messages))
|
||
for _, m := range messages {
|
||
switch m.Role {
|
||
case RoleTool:
|
||
out = append(out, map[string]any{
|
||
"role": "tool",
|
||
"tool_call_id": m.ToolCallID,
|
||
"content": m.Content,
|
||
})
|
||
case RoleAssistant:
|
||
// 回放捕获的思考（Message.Thinking 由 SSE 的 reasoning_content 捕获而来）：
|
||
// 交错思考模型要求历史 assistant 原样带回 <think> 块。没有思考块的
|
||
// 消息零改动——非思考模型的上下文不受影响。
|
||
content := m.Content
|
||
if think := replayThink(m); think != "" {
|
||
content = "<think>\n" + think + "\n</think>\n\n" + content
|
||
}
|
||
item := map[string]any{"role": "assistant", "content": content}
|
||
if len(m.ToolCalls) > 0 {
|
||
calls := make([]map[string]any, 0, len(m.ToolCalls))
|
||
for _, tc := range m.ToolCalls {
|
||
calls = append(calls, map[string]any{
|
||
"id": tc.ID,
|
||
"type": "function",
|
||
"function": map[string]any{
|
||
"name": tc.Name,
|
||
"arguments": tc.Arguments,
|
||
},
|
||
})
|
||
}
|
||
item["tool_calls"] = calls
|
||
}
|
||
out = append(out, item)
|
||
default:
|
||
out = append(out, map[string]any{
|
||
"role": string(m.Role),
|
||
"content": messageContent(m),
|
||
})
|
||
}
|
||
}
|
||
return out
|
||
}
|
||
|
||
func toProviderTools(tools []ToolDecl) []map[string]any {
|
||
out := make([]map[string]any, 0, len(tools))
|
||
for _, t := range tools {
|
||
params := t.Parameters
|
||
if len(params) == 0 {
|
||
params = json.RawMessage(`{"type":"object","properties":{}}`)
|
||
}
|
||
out = append(out, map[string]any{
|
||
"type": "function",
|
||
"function": map[string]any{
|
||
"name": t.Name,
|
||
"description": t.Description,
|
||
"parameters": json.RawMessage(params),
|
||
},
|
||
})
|
||
}
|
||
return out
|
||
}
|