320 lines
8.0 KiB
Go
320 lines
8.0 KiB
Go
package ai
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
|
|
"ai-agent-scaffold-go/internal/domain/agent/ports"
|
|
)
|
|
|
|
type OpenAIClient struct {
|
|
httpClient *http.Client
|
|
completionsURL string
|
|
apiKey string
|
|
model string
|
|
}
|
|
|
|
type OpenAIToolDef struct {
|
|
Name string
|
|
Description string
|
|
}
|
|
|
|
func NewOpenAIClient(completionsURL, apiKey, model string, requestTimeout time.Duration) *OpenAIClient {
|
|
if requestTimeout <= 0 {
|
|
requestTimeout = 5 * time.Minute
|
|
}
|
|
return &OpenAIClient{
|
|
httpClient: &http.Client{Timeout: requestTimeout},
|
|
completionsURL: completionsURL,
|
|
apiKey: apiKey,
|
|
model: model,
|
|
}
|
|
}
|
|
|
|
func (c *OpenAIClient) Generate(ctx context.Context, messages []ports.ChatMessage, tools []OpenAIToolDef) (ports.ChatReply, error) {
|
|
body, err := buildRequestBody(c.model, messages, tools, false)
|
|
if err != nil {
|
|
return ports.ChatReply{}, err
|
|
}
|
|
resp, err := c.do(ctx, body)
|
|
if err != nil {
|
|
return ports.ChatReply{}, err
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
raw, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return ports.ChatReply{}, fmt.Errorf("openai read body: %w", err)
|
|
}
|
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
|
return ports.ChatReply{}, fmt.Errorf("openai upstream %d: %s", resp.StatusCode, truncate(string(raw), 400))
|
|
}
|
|
|
|
var parsed openaiCompletion
|
|
if err := json.Unmarshal(raw, &parsed); err != nil {
|
|
return ports.ChatReply{}, fmt.Errorf("openai decode: %w", err)
|
|
}
|
|
if len(parsed.Choices) == 0 {
|
|
return ports.ChatReply{}, fmt.Errorf("openai response has no choices")
|
|
}
|
|
choice := parsed.Choices[0].Message
|
|
reply := ports.ChatReply{Content: choice.Content}
|
|
for _, tc := range choice.ToolCalls {
|
|
reply.ToolCalls = append(reply.ToolCalls, ports.ChatToolCall{
|
|
ID: tc.ID,
|
|
Name: tc.Function.Name,
|
|
Arguments: tc.Function.Arguments,
|
|
})
|
|
}
|
|
return reply, nil
|
|
}
|
|
|
|
func (c *OpenAIClient) Stream(ctx context.Context, messages []ports.ChatMessage, tools []OpenAIToolDef) (<-chan ports.ChatStreamEvent, <-chan error) {
|
|
events := make(chan ports.ChatStreamEvent, 8)
|
|
errs := make(chan error, 1)
|
|
go func() {
|
|
defer close(events)
|
|
defer close(errs)
|
|
|
|
body, err := buildRequestBody(c.model, messages, tools, true)
|
|
if err != nil {
|
|
errs <- err
|
|
return
|
|
}
|
|
resp, err := c.do(ctx, body)
|
|
if err != nil {
|
|
errs <- err
|
|
return
|
|
}
|
|
defer resp.Body.Close()
|
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
|
raw, _ := io.ReadAll(resp.Body)
|
|
errs <- fmt.Errorf("openai upstream %d: %s", resp.StatusCode, truncate(string(raw), 400))
|
|
return
|
|
}
|
|
|
|
toolCallBuf := map[int]*ports.ChatToolCall{}
|
|
reader := bufio.NewReader(resp.Body)
|
|
for {
|
|
line, err := reader.ReadString('\n')
|
|
if err != nil {
|
|
if err == io.EOF {
|
|
emitToolCalls(events, toolCallBuf)
|
|
events <- ports.ChatStreamEvent{Done: true}
|
|
return
|
|
}
|
|
errs <- fmt.Errorf("openai stream read: %w", err)
|
|
return
|
|
}
|
|
line = strings.TrimRight(line, "\r\n")
|
|
if line == "" || !strings.HasPrefix(line, "data:") {
|
|
continue
|
|
}
|
|
payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
|
|
if payload == "[DONE]" {
|
|
emitToolCalls(events, toolCallBuf)
|
|
events <- ports.ChatStreamEvent{Done: true}
|
|
return
|
|
}
|
|
var chunk openaiStreamChunk
|
|
if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
|
|
errs <- fmt.Errorf("openai stream decode: %w", err)
|
|
return
|
|
}
|
|
if len(chunk.Choices) == 0 {
|
|
continue
|
|
}
|
|
delta := chunk.Choices[0].Delta
|
|
if delta.Content != "" {
|
|
select {
|
|
case events <- ports.ChatStreamEvent{Delta: delta.Content}:
|
|
case <-ctx.Done():
|
|
errs <- ctx.Err()
|
|
return
|
|
}
|
|
}
|
|
for _, tc := range delta.ToolCalls {
|
|
current, ok := toolCallBuf[tc.Index]
|
|
if !ok {
|
|
current = &ports.ChatToolCall{}
|
|
toolCallBuf[tc.Index] = current
|
|
}
|
|
if tc.ID != "" {
|
|
current.ID = tc.ID
|
|
}
|
|
if tc.Function.Name != "" {
|
|
current.Name = tc.Function.Name
|
|
}
|
|
if tc.Function.Arguments != "" {
|
|
current.Arguments += tc.Function.Arguments
|
|
}
|
|
}
|
|
}
|
|
}()
|
|
return events, errs
|
|
}
|
|
|
|
func emitToolCalls(events chan<- ports.ChatStreamEvent, buf map[int]*ports.ChatToolCall) {
|
|
if len(buf) == 0 {
|
|
return
|
|
}
|
|
calls := make([]ports.ChatToolCall, 0, len(buf))
|
|
for i := 0; i < len(buf); i++ {
|
|
if call, ok := buf[i]; ok && call != nil {
|
|
calls = append(calls, *call)
|
|
}
|
|
}
|
|
if len(calls) > 0 {
|
|
events <- ports.ChatStreamEvent{ToolCalls: calls}
|
|
}
|
|
}
|
|
|
|
func (c *OpenAIClient) do(ctx context.Context, body []byte) (*http.Response, error) {
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.completionsURL, bytes.NewReader(body))
|
|
if err != nil {
|
|
return nil, fmt.Errorf("openai build request: %w", err)
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("Accept", "application/json")
|
|
if c.apiKey != "" {
|
|
req.Header.Set("Authorization", "Bearer "+c.apiKey)
|
|
}
|
|
resp, err := c.httpClient.Do(req)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("openai http: %w", err)
|
|
}
|
|
return resp, nil
|
|
}
|
|
|
|
func buildRequestBody(modelName string, messages []ports.ChatMessage, tools []OpenAIToolDef, stream bool) ([]byte, error) {
|
|
payload := map[string]any{
|
|
"model": modelName,
|
|
"messages": encodeMessages(messages),
|
|
"stream": stream,
|
|
}
|
|
if len(tools) > 0 {
|
|
payload["tools"] = encodeTools(tools)
|
|
}
|
|
raw, err := json.Marshal(payload)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("openai marshal request: %w", err)
|
|
}
|
|
return raw, nil
|
|
}
|
|
|
|
func encodeMessages(messages []ports.ChatMessage) []map[string]any {
|
|
encoded := make([]map[string]any, 0, len(messages))
|
|
for _, m := range messages {
|
|
entry := map[string]any{"role": string(m.Role)}
|
|
if m.Content != "" {
|
|
entry["content"] = m.Content
|
|
} else if m.Role != ports.ChatRoleAssistant || len(m.ToolCalls) == 0 {
|
|
entry["content"] = ""
|
|
}
|
|
if m.Name != "" {
|
|
entry["name"] = m.Name
|
|
}
|
|
if m.ToolCallID != "" {
|
|
entry["tool_call_id"] = m.ToolCallID
|
|
}
|
|
if len(m.ToolCalls) > 0 {
|
|
calls := make([]map[string]any, 0, len(m.ToolCalls))
|
|
for _, c := range m.ToolCalls {
|
|
calls = append(calls, map[string]any{
|
|
"id": c.ID,
|
|
"type": "function",
|
|
"function": map[string]any{
|
|
"name": c.Name,
|
|
"arguments": c.Arguments,
|
|
},
|
|
})
|
|
}
|
|
entry["tool_calls"] = calls
|
|
}
|
|
encoded = append(encoded, entry)
|
|
}
|
|
return encoded
|
|
}
|
|
|
|
func encodeTools(tools []OpenAIToolDef) []map[string]any {
|
|
out := make([]map[string]any, 0, len(tools))
|
|
for _, t := range tools {
|
|
out = append(out, map[string]any{
|
|
"type": "function",
|
|
"function": map[string]any{
|
|
"name": t.Name,
|
|
"description": fallbackDescription(t),
|
|
"parameters": map[string]any{
|
|
"type": "object",
|
|
"properties": map[string]any{
|
|
"query": map[string]any{
|
|
"type": "string",
|
|
"description": "自由文本输入或检索查询，工具会按其语义解释",
|
|
},
|
|
},
|
|
},
|
|
},
|
|
})
|
|
}
|
|
return out
|
|
}
|
|
|
|
func fallbackDescription(t OpenAIToolDef) string {
|
|
if strings.TrimSpace(t.Description) != "" {
|
|
return t.Description
|
|
}
|
|
return "外部工具 " + t.Name + "，参数 query 为自由文本"
|
|
}
|
|
|
|
func truncate(s string, max int) string {
|
|
if len(s) <= max {
|
|
return s
|
|
}
|
|
return s[:max] + "..."
|
|
}
|
|
|
|
type openaiCompletion struct {
|
|
Choices []struct {
|
|
Message struct {
|
|
Content string `json:"content"`
|
|
ToolCalls []openaiToolCallV1 `json:"tool_calls"`
|
|
} `json:"message"`
|
|
} `json:"choices"`
|
|
}
|
|
|
|
type openaiStreamChunk struct {
|
|
Choices []struct {
|
|
Delta struct {
|
|
Content string `json:"content"`
|
|
ToolCalls []openaiStreamToolCall `json:"tool_calls"`
|
|
} `json:"delta"`
|
|
FinishReason string `json:"finish_reason"`
|
|
} `json:"choices"`
|
|
}
|
|
|
|
type openaiToolCallV1 struct {
|
|
ID string `json:"id"`
|
|
Type string `json:"type"`
|
|
Function struct {
|
|
Name string `json:"name"`
|
|
Arguments string `json:"arguments"`
|
|
} `json:"function"`
|
|
}
|
|
|
|
type openaiStreamToolCall struct {
|
|
Index int `json:"index"`
|
|
ID string `json:"id"`
|
|
Type string `json:"type"`
|
|
Function struct {
|
|
Name string `json:"name"`
|
|
Arguments string `json:"arguments"`
|
|
} `json:"function"`
|
|
}
|