255 lines
7.9 KiB
Go
255 lines
7.9 KiB
Go
// Package llm provides a thin OpenAI-compatible client used by the agent
|
|
// runtime. It hides the openai-go SDK types behind a small, provider-agnostic
|
|
// surface so that swapping the backend (a different SDK, a custom HTTP client,
|
|
// ...) only touches this package.
|
|
//
|
|
// The client speaks the standard Chat Completions protocol with tool calling,
|
|
// so any OpenAI-compatible endpoint (OpenAI, DeepSeek, Moonshot, vLLM, Ollama,
|
|
// ...) works by configuring base_url + api_key + model in the admin console.
|
|
package llm
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
|
|
"github.com/openai/openai-go"
|
|
"github.com/openai/openai-go/option"
|
|
"github.com/openai/openai-go/shared"
|
|
)
|
|
|
|
// Config describes one model endpoint. Built from a store.ModelConfig row at
|
|
// call time; api_key is decrypted by the caller.
|
|
type Config struct {
|
|
BaseURL string
|
|
APIKey string
|
|
Model string
|
|
MaxTokens int // 0 = leave unset
|
|
}
|
|
|
|
// Role labels a message. Matches the OpenAI convention.
|
|
type Role string
|
|
|
|
const (
|
|
RoleSystem Role = "system"
|
|
RoleUser Role = "user"
|
|
RoleAssistant Role = "assistant"
|
|
RoleTool Role = "tool"
|
|
)
|
|
|
|
// Message is a single chat message. For assistant messages with tool calls,
|
|
// ToolCalls is set and Content may be empty. For role=tool, ToolCallID links
|
|
// the result back to the call it answers.
|
|
type Message struct {
|
|
Role Role `json:"role"`
|
|
Content string `json:"content,omitempty"`
|
|
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
|
|
ToolCallID string `json:"tool_call_id,omitempty"`
|
|
// Name is the tool name for role=tool; some providers expect it.
|
|
Name string `json:"name,omitempty"`
|
|
}
|
|
|
|
// ToolCall represents one tool invocation requested by the assistant.
|
|
type ToolCall struct {
|
|
ID string `json:"id"`
|
|
Name string `json:"name"`
|
|
Arguments map[string]any `json:"arguments"`
|
|
}
|
|
|
|
// Function describes one tool the model may call, in JSON Schema form.
|
|
type Function struct {
|
|
Name string `json:"name"`
|
|
Description string `json:"description"`
|
|
// Parameters is a raw JSON Schema document for the arguments.
|
|
Parameters []byte `json:"parameters"`
|
|
}
|
|
|
|
// Client is a stateless caller bound to one Config. Build one per request.
|
|
type Client struct {
|
|
cfg Config
|
|
oc openai.Client
|
|
}
|
|
|
|
// New builds a client for the given config.
|
|
func New(cfg Config) *Client {
|
|
opts := []option.RequestOption{
|
|
option.WithAPIKey(cfg.APIKey),
|
|
}
|
|
if cfg.BaseURL != "" {
|
|
opts = append(opts, option.WithBaseURL(cfg.BaseURL))
|
|
}
|
|
return &Client{cfg: cfg, oc: openai.NewClient(opts...)}
|
|
}
|
|
|
|
// StreamEvent is what RunStream yields. Exactly one field is non-zero per
|
|
// event. The caller routes them to the SSE layer (token / tool_calls / done).
|
|
type StreamEvent struct {
|
|
// Delta is an incremental assistant text token (content streaming).
|
|
Delta string
|
|
// ToolCalls is set when the assistant requests tool calls in this chunk.
|
|
// For a streaming response these accumulate until finish_reason == "tool_calls".
|
|
ToolCalls []ToolCall
|
|
// FinishReason is set on the terminal event ("stop", "tool_calls", ...).
|
|
FinishReason string
|
|
// Err carries a non-stream error (auth, network, parse).
|
|
Err error
|
|
}
|
|
|
|
// RunStream issues a streaming chat completion and invokes onEvent for each
|
|
// chunk. It returns the fully-assembled assistant message (content + any tool
|
|
// calls) so the caller can append it to history. The context controls cancel.
|
|
func (c *Client) RunStream(ctx context.Context, messages []Message, tools []Function, onEvent func(StreamEvent)) (Message, error) {
|
|
if onEvent == nil {
|
|
onEvent = func(StreamEvent) {}
|
|
}
|
|
params := openai.ChatCompletionNewParams{
|
|
Model: openai.ChatModel(c.cfg.Model),
|
|
Messages: toOpenAIMessages(messages),
|
|
Tools: toOpenAITools(tools),
|
|
}
|
|
if c.cfg.MaxTokens > 0 {
|
|
params.MaxTokens = openai.Int(int64(c.cfg.MaxTokens))
|
|
}
|
|
|
|
stream := c.oc.Chat.Completions.NewStreaming(ctx, params)
|
|
|
|
var (
|
|
contentBuf string
|
|
toolBuf []openai.ChatCompletionChunkChoiceDeltaToolCall
|
|
finish string
|
|
)
|
|
for stream.Next() {
|
|
chunk := stream.Current()
|
|
if len(chunk.Choices) == 0 {
|
|
continue
|
|
}
|
|
ch := chunk.Choices[0]
|
|
if ch.FinishReason != "" {
|
|
finish = ch.FinishReason
|
|
}
|
|
delta := ch.Delta
|
|
if delta.Content != "" {
|
|
contentBuf += delta.Content
|
|
onEvent(StreamEvent{Delta: delta.Content})
|
|
}
|
|
// Accumulate tool-call fragments by index. The SDK guarantees a stable
|
|
// index across chunks for the same call; arguments stream in pieces.
|
|
for _, tc := range delta.ToolCalls {
|
|
for len(toolBuf) <= int(tc.Index) {
|
|
toolBuf = append(toolBuf, openai.ChatCompletionChunkChoiceDeltaToolCall{})
|
|
}
|
|
slot := &toolBuf[tc.Index]
|
|
slot.Index = tc.Index
|
|
if tc.ID != "" {
|
|
slot.ID = tc.ID
|
|
}
|
|
if tc.Function.Name != "" {
|
|
slot.Function.Name = tc.Function.Name
|
|
}
|
|
slot.Function.Arguments += tc.Function.Arguments
|
|
}
|
|
}
|
|
if err := stream.Err(); err != nil {
|
|
// API errors arrive here as *openai.Error with a useful message.
|
|
onEvent(StreamEvent{Err: err})
|
|
return Message{}, fmt.Errorf("llm stream: %w", err)
|
|
}
|
|
if finish == "" {
|
|
finish = "stop"
|
|
}
|
|
|
|
// Build the assembled tool calls (if any) and surface a single tool_calls
|
|
// event so the caller can drive the execution loop.
|
|
var calls []ToolCall
|
|
for _, tc := range toolBuf {
|
|
if tc.ID == "" && tc.Function.Name == "" {
|
|
continue
|
|
}
|
|
args, _ := parseArguments(tc.Function.Arguments)
|
|
calls = append(calls, ToolCall{
|
|
ID: tc.ID,
|
|
Name: tc.Function.Name,
|
|
Arguments: args,
|
|
})
|
|
}
|
|
if len(calls) > 0 {
|
|
onEvent(StreamEvent{ToolCalls: calls, FinishReason: finish})
|
|
} else {
|
|
onEvent(StreamEvent{FinishReason: finish})
|
|
}
|
|
|
|
return Message{
|
|
Role: RoleAssistant,
|
|
Content: contentBuf,
|
|
ToolCalls: calls,
|
|
}, nil
|
|
}
|
|
|
|
// toOpenAIMessages converts our Message slice into the SDK's union type.
|
|
func toOpenAIMessages(ms []Message) []openai.ChatCompletionMessageParamUnion {
|
|
out := make([]openai.ChatCompletionMessageParamUnion, 0, len(ms))
|
|
for _, m := range ms {
|
|
switch m.Role {
|
|
case RoleSystem:
|
|
out = append(out, openai.SystemMessage(m.Content))
|
|
case RoleUser:
|
|
out = append(out, openai.UserMessage(m.Content))
|
|
case RoleAssistant:
|
|
msg := &openai.ChatCompletionAssistantMessageParam{}
|
|
msg.Content.OfString = openai.String(m.Content)
|
|
for _, tc := range m.ToolCalls {
|
|
args := marshalArguments(tc.Arguments)
|
|
msg.ToolCalls = append(msg.ToolCalls, openai.ChatCompletionMessageToolCallParam{
|
|
ID: tc.ID,
|
|
Function: openai.ChatCompletionMessageToolCallFunctionParam{
|
|
Name: tc.Name,
|
|
Arguments: args,
|
|
},
|
|
})
|
|
}
|
|
out = append(out, openai.ChatCompletionMessageParamUnion{
|
|
OfAssistant: msg,
|
|
})
|
|
case RoleTool:
|
|
out = append(out, openai.ToolMessage(m.Content, m.ToolCallID))
|
|
}
|
|
}
|
|
return out
|
|
}
|
|
|
|
// toOpenAITools converts our Function list into SDK tool params. The
|
|
// parameters JSON Schema is decoded into the SDK's map form.
|
|
func toOpenAITools(fs []Function) []openai.ChatCompletionToolParam {
|
|
if len(fs) == 0 {
|
|
return nil
|
|
}
|
|
out := make([]openai.ChatCompletionToolParam, 0, len(fs))
|
|
for _, f := range fs {
|
|
params := f.Parameters
|
|
if len(params) == 0 {
|
|
params = []byte(`{"type":"object","properties":{}}`)
|
|
}
|
|
var schemaMap map[string]any
|
|
// A malformed schema is non-fatal: fall back to an empty object so a
|
|
// bad skill file never crashes a chat request.
|
|
if err := json.Unmarshal(params, &schemaMap); err != nil || schemaMap == nil {
|
|
schemaMap = map[string]any{"type": "object", "properties": map[string]any{}}
|
|
}
|
|
fd := shared.FunctionDefinitionParam{
|
|
Name: f.Name,
|
|
Parameters: shared.FunctionParameters(schemaMap),
|
|
}
|
|
if f.Description != "" {
|
|
fd.Description = openai.String(f.Description)
|
|
}
|
|
out = append(out, openai.ChatCompletionToolParam{Function: fd})
|
|
}
|
|
return out
|
|
}
|
|
|
|
// ErrNoResponse is returned when the stream produced no usable content and no
|
|
// tool calls (e.g. content filter). The caller can map this to a user-facing
|
|
// "无法生成回复" message.
|
|
var ErrNoResponse = errors.New("model returned no content")
|