WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes
Anydoc / anydoc parser engine (push) Canceled after 0s
App / Format, vet, test, and build (push) Canceled after 0s
CLI / test (macos-latest / go 1.26) (push) Canceled after 0s
CLI / test (ubuntu-latest / go 1.26) (push) Canceled after 0s
CLI / test (windows-latest / go 1.26) (push) Canceled after 0s
Build and Push Docker Image / build-ui (push) Canceled after 0s
Build and Push Docker Image / build-docreader (push) Canceled after 0s
Build and Push Docker Image / build-sandbox (push) Canceled after 0s
Build and Push Docker Image / build-app (amd64, linux/amd64, ubuntu-latest) (push) Canceled after 0s
Build and Push Docker Image / build-app (arm64, linux/arm64, ubuntu-24.04-arm) (push) Canceled after 0s
DocReader / Python tests and Go client tests (push) Canceled after 0s
Build & Publish (dsh-weknora) / Typecheck, test, pack (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (pinned) (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (latest, advisory) (push) Canceled after 0s
Frontend / Test, type-check, and build (push) Canceled after 0s
Go Lint Cache / warm lint cache (.) (push) Canceled after 0s
Go Lint Cache / warm lint cache (cli) (push) Canceled after 0s
Go Lint Cache / warm lint cache (client) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.10) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.11) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.12) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.13) (push) Canceled after 0s
Build and Push Docker Image / merge (push) Canceled after 0s
Build & Publish (dsh-weknora) / Publish to npm (push) Canceled after 0s
Build & Publish (mcp-server) / Build sdist + wheel (push) Canceled after 0s
Build & Publish (mcp-server) / Publish to PyPI (push) Canceled after 0s

This commit is contained in:
yawan
2026-09-22 16:19:06 +08:00
commit f7fc992f57
3929 changed files with 1006607 additions and 0 deletions
+65
View File
@@ -0,0 +1,65 @@
package asr
import (
"context"
"github.com/Tencent/WeKnora/internal/types"
)
// Segment represents a transcribed segment with timestamps.
type Segment struct {
Start float64 `json:"start"`
End float64 `json:"end"`
Text string `json:"text"`
}
// TranscriptionResult holds the full text and its segments.
type TranscriptionResult struct {
Text string `json:"text"`
Segments []Segment `json:"segments,omitempty"`
}
// ASR defines the interface for Automatic Speech Recognition model operations.
type ASR interface {
// Transcribe sends audio bytes to the ASR model and returns the transcribed text and segments.
Transcribe(ctx context.Context, audioBytes []byte, fileName string) (*TranscriptionResult, error)
GetModelName() string
GetModelID() string
}
// Config holds the configuration needed to create an ASR instance.
type Config struct {
Source types.ModelSource
BaseURL string
ModelName string
APIKey string
ModelID string
Language string // optional: specify language for transcription
// CustomHeaders 允许在调用远程 API 时附加自定义 HTTP 请求头(类似 OpenAI Python SDK 的 extra_headers)。
CustomHeaders map[string]string
}
// ConfigFromModel 根据 types.Model 构造 asr.Config。
// 生产路径(从 DB 拉起)和测试连接路径(临时表单)共享这份映射。
// 当前 ASR 不涉及 WeKnoraCloud 凭证,所以签名不含 appID/appSecret。
func ConfigFromModel(m *types.Model) *Config {
if m == nil {
return nil
}
return &Config{
ModelID: m.ID,
APIKey: m.Parameters.APIKey,
BaseURL: m.Parameters.BaseURL,
ModelName: m.Name,
Source: m.Source,
CustomHeaders: m.Parameters.CustomHeaders,
}
}
// NewASR creates an ASR instance based on the provided configuration.
// All ASR vendors use the OpenAI-compatible /v1/audio/transcriptions API.
func NewASR(config *Config) (ASR, error) {
a, err := NewOpenAIASR(config)
return wrapASRLangfuse(a, err)
}
@@ -0,0 +1,36 @@
package asr
import (
"testing"
"github.com/Tencent/WeKnora/internal/types"
)
func TestConfigFromModel(t *testing.T) {
m := &types.Model{
ID: "asr-1",
Name: "whisper-1",
Source: types.ModelSourceRemote,
Parameters: types.ModelParameters{
BaseURL: "https://api.example.com/v1",
APIKey: "sk",
CustomHeaders: map[string]string{"X": "y"},
},
}
cfg := ConfigFromModel(m)
if cfg == nil || cfg.ModelID != "asr-1" || cfg.ModelName != "whisper-1" {
t.Fatalf("identity mismatch: %+v", cfg)
}
if cfg.BaseURL != "https://api.example.com/v1" || cfg.APIKey != "sk" {
t.Errorf("connection fields mismatch: %+v", cfg)
}
if cfg.CustomHeaders["X"] != "y" {
t.Errorf("CustomHeaders not propagated: %+v", cfg.CustomHeaders)
}
}
func TestConfigFromModel_Nil(t *testing.T) {
if got := ConfigFromModel(nil); got != nil {
t.Fatalf("expected nil, got %+v", got)
}
}
+79
View File
@@ -0,0 +1,79 @@
package asr
import (
"context"
"github.com/Tencent/WeKnora/internal/tracing/langfuse"
)
// langfuseASR wraps an ASR implementation and reports each Transcribe call
// as a Langfuse generation observation. Audio bytes are not uploaded — we
// record file name, audio size and duration (derived from segment end times)
// so usage can be billed by second/minute, which is how most ASR providers
// price their services.
type langfuseASR struct {
inner ASR
}
func (l *langfuseASR) GetModelName() string { return l.inner.GetModelName() }
func (l *langfuseASR) GetModelID() string { return l.inner.GetModelID() }
func (l *langfuseASR) Transcribe(ctx context.Context, audioBytes []byte, fileName string) (*TranscriptionResult, error) {
mgr := langfuse.GetManager()
if !mgr.Enabled() {
return l.inner.Transcribe(ctx, audioBytes, fileName)
}
genCtx, gen := mgr.StartGeneration(ctx, langfuse.GenerationOptions{
Name: "asr.transcribe",
Model: l.inner.GetModelName(),
Input: map[string]interface{}{
"file_name": fileName,
"audio_size": len(audioBytes),
},
Metadata: map[string]interface{}{
"model_id": l.inner.GetModelID(),
"audio_size": len(audioBytes),
},
})
result, err := l.inner.Transcribe(genCtx, audioBytes, fileName)
output := map[string]interface{}{}
var duration float64
if result != nil {
output["text"] = result.Text
output["segment_count"] = len(result.Segments)
if n := len(result.Segments); n > 0 {
duration = result.Segments[n-1].End
output["duration_seconds"] = duration
}
}
// ASR is billed per second; we emit the audio duration (when available)
// in the "Output" side of usage so users can configure per-model pricing
// in Langfuse (e.g. $0.006 per minute for Whisper-1).
var usage *langfuse.TokenUsage
if duration > 0 {
seconds := int(duration + 0.5)
usage = &langfuse.TokenUsage{
Output: seconds,
Total: seconds,
Unit: "SECONDS",
}
}
gen.Finish(output, usage, err)
return result, err
}
// wrapASRLangfuse applies the Langfuse decorator when the manager is enabled.
func wrapASRLangfuse(a ASR, err error) (ASR, error) {
if err != nil || a == nil {
return a, err
}
if !langfuse.GetManager().Enabled() {
return a, nil
}
return &langfuseASR{inner: a}, nil
}
+150
View File
@@ -0,0 +1,150 @@
package asr
import (
"bytes"
"context"
"encoding/json"
"fmt"
"path/filepath"
"strings"
"time"
"github.com/Tencent/WeKnora/internal/logger"
secutils "github.com/Tencent/WeKnora/internal/utils"
openai "github.com/sashabaranov/go-openai"
)
const (
asrDefaultTimeout = 300 * time.Second // audio transcription can be slow
)
// OpenAIASR implements ASR via an OpenAI-compatible audio transcriptions API.
type OpenAIASR struct {
modelName string
modelID string
client *openai.Client
baseURL string
language string
}
// NewOpenAIASR creates an OpenAI-compatible ASR instance.
func NewOpenAIASR(config *Config) (*OpenAIASR, error) {
if err := validateASRBaseURL(config.BaseURL); err != nil {
return nil, err
}
apiCfg := openai.DefaultConfig(config.APIKey)
if config.BaseURL != "" {
apiCfg.BaseURL = config.BaseURL
}
httpClient := newASRHTTPClient(asrDefaultTimeout)
// 注入用户自定义 HTTP header(类似 OpenAI Python SDK 的 extra_headers)
if len(config.CustomHeaders) > 0 {
apiCfg.HTTPClient = secutils.WrapHTTPClientWithHeaders(httpClient, config.CustomHeaders)
} else {
apiCfg.HTTPClient = httpClient
}
return &OpenAIASR{
modelName: config.ModelName,
modelID: config.ModelID,
client: openai.NewClientWithConfig(apiCfg),
baseURL: config.BaseURL,
language: config.Language,
}, nil
}
// Transcribe sends audio bytes to the OpenAI-compatible audio transcriptions API.
func (s *OpenAIASR) Transcribe(ctx context.Context, audioBytes []byte, fileName string) (*TranscriptionResult, error) {
if len(audioBytes) == 0 {
return nil, fmt.Errorf("audio bytes are empty")
}
// Ensure fileName has an extension for proper MIME type detection
if fileName == "" {
fileName = "audio.mp3"
}
logger.Infof(ctx, "[ASR] Calling OpenAI-compatible transcription API, model=%s, baseURL=%s, audioSize=%d, file=%s",
s.modelName, s.baseURL, len(audioBytes), fileName)
req := openai.AudioRequest{
Model: s.modelName,
FilePath: fileName,
Reader: bytes.NewReader(audioBytes),
Format: openai.AudioResponseFormatVerboseJSON,
}
if s.language != "" {
req.Language = s.language
}
resp, err := s.client.CreateTranscription(ctx, req)
if err != nil {
return nil, fmt.Errorf("ASR transcription request failed: %w", err)
}
jsonData, err := json.Marshal(resp)
if err != nil {
return nil, fmt.Errorf("marshal transcription response: %w", err)
}
logger.Debugf(ctx, "[ASR] Transcription response: %s", string(jsonData))
text := strings.TrimSpace(resp.Text)
logger.Infof(ctx, "[ASR] Transcription completed, text length=%d", len(text))
var segments []Segment
for _, s := range resp.Segments {
segments = append(segments, Segment{
Start: s.Start,
End: s.End,
Text: strings.TrimSpace(s.Text),
})
}
return &TranscriptionResult{
Text: text,
Segments: segments,
}, nil
}
func (s *OpenAIASR) GetModelName() string { return s.modelName }
func (s *OpenAIASR) GetModelID() string { return s.modelID }
// DetectAudioFormat returns a file extension hint for the given audio bytes.
func DetectAudioFormat(data []byte, fileName string) string {
if fileName != "" {
ext := strings.ToLower(filepath.Ext(fileName))
if ext != "" {
return ext
}
}
// Try to detect from magic bytes
if len(data) >= 4 {
// MP3: starts with ID3 or 0xFF 0xFB
if data[0] == 'I' && data[1] == 'D' && data[2] == '3' {
return ".mp3"
}
if data[0] == 0xFF && (data[1]&0xE0) == 0xE0 {
return ".mp3"
}
// FLAC: starts with "fLaC"
if data[0] == 'f' && data[1] == 'L' && data[2] == 'a' && data[3] == 'C' {
return ".flac"
}
// OGG: starts with "OggS"
if data[0] == 'O' && data[1] == 'g' && data[2] == 'g' && data[3] == 'S' {
return ".ogg"
}
// WAV: starts with "RIFF"
if data[0] == 'R' && data[1] == 'I' && data[2] == 'F' && data[3] == 'F' {
return ".wav"
}
}
// M4A: check for ftyp box
if len(data) >= 8 && data[4] == 'f' && data[5] == 't' && data[6] == 'y' && data[7] == 'p' {
return ".m4a"
}
return ".mp3" // default
}
+25
View File
@@ -0,0 +1,25 @@
package asr
import (
"fmt"
"net/http"
"time"
secutils "github.com/Tencent/WeKnora/internal/utils"
)
func validateASRBaseURL(baseURL string) error {
if baseURL == "" {
return nil
}
if err := secutils.ValidateURLForSSRF(baseURL); err != nil {
return fmt.Errorf("base URL SSRF check failed: %w", err)
}
return nil
}
func newASRHTTPClient(timeout time.Duration) *http.Client {
cfg := secutils.DefaultSSRFSafeHTTPClientConfig()
cfg.Timeout = timeout
return secutils.NewSSRFSafeHTTPClient(cfg)
}
@@ -0,0 +1,26 @@
package asr
import (
"testing"
secutils "github.com/Tencent/WeKnora/internal/utils"
)
func withASRSSRFWhitelist(t *testing.T, raw string) {
t.Helper()
t.Setenv("SSRF_WHITELIST", raw)
secutils.ResetSSRFWhitelistForTest()
t.Cleanup(secutils.ResetSSRFWhitelistForTest)
}
func TestOpenAIASRRejectsInternalBaseURL(t *testing.T) {
withASRSSRFWhitelist(t, "")
_, err := NewOpenAIASR(&Config{
BaseURL: "http://169.254.169.254/latest/meta-data/",
ModelName: "asr-test",
})
if err == nil {
t.Fatalf("NewOpenAIASR returned nil error for blocked internal BaseURL")
}
}
+580
View File
@@ -0,0 +1,580 @@
package chat
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"net/url"
"strings"
"github.com/Tencent/WeKnora/internal/models/provider"
"github.com/Tencent/WeKnora/internal/models/requestbudget"
"github.com/Tencent/WeKnora/internal/types"
secutils "github.com/Tencent/WeKnora/internal/utils"
)
const anthropicVersion = "2023-06-01"
type AnthropicChat struct {
modelName string
modelID string
baseURL string
apiKey string
customHeaders map[string]string
}
type anthropicCacheControl struct {
Type string `json:"type"`
TTL string `json:"ttl,omitempty"`
}
type anthropicContentBlock struct {
Type string `json:"type"`
Text string `json:"text,omitempty"`
CacheControl *anthropicCacheControl `json:"cache_control,omitempty"`
}
type anthropicMessage struct {
Role string `json:"role"`
Content any `json:"content"`
}
type anthropicRequest struct {
Model string `json:"model"`
MaxTokens int `json:"max_tokens"`
Stream bool `json:"stream,omitempty"`
System any `json:"system,omitempty"`
Messages []anthropicMessage `json:"messages"`
Temperature *float64 `json:"temperature,omitempty"`
TopP *float64 `json:"top_p,omitempty"`
}
type anthropicResponse struct {
ID string `json:"id"`
Type string `json:"type"`
Role string `json:"role"`
Content []struct {
Type string `json:"type"`
Text string `json:"text"`
} `json:"content"`
StopReason string `json:"stop_reason"`
Usage struct {
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
CacheCreationInputTokens *int `json:"cache_creation_input_tokens"`
CacheReadInputTokens *int `json:"cache_read_input_tokens"`
} `json:"usage"`
Error *struct {
Type string `json:"type"`
Message string `json:"message"`
} `json:"error,omitempty"`
}
type anthropicStreamEvent struct {
Type string `json:"type"`
Message *struct {
Usage struct {
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
CacheCreationInputTokens *int `json:"cache_creation_input_tokens"`
CacheReadInputTokens *int `json:"cache_read_input_tokens"`
} `json:"usage"`
} `json:"message,omitempty"`
Delta *struct {
Type string `json:"type"`
Text string `json:"text"`
StopReason string `json:"stop_reason"`
} `json:"delta,omitempty"`
Usage *struct {
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
CacheCreationInputTokens *int `json:"cache_creation_input_tokens"`
CacheReadInputTokens *int `json:"cache_read_input_tokens"`
} `json:"usage,omitempty"`
Error *struct {
Type string `json:"type"`
Message string `json:"message"`
} `json:"error,omitempty"`
}
func NewAnthropicChat(config *ChatConfig) (*AnthropicChat, error) {
if config.BaseURL != "" {
if err := secutils.ValidateURLForSSRF(config.BaseURL); err != nil {
return nil, fmt.Errorf("baseURL SSRF check failed: %w", err)
}
}
if strings.TrimSpace(config.APIKey) == "" {
return nil, fmt.Errorf("Anthropic provider: API key is required")
}
baseURL := strings.TrimRight(config.BaseURL, "/")
if baseURL == "" {
baseURL = provider.AnthropicBaseURL
}
return &AnthropicChat{
modelName: config.ModelName,
modelID: config.ModelID,
baseURL: baseURL,
apiKey: config.APIKey,
customHeaders: config.CustomHeaders,
}, nil
}
func (c *AnthropicChat) Chat(ctx context.Context, messages []Message, opts *ChatOptions) (*types.ChatResponse, error) {
ctx = requestbudget.WithModel(ctx, c.modelID)
if err := requestbudget.Check(ctx); err != nil {
return nil, err
}
reqBody := c.buildRequest(ctx, messages, opts)
jsonData, err := json.Marshal(reqBody)
if err != nil {
return nil, fmt.Errorf("marshal request: %w", err)
}
ctx, cancel := withLLMTimeout(ctx, defaultChatTimeout)
defer cancel()
endpoint := c.endpoint()
if err := secutils.ValidateURLForSSRF(endpoint); err != nil {
return nil, fmt.Errorf("endpoint SSRF check failed: %w", err)
}
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewBuffer(jsonData))
if err != nil {
return nil, fmt.Errorf("create request: %w", err)
}
httpReq.Header.Set("Content-Type", "application/json")
httpReq.Header.Set("x-api-key", c.apiKey)
httpReq.Header.Set("anthropic-version", anthropicVersion)
secutils.ApplyCustomHeaders(httpReq, c.customHeaders)
resp, err := rawHTTPClient.Do(httpReq)
if err != nil {
return nil, fmt.Errorf("send request: %w", err)
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("read response: %w", err)
}
if strings.Contains(strings.ToLower(resp.Header.Get("Content-Type")), "text/event-stream") {
chatResp, err := parseAnthropicSSE(bytes.NewReader(body))
if err != nil {
return nil, err
}
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, chatResp.Content)
}
logUsage(ctx, c.modelName, &chatResp.Usage)
return chatResp, nil
}
var chatResp anthropicResponse
if err := json.Unmarshal(body, &chatResp); err != nil {
return nil, fmt.Errorf("decode response: %w", err)
}
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
if chatResp.Error != nil && chatResp.Error.Message != "" {
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, chatResp.Error.Message)
}
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
result := c.parseResponse(&chatResp)
logUsage(ctx, c.modelName, &result.Usage)
return result, nil
}
func (c *AnthropicChat) ChatStream(ctx context.Context, messages []Message, opts *ChatOptions) (<-chan types.StreamResponse, error) {
ctx = requestbudget.WithModel(ctx, c.modelID)
if err := requestbudget.Check(ctx); err != nil {
return nil, err
}
reqBody := c.buildRequest(ctx, messages, opts)
reqBody.Stream = true
jsonData, err := json.Marshal(reqBody)
if err != nil {
return nil, fmt.Errorf("marshal request: %w", err)
}
endpoint := c.endpoint()
if err := secutils.ValidateURLForSSRF(endpoint); err != nil {
return nil, fmt.Errorf("endpoint SSRF check failed: %w", err)
}
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewBuffer(jsonData))
if err != nil {
return nil, fmt.Errorf("create request: %w", err)
}
httpReq.Header.Set("Content-Type", "application/json")
httpReq.Header.Set("Accept", "text/event-stream")
httpReq.Header.Set("x-api-key", c.apiKey)
httpReq.Header.Set("anthropic-version", anthropicVersion)
secutils.ApplyCustomHeaders(httpReq, c.customHeaders)
resp, err := rawHTTPClient.Do(httpReq)
if err != nil {
return nil, fmt.Errorf("send request: %w", err)
}
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
body, _ := io.ReadAll(resp.Body)
resp.Body.Close()
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
streamChan := make(chan types.StreamResponse)
go processAnthropicStream(ctx, c.modelName, resp, streamChan)
return streamChan, nil
}
func (c *AnthropicChat) GetModelName() string {
return c.modelName
}
func (c *AnthropicChat) GetModelID() string {
return c.modelID
}
func (c *AnthropicChat) endpoint() string {
baseURL := strings.TrimRight(c.baseURL, "/")
if isAnthropicMessagesEndpoint(baseURL) {
return baseURL
}
if isAnthropicVersionedBaseURL(baseURL) {
return baseURL + "/messages"
}
return baseURL + "/v1/messages"
}
func isAnthropicMessagesEndpoint(baseURL string) bool {
u, err := url.Parse(baseURL)
if err != nil {
return false
}
path := strings.TrimRight(u.Path, "/")
return strings.HasSuffix(path, "/messages")
}
func isAnthropicVersionedBaseURL(baseURL string) bool {
u, err := url.Parse(baseURL)
if err != nil {
return false
}
path := strings.TrimRight(u.Path, "/")
return strings.HasSuffix(path, "/v1") || strings.HasSuffix(path, "/v1beta")
}
func (c *AnthropicChat) buildRequest(_ context.Context, messages []Message, opts *ChatOptions) anthropicRequest {
req := anthropicRequest{
Model: c.modelName,
MaxTokens: 1024,
Messages: make([]anthropicMessage, 0, len(messages)),
}
if opts != nil {
if opts.MaxTokens > 0 {
req.MaxTokens = opts.MaxTokens
} else if opts.MaxCompletionTokens > 0 {
req.MaxTokens = opts.MaxCompletionTokens
}
if opts.Temperature > 0 {
temperature := opts.Temperature
req.Temperature = &temperature
}
if opts.TopP > 0 {
topP := opts.TopP
req.TopP = &topP
}
}
var systemParts []string
for _, msg := range messages {
content := strings.TrimSpace(msg.Content)
if content == "" {
content = textFromMultiContent(msg.MultiContent)
}
if content == "" {
continue
}
switch msg.Role {
case "system":
systemParts = append(systemParts, content)
case "assistant":
req.Messages = append(req.Messages, anthropicMessage{Role: "assistant", Content: content})
case "user":
req.Messages = append(req.Messages, anthropicMessage{Role: "user", Content: content})
default:
req.Messages = append(req.Messages, anthropicMessage{Role: "user", Content: content})
}
}
systemText := strings.Join(systemParts, "\n\n")
retention := resolveCacheRetention(opts)
marker := cacheControlFor(retention, "1h")
if marker == nil {
req.System = systemText
return req
}
if systemText != "" {
req.System = []anthropicContentBlock{{
Type: "text",
Text: systemText,
CacheControl: &anthropicCacheControl{Type: marker.Type, TTL: marker.TTL},
}}
}
if len(req.Messages) > 0 {
last := &req.Messages[len(req.Messages)-1]
if text, ok := last.Content.(string); ok && text != "" {
last.Content = []anthropicContentBlock{{
Type: "text",
Text: text,
CacheControl: &anthropicCacheControl{Type: marker.Type, TTL: marker.TTL},
}}
}
}
return req
}
func textFromMultiContent(parts []MessageContentPart) string {
if len(parts) == 0 {
return ""
}
textParts := make([]string, 0, len(parts))
for _, part := range parts {
if part.Type == "text" && strings.TrimSpace(part.Text) != "" {
textParts = append(textParts, strings.TrimSpace(part.Text))
}
}
return strings.Join(textParts, "\n")
}
func (c *AnthropicChat) parseResponse(resp *anthropicResponse) *types.ChatResponse {
parts := make([]string, 0, len(resp.Content))
for _, part := range resp.Content {
if part.Type == "text" && part.Text != "" {
parts = append(parts, part.Text)
}
}
inputTokens := resp.Usage.InputTokens
outputTokens := resp.Usage.OutputTokens
cacheRead := valueOrZero(resp.Usage.CacheReadInputTokens)
cacheWrite := valueOrZero(resp.Usage.CacheCreationInputTokens)
promptTokens := inputTokens + cacheRead + cacheWrite
usage := types.TokenUsage{
PromptTokens: promptTokens,
CompletionTokens: outputTokens,
TotalTokens: promptTokens + outputTokens,
}
usage.SetPromptCacheUsage(cacheRead, cacheWrite, max(0, promptTokens-cacheRead),
resp.Usage.CacheReadInputTokens != nil || resp.Usage.CacheCreationInputTokens != nil)
return &types.ChatResponse{
Content: strings.Join(parts, ""),
FinishReason: resp.StopReason,
Usage: usage,
}
}
func parseAnthropicSSE(reader io.Reader) (*types.ChatResponse, error) {
sseReader := NewSSEReader(reader)
var contentParts []string
var finishReason string
var inputTokens int
var outputTokens int
var cacheReadTokens int
var cacheWriteTokens int
var cacheReported bool
for {
event, err := sseReader.ReadEvent()
if err == io.EOF {
break
}
if err != nil {
return nil, fmt.Errorf("read SSE response: %w", err)
}
if event.Done {
break
}
if len(event.Data) == 0 {
continue
}
var streamEvent anthropicStreamEvent
if err := json.Unmarshal(event.Data, &streamEvent); err != nil {
return nil, fmt.Errorf("decode SSE response: %w", err)
}
if streamEvent.Error != nil && streamEvent.Error.Message != "" {
return nil, fmt.Errorf("API stream error: %s", streamEvent.Error.Message)
}
if streamEvent.Message != nil {
inputTokens = max(inputTokens, streamEvent.Message.Usage.InputTokens)
outputTokens = max(outputTokens, streamEvent.Message.Usage.OutputTokens)
cacheReadTokens, cacheWriteTokens, cacheReported = mergeAnthropicCacheCounters(
cacheReadTokens, cacheWriteTokens, cacheReported,
streamEvent.Message.Usage.CacheReadInputTokens,
streamEvent.Message.Usage.CacheCreationInputTokens,
)
}
if streamEvent.Delta != nil {
if streamEvent.Delta.Type == "text_delta" && streamEvent.Delta.Text != "" {
contentParts = append(contentParts, streamEvent.Delta.Text)
}
if streamEvent.Delta.StopReason != "" {
finishReason = streamEvent.Delta.StopReason
}
}
if streamEvent.Usage != nil {
inputTokens = max(inputTokens, streamEvent.Usage.InputTokens)
outputTokens = max(outputTokens, streamEvent.Usage.OutputTokens)
cacheReadTokens, cacheWriteTokens, cacheReported = mergeAnthropicCacheCounters(
cacheReadTokens, cacheWriteTokens, cacheReported,
streamEvent.Usage.CacheReadInputTokens,
streamEvent.Usage.CacheCreationInputTokens,
)
}
}
promptTokens := inputTokens + cacheReadTokens + cacheWriteTokens
usage := types.TokenUsage{
PromptTokens: promptTokens,
CompletionTokens: outputTokens,
TotalTokens: promptTokens + outputTokens,
}
usage.SetPromptCacheUsage(cacheReadTokens, cacheWriteTokens,
max(0, promptTokens-cacheReadTokens), cacheReported)
return &types.ChatResponse{
Content: strings.Join(contentParts, ""),
FinishReason: finishReason,
Usage: usage,
}, nil
}
func processAnthropicStream(ctx context.Context, model string, resp *http.Response, streamChan chan types.StreamResponse) {
defer close(streamChan)
defer resp.Body.Close()
sseReader := NewSSEReader(resp.Body)
var usage *types.TokenUsage
var finishReason string
for {
event, err := sseReader.ReadEvent()
if err != nil {
if err == io.EOF {
logUsage(ctx, model, usage)
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeAnswer,
Content: "",
Done: true,
Usage: usage,
FinishReason: finishReason,
}
} else {
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeError,
Content: err.Error(),
Done: true,
}
}
return
}
if event.Done {
logUsage(ctx, model, usage)
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeAnswer,
Content: "",
Done: true,
Usage: usage,
FinishReason: finishReason,
}
return
}
if len(event.Data) == 0 {
continue
}
var streamEvent anthropicStreamEvent
if err := json.Unmarshal(event.Data, &streamEvent); err != nil {
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeError,
Content: fmt.Sprintf("decode SSE response: %v", err),
Done: true,
}
return
}
if streamEvent.Error != nil && streamEvent.Error.Message != "" {
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeError,
Content: streamEvent.Error.Message,
Done: true,
}
return
}
if streamEvent.Message != nil {
usage = mergeAnthropicUsage(usage, streamEvent.Message.Usage.InputTokens,
streamEvent.Message.Usage.OutputTokens,
streamEvent.Message.Usage.CacheReadInputTokens,
streamEvent.Message.Usage.CacheCreationInputTokens)
}
if streamEvent.Delta != nil {
if streamEvent.Delta.StopReason != "" {
finishReason = streamEvent.Delta.StopReason
}
if streamEvent.Delta.Type == "text_delta" && streamEvent.Delta.Text != "" {
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeAnswer,
Content: streamEvent.Delta.Text,
Done: false,
}
}
}
if streamEvent.Usage != nil {
usage = mergeAnthropicUsage(usage, streamEvent.Usage.InputTokens,
streamEvent.Usage.OutputTokens,
streamEvent.Usage.CacheReadInputTokens,
streamEvent.Usage.CacheCreationInputTokens)
}
}
}
func mergeAnthropicUsage(
current *types.TokenUsage,
inputTokens, outputTokens int,
cacheRead, cacheWrite *int,
) *types.TokenUsage {
if current == nil {
current = &types.TokenUsage{}
}
read, write, reported := mergeAnthropicCacheCounters(
current.CacheReadTokens, current.CacheWriteTokens, current.CacheReported,
cacheRead, cacheWrite,
)
uncachedInput := max(0, current.PromptTokens-current.CacheReadTokens-current.CacheWriteTokens)
uncachedInput = max(uncachedInput, inputTokens)
current.PromptTokens = uncachedInput + read + write
current.CompletionTokens = max(current.CompletionTokens, outputTokens)
current.TotalTokens = current.PromptTokens + current.CompletionTokens
current.SetPromptCacheUsage(read, write, max(0, current.PromptTokens-read), reported)
return current
}
func mergeAnthropicCacheCounters(
currentRead, currentWrite int,
currentReported bool,
cacheRead, cacheWrite *int,
) (read, write int, reported bool) {
read = currentRead
write = currentWrite
reported = currentReported || cacheRead != nil || cacheWrite != nil
if cacheRead != nil {
read = max(read, *cacheRead)
}
if cacheWrite != nil {
write = max(write, *cacheWrite)
}
return read, write, reported
}
+331
View File
@@ -0,0 +1,331 @@
package chat
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"github.com/Tencent/WeKnora/internal/models/provider"
"github.com/Tencent/WeKnora/internal/types"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestAnthropicChat(t *testing.T) {
t.Setenv("SSRF_WHITELIST", "127.0.0.1")
var capturedHeaders http.Header
var capturedRequest anthropicRequest
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
assert.Equal(t, "/v1/messages", r.URL.Path)
capturedHeaders = r.Header.Clone()
require.NoError(t, json.NewDecoder(r.Body).Decode(&capturedRequest))
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{
"id":"msg_123",
"type":"message",
"role":"assistant",
"content":[{"type":"text","text":"hello"}],
"stop_reason":"end_turn",
"usage":{"input_tokens":3,"output_tokens":2}
}`))
}))
defer server.Close()
chat, err := NewAnthropicChat(&ChatConfig{
Source: types.ModelSourceRemote,
BaseURL: server.URL,
ModelName: "claude-sonnet-4-5",
APIKey: "test-key",
Provider: string(provider.ProviderAnthropic),
CustomHeaders: map[string]string{
"anthropic-beta": "test-beta",
},
})
require.NoError(t, err)
resp, err := chat.Chat(context.Background(), []Message{
{Role: "system", Content: "You are helpful."},
{Role: "user", Content: "Hi"},
}, &ChatOptions{MaxTokens: 7, Temperature: 0.2})
require.NoError(t, err)
assert.Equal(t, "test-key", capturedHeaders.Get("x-api-key"))
assert.Equal(t, anthropicVersion, capturedHeaders.Get("anthropic-version"))
assert.Equal(t, "test-beta", capturedHeaders.Get("anthropic-beta"))
assert.Equal(t, "claude-sonnet-4-5", capturedRequest.Model)
assert.Equal(t, 7, capturedRequest.MaxTokens)
systemText, systemCache := firstTextAndCache(t, capturedRequest.System)
assert.Equal(t, "You are helpful.", systemText)
require.Equal(t, "ephemeral", systemCache["type"])
require.Len(t, capturedRequest.Messages, 1)
assert.Equal(t, "user", capturedRequest.Messages[0].Role)
userText, userCache := firstTextAndCache(t, capturedRequest.Messages[0].Content)
assert.Equal(t, "Hi", userText)
require.Equal(t, "ephemeral", userCache["type"])
assert.Equal(t, "hello", resp.Content)
assert.Equal(t, "end_turn", resp.FinishReason)
assert.Equal(t, 3, resp.Usage.PromptTokens)
assert.Equal(t, 2, resp.Usage.CompletionTokens)
assert.Equal(t, 5, resp.Usage.TotalTokens)
}
func TestAnthropicChat_CacheUsage(t *testing.T) {
t.Setenv("SSRF_WHITELIST", "127.0.0.1")
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{
"id":"msg_cached","type":"message","role":"assistant",
"content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn",
"usage":{"input_tokens":24,"output_tokens":2,"cache_creation_input_tokens":100,"cache_read_input_tokens":900}
}`))
}))
defer server.Close()
chat, err := NewAnthropicChat(&ChatConfig{
Source: types.ModelSourceRemote, BaseURL: server.URL, ModelName: "claude-sonnet-4-5",
APIKey: "test-key", Provider: string(provider.ProviderAnthropic),
})
require.NoError(t, err)
resp, err := chat.Chat(context.Background(), []Message{{Role: "user", Content: "Hi"}}, nil)
require.NoError(t, err)
assert.Equal(t, 1024, resp.Usage.PromptTokens)
assert.Equal(t, 900, resp.Usage.CacheReadTokens)
assert.Equal(t, 100, resp.Usage.CacheWriteTokens)
assert.Equal(t, 124, resp.Usage.CacheMissTokens)
assert.True(t, resp.Usage.CacheReported)
assert.Equal(t, types.PromptCacheStatusHit, resp.Usage.CacheStatus)
}
func TestAnthropicChat_FullEndpoint(t *testing.T) {
t.Setenv("SSRF_WHITELIST", "127.0.0.1")
var capturedPath string
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
capturedPath = r.URL.Path
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{
"id":"msg_123",
"type":"message",
"role":"assistant",
"content":[{"type":"text","text":"hello"}],
"stop_reason":"end_turn",
"usage":{"input_tokens":3,"output_tokens":2}
}`))
}))
defer server.Close()
chat, err := NewAnthropicChat(&ChatConfig{
Source: types.ModelSourceRemote,
BaseURL: server.URL + "/api/proxy/forward",
ModelName: "gpt-5.5",
APIKey: "test-key",
Provider: string(provider.ProviderAnthropic),
})
require.NoError(t, err)
resp, err := chat.Chat(context.Background(), []Message{{Role: "user", Content: "Hi"}}, nil)
require.NoError(t, err)
assert.Equal(t, "/api/proxy/forward/v1/messages", capturedPath)
assert.Equal(t, "hello", resp.Content)
}
func TestAnthropicChat_MessagesEndpoint(t *testing.T) {
t.Setenv("SSRF_WHITELIST", "127.0.0.1")
var capturedPath string
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
capturedPath = r.URL.Path
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{
"id":"msg_123",
"type":"message",
"role":"assistant",
"content":[{"type":"text","text":"hello"}],
"stop_reason":"end_turn",
"usage":{"input_tokens":3,"output_tokens":2}
}`))
}))
defer server.Close()
chat, err := NewAnthropicChat(&ChatConfig{
Source: types.ModelSourceRemote,
BaseURL: server.URL + "/api/v1/messages",
ModelName: "gpt-5.5",
APIKey: "test-key",
Provider: string(provider.ProviderAnthropic),
})
require.NoError(t, err)
resp, err := chat.Chat(context.Background(), []Message{{Role: "user", Content: "Hi"}}, nil)
require.NoError(t, err)
assert.Equal(t, "/api/v1/messages", capturedPath)
assert.Equal(t, "hello", resp.Content)
}
func TestAnthropicChat_SSEResponse(t *testing.T) {
t.Setenv("SSRF_WHITELIST", "127.0.0.1")
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
assert.Equal(t, "/v1/messages", r.URL.Path)
w.Header().Set("Content-Type", "text/event-stream")
_, _ = w.Write([]byte(`event: message_start
data: {"type":"message_start","message":{"usage":{"input_tokens":14,"output_tokens":0,"cache_creation_input_tokens":0,"cache_read_input_tokens":100}}}
event: content_block_delta
data: {"type":"content_block_delta","delta":{"type":"text_delta","text":"pong"}}
event: message_delta
data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":5}}
event: message_stop
data: {"type":"message_stop"}
`))
}))
defer server.Close()
chat, err := NewAnthropicChat(&ChatConfig{
Source: types.ModelSourceRemote,
BaseURL: server.URL,
ModelName: "gpt-5.5",
APIKey: "test-key",
Provider: string(provider.ProviderAnthropic),
})
require.NoError(t, err)
resp, err := chat.Chat(context.Background(), []Message{{Role: "user", Content: "ping"}}, nil)
require.NoError(t, err)
assert.Equal(t, "pong", resp.Content)
assert.Equal(t, "end_turn", resp.FinishReason)
assert.Equal(t, 114, resp.Usage.PromptTokens)
assert.Equal(t, 5, resp.Usage.CompletionTokens)
assert.Equal(t, 119, resp.Usage.TotalTokens)
assert.Equal(t, 100, resp.Usage.CacheReadTokens)
assert.Equal(t, 14, resp.Usage.CacheMissTokens)
}
func TestAnthropicChat_ChatStream(t *testing.T) {
t.Setenv("SSRF_WHITELIST", "127.0.0.1")
var capturedRequest anthropicRequest
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
assert.Equal(t, "/v1/messages", r.URL.Path)
assert.Equal(t, "text/event-stream", r.Header.Get("Accept"))
require.NoError(t, json.NewDecoder(r.Body).Decode(&capturedRequest))
w.Header().Set("Content-Type", "text/event-stream")
_, _ = w.Write([]byte(`event: message_start
data: {"type":"message_start","message":{"usage":{"input_tokens":14,"output_tokens":0,"cache_creation_input_tokens":0,"cache_read_input_tokens":100}}}
event: content_block_delta
data: {"type":"content_block_delta","delta":{"type":"text_delta","text":"pong"}}
event: message_delta
data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":5}}
event: message_stop
data: {"type":"message_stop"}
`))
}))
defer server.Close()
chat, err := NewAnthropicChat(&ChatConfig{
Source: types.ModelSourceRemote,
BaseURL: server.URL,
ModelName: "gpt-5.5",
APIKey: "test-key",
Provider: string(provider.ProviderAnthropic),
})
require.NoError(t, err)
ch, err := chat.ChatStream(context.Background(), []Message{{Role: "user", Content: "ping"}}, nil)
require.NoError(t, err)
var chunks []types.StreamResponse
for chunk := range ch {
chunks = append(chunks, chunk)
}
require.Len(t, chunks, 2)
assert.True(t, capturedRequest.Stream)
assert.Equal(t, "pong", chunks[0].Content)
assert.False(t, chunks[0].Done)
assert.True(t, chunks[1].Done)
assert.Equal(t, "end_turn", chunks[1].FinishReason)
require.NotNil(t, chunks[1].Usage)
assert.Equal(t, 114, chunks[1].Usage.PromptTokens)
assert.Equal(t, 5, chunks[1].Usage.CompletionTokens)
assert.Equal(t, 119, chunks[1].Usage.TotalTokens)
assert.Equal(t, 100, chunks[1].Usage.CacheReadTokens)
assert.Equal(t, 14, chunks[1].Usage.CacheMissTokens)
}
func TestNewRemoteChat_AnthropicProvider(t *testing.T) {
chat, err := NewRemoteChat(&ChatConfig{
Source: types.ModelSourceRemote,
ModelName: "claude-sonnet-4-5",
APIKey: "test-key",
Provider: string(provider.ProviderAnthropic),
})
require.NoError(t, err)
_, ok := chat.(*AnthropicChat)
assert.True(t, ok)
}
func TestAnthropicChat_CacheRetentionNoneKeepsPlainStrings(t *testing.T) {
t.Setenv("SSRF_WHITELIST", "127.0.0.1")
var capturedRequest anthropicRequest
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
require.NoError(t, json.NewDecoder(r.Body).Decode(&capturedRequest))
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{
"id":"msg_123","type":"message","role":"assistant",
"content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn",
"usage":{"input_tokens":3,"output_tokens":2}
}`))
}))
defer server.Close()
chat, err := NewAnthropicChat(&ChatConfig{
Source: types.ModelSourceRemote, BaseURL: server.URL, ModelName: "claude-sonnet-4-5",
APIKey: "test-key", Provider: string(provider.ProviderAnthropic),
})
require.NoError(t, err)
_, err = chat.Chat(context.Background(), []Message{
{Role: "system", Content: "You are helpful."},
{Role: "user", Content: "Hi"},
}, &ChatOptions{CacheRetention: CacheRetentionNone})
require.NoError(t, err)
assert.Equal(t, "You are helpful.", capturedRequest.System)
assert.Equal(t, "Hi", capturedRequest.Messages[0].Content)
}
func firstTextAndCache(t *testing.T, v any) (string, map[string]any) {
t.Helper()
switch x := v.(type) {
case string:
t.Fatalf("expected cache_control content blocks, got plain string %q", x)
return "", nil
case []any:
require.NotEmpty(t, x)
block, ok := x[0].(map[string]any)
require.True(t, ok)
text, _ := block["text"].(string)
cache, _ := block["cache_control"].(map[string]any)
require.NotNil(t, cache)
return text, cache
default:
t.Fatalf("unexpected content type %T", v)
return "", nil
}
}
+195
View File
@@ -0,0 +1,195 @@
package chat
import (
"context"
"encoding/json"
"fmt"
"strings"
"github.com/Tencent/WeKnora/internal/models/provider"
"github.com/Tencent/WeKnora/internal/models/utils/ollama"
"github.com/Tencent/WeKnora/internal/types"
)
// Tool represents a function/tool definition
type Tool struct {
Type string `json:"type"` // "function"
Function FunctionDef `json:"function"`
}
// FunctionDef represents a function definition
type FunctionDef struct {
Name string `json:"name"`
Description string `json:"description"`
Parameters json.RawMessage `json:"parameters"`
}
// ChatOptions 聊天选项
type ChatOptions struct {
Temperature float64 `json:"temperature"` // 温度参数
TopP float64 `json:"top_p"` // Top P 参数
Seed int `json:"seed"` // 随机种子
MaxTokens int `json:"max_tokens"` // 最大 token 数
MaxCompletionTokens int `json:"max_completion_tokens"` // 最大完成 token 数
FrequencyPenalty float64 `json:"frequency_penalty"` // 频率惩罚
PresencePenalty float64 `json:"presence_penalty"` // 存在惩罚
Thinking *bool `json:"thinking"` // 是否启用思考
Tools []Tool `json:"tools,omitempty"` // 可用工具列表
ToolChoice string `json:"tool_choice,omitempty"` // "auto", "required", "none", or specific tool
ParallelToolCalls *bool `json:"parallel_tool_calls,omitempty"` // 是否允许并行工具调用(默认 nil 表示由模型决定)
Format json.RawMessage `json:"format,omitempty"` // 响应格式定义
// PromptCacheKey is the provider routing key (OpenAI prompt_cache_key).
// Empty falls back to the session ID on the call context.
PromptCacheKey string `json:"-"`
// CacheRetention controls provider prompt-cache TTL. none disables cache
// markers; empty/short is the default 5-minute cache; long requests 1h/24h
// where the provider accepts it.
CacheRetention CacheRetention `json:"-"`
}
// MessageContentPart represents a part of multi-content message
type MessageContentPart struct {
Type string `json:"type"` // "text" or "image_url"
Text string `json:"text,omitempty"` // For type="text"
ImageURL *ImageURL `json:"image_url,omitempty"` // For type="image_url"
}
// ImageURL represents the image URL structure
type ImageURL struct {
URL string `json:"url"` // URL or base64 data URI
Detail string `json:"detail,omitempty"` // "auto", "low", "high"
}
// MessageKind marks messages the engine synthesized rather than received from
// the user or the model. Compaction needs to tell its own summary apart from a
// real user turn: a summary that looks like ordinary history gets fed back into
// the next summarization pass and degrades into a summary of a summary.
type MessageKind string
// MessageKindCompactionSummary marks the message that replaces compacted
// history. It carries the `user` role because that is where providers expect
// conversation history, so the role alone cannot identify it.
const MessageKindCompactionSummary MessageKind = "compaction_summary"
// Message 表示聊天消息
type Message struct {
Role string `json:"role"` // 角色:system, user, assistant, tool
Content string `json:"content"` // 消息内容
MultiContent []MessageContentPart `json:"multi_content,omitempty"` // 多内容消息(文本+图片)
Name string `json:"name,omitempty"` // Function/tool name (for tool role)
ToolCallID string `json:"tool_call_id,omitempty"` // Tool call ID (for tool role)
ToolCalls []ToolCall `json:"tool_calls,omitempty"` // Tool calls (for assistant role)
Images []string `json:"images,omitempty"` // Image URLs for multimodal (only for current user message)
// ReasoningContent 是 assistant 推理类模型(DeepSeek thinking、小米 MiMo、vLLM reasoning 等)
// 上一轮输出的思考内容。部分供应商(MiMo、DeepSeek V3.2/V4 thinking 模式)要求多轮对话中
// 把 assistant 的 reasoning_content 原样回传,否则会以 400 拒绝请求;其他不要求的供应商
// 会忽略未知字段,无副作用。
ReasoningContent string `json:"reasoning_content,omitempty"`
// Kind is engine-internal bookkeeping. `json:"-"` keeps it off the wire:
// providers reject unknown message fields on some endpoints, and this one
// means nothing to them anyway.
Kind MessageKind `json:"-"`
}
// ToolCall represents a tool call in a message
type ToolCall struct {
ID string `json:"id"`
Type string `json:"type"` // "function"
Function FunctionCall `json:"function"`
ProviderMetadata types.ToolCallMetadata `json:"provider_metadata,omitempty"`
}
// FunctionCall represents a function call
type FunctionCall struct {
Name string `json:"name"`
Arguments string `json:"arguments"` // JSON string
}
// Chat 定义了聊天接口
type Chat interface {
// Chat 进行非流式聊天
Chat(ctx context.Context, messages []Message, opts *ChatOptions) (*types.ChatResponse, error)
// ChatStream 进行流式聊天
ChatStream(ctx context.Context, messages []Message, opts *ChatOptions) (<-chan types.StreamResponse, error)
// GetModelName 获取模型名称
GetModelName() string
// GetModelID 获取模型ID
GetModelID() string
}
type ChatConfig struct {
Source types.ModelSource
BaseURL string
ModelName string
APIKey string
ModelID string
Provider string
// MaxConcurrency caps concurrent background calls to this model; 0 falls
// back to the process-wide default (see limiter.GateN).
MaxConcurrency int
ExtraConfig map[string]string
// CustomHeaders 允许在调用远程 OpenAI 兼容 API 时附加自定义 HTTP 请求头(类似 OpenAI Python SDK 的 extra_headers)。
CustomHeaders map[string]string
AppID string
AppSecret string // 加密值,由工厂函数调用方传入,在 NewWeKnoraCloudChat 中使用前已解密
}
// ConfigFromModel 根据 types.Model 构造 ChatConfig。
// 保证生产路径(service 层根据 DB 中的模型配置拉起实例)和测试路径
// (handler 层根据前端表单临时拉起实例)走完全相同的字段映射,避免重复样板。
// appID / appSecret 是已经解密/解析好的 WeKnoraCloud 凭证,调用方负责传入。
func ConfigFromModel(m *types.Model, appID, appSecret string) *ChatConfig {
if m == nil {
return nil
}
return &ChatConfig{
ModelID: m.ID,
APIKey: m.Parameters.APIKey,
BaseURL: m.Parameters.BaseURL,
ModelName: m.Name,
Source: m.Source,
Provider: m.Parameters.Provider,
MaxConcurrency: m.Parameters.MaxConcurrency,
ExtraConfig: m.Parameters.ExtraConfig,
CustomHeaders: m.Parameters.CustomHeaders,
AppID: appID,
AppSecret: appSecret,
}
}
// NewChat 创建聊天实例
func NewChat(config *ChatConfig, ollamaService *ollama.OllamaService) (Chat, error) {
var c Chat
var err error
switch strings.ToLower(string(config.Source)) {
case string(types.ModelSourceLocal):
c, err = NewOllamaChat(config, ollamaService)
case string(types.ModelSourceRemote):
c, err = NewRemoteChat(config)
default:
return nil, fmt.Errorf("unsupported chat model source: %s", config.Source)
}
c, err = wrapChatRequestBudget(c, err)
c, err = wrapChatDebug(c, err)
c, err = wrapChatLangfuse(c, err)
// Outermost: hold the per-model concurrency slot only around the real
// provider round-trip, so the wait is excluded from debug/langfuse timing.
return wrapChatConcurrency(c, config.MaxConcurrency, err)
}
// NewRemoteChat 根据 provider 创建远程聊天实例。
// Anthropic 走独立的 Messages 协议实现;其余 OpenAI 兼容供应商统一由
// RemoteAPIChat 处理,provider 特定行为在构造时通过 providerAdapter 解析。
func NewRemoteChat(config *ChatConfig) (Chat, error) {
providerName := provider.ProviderName(config.Provider)
if providerName == "" {
providerName = provider.DetectProvider(config.BaseURL)
}
if providerName == provider.ProviderAnthropic {
return NewAnthropicChat(config)
}
return NewRemoteAPIChat(config)
}
@@ -0,0 +1,82 @@
package chat
import (
"context"
"github.com/Tencent/WeKnora/internal/models/limiter"
"github.com/Tencent/WeKnora/internal/types"
)
// Model provider budgets are the real bottleneck shared by every LLM-backed
// background stage (summary / question / graph / multimodal enrichment), which
// all target the same model. This governor caps concurrent calls per model at
// the client layer — the one place that sees all task types — instead of at the
// asynq queue layer, whose weights are scheduling priority rather than
// throttling.
//
// Only background (asynq worker) calls are throttled; interactive chat is left
// untouched (see types.IsBackgroundTask), so a document-ingestion storm cannot
// exhaust the provider yet user-facing latency is never gated behind the
// semaphore. The governor singleton itself lives in the limiter package so chat
// and vlm share the same limiter and per-model budget.
// concurrencyChat throttles background LLM calls through a per-model
// distributed semaphore. It is the outermost wrapper so the slot is held only
// around the actual provider round-trip and the wait time is excluded from the
// inner debug/langfuse timing.
type concurrencyChat struct {
inner Chat
// limit is this model's configured per-model background cap; 0 falls back
// to the process-wide default (see limiter.GateN).
limit int
}
func (w *concurrencyChat) GetModelName() string { return w.inner.GetModelName() }
func (w *concurrencyChat) GetModelID() string { return w.inner.GetModelID() }
func (w *concurrencyChat) Chat(ctx context.Context, messages []Message, opts *ChatOptions) (*types.ChatResponse, error) {
release := limiter.GateNamedN(ctx, w.inner.GetModelID(), w.inner.GetModelName(), w.limit)
defer release()
return w.inner.Chat(ctx, messages, opts)
}
func (w *concurrencyChat) ChatStream(ctx context.Context, messages []Message, opts *ChatOptions) (<-chan types.StreamResponse, error) {
release := limiter.GateNamedN(ctx, w.inner.GetModelID(), w.inner.GetModelName(), w.limit)
ch, err := w.inner.ChatStream(ctx, messages, opts)
if err != nil || ch == nil {
release()
return ch, err
}
// Hold the slot until the stream fully drains, then release. If the
// consumer abandons the stream (stops reading out) we would otherwise
// block forever on the send and never release the slot; select on
// ctx.Done() so a cancelled call frees its slot promptly, and drain the
// inner channel in the background so the upstream producer can exit.
out := make(chan types.StreamResponse)
go func() {
defer close(out)
defer release()
for resp := range ch {
select {
case out <- resp:
case <-ctx.Done():
go func() {
for range ch {
}
}()
return
}
}
}()
return out, nil
}
// wrapChatConcurrency installs the background concurrency governor as the
// outermost Chat decorator. It is always applied; when no limiter is installed
// or the call is interactive, the wrapper is a cheap passthrough.
func wrapChatConcurrency(c Chat, limit int, err error) (Chat, error) {
if err != nil || c == nil {
return c, err
}
return &concurrencyChat{inner: c, limit: limit}, nil
}
@@ -0,0 +1,93 @@
package chat
import (
"context"
"testing"
"time"
"github.com/Tencent/WeKnora/internal/models/limiter"
"github.com/Tencent/WeKnora/internal/types"
)
// fakeChat is a minimal Chat whose stream emits continuously until ctx is done,
// so we can exercise the concurrency wrapper's slot lifecycle.
type fakeChat struct{ id string }
func (f *fakeChat) GetModelName() string { return f.id }
func (f *fakeChat) GetModelID() string { return f.id }
func (f *fakeChat) Chat(ctx context.Context, _ []Message, _ *ChatOptions) (*types.ChatResponse, error) {
return &types.ChatResponse{}, nil
}
func (f *fakeChat) ChatStream(ctx context.Context, _ []Message, _ *ChatOptions) (<-chan types.StreamResponse, error) {
ch := make(chan types.StreamResponse)
go func() {
defer close(ch)
for {
select {
case <-ctx.Done():
return
case ch <- types.StreamResponse{}:
}
}
}()
return ch, nil
}
// TestConcurrencyChatInteractiveNotGated verifies interactive calls bypass the
// governor entirely even at limit 1.
func TestConcurrencyChatInteractiveNotGated(t *testing.T) {
t.Cleanup(func() { limiter.SetGovernor(nil, 0) })
limiter.SetGovernor(limiter.NewLocalLimiter(), 1)
w := &concurrencyChat{inner: &fakeChat{id: "model-x"}}
// No background marker: neither call should be throttled.
if _, err := w.Chat(context.Background(), nil, nil); err != nil {
t.Fatalf("interactive chat: %v", err)
}
if _, err := w.Chat(context.Background(), nil, nil); err != nil {
t.Fatalf("interactive chat 2: %v", err)
}
}
// TestConcurrencyChatStreamReleasesOnAbandon verifies that when the consumer
// stops reading and cancels the context, the held slot is released rather than
// leaked (the #9 fix).
func TestConcurrencyChatStreamReleasesOnAbandon(t *testing.T) {
t.Cleanup(func() { limiter.SetGovernor(nil, 0) })
limiter.SetGovernor(limiter.NewLocalLimiter(), 1)
const id = "model-y"
w := &concurrencyChat{inner: &fakeChat{id: id}}
streamCtx, cancel := context.WithCancel(types.WithBackgroundTask(context.Background()))
out, err := w.ChatStream(streamCtx, nil, nil)
if err != nil {
t.Fatalf("stream: %v", err)
}
// Consume one item so the relay goroutine is running and holding the slot.
<-out
// The single slot is now held: a background acquire must block.
blocked := make(chan struct{})
go func() {
rel := limiter.Gate(types.WithBackgroundTask(context.Background()), id)
rel()
close(blocked)
}()
select {
case <-blocked:
t.Fatal("slot should be held by the live stream")
case <-time.After(50 * time.Millisecond):
}
// Abandon the stream: stop reading and cancel. The relay must release.
cancel()
select {
case <-blocked:
// released and reacquired successfully
case <-time.After(2 * time.Second):
t.Fatal("stream slot leaked: not released after consumer abandoned + cancelled")
}
}
@@ -0,0 +1,59 @@
package chat
import (
"testing"
"github.com/Tencent/WeKnora/internal/types"
)
// TestConfigFromModel 验证 ConfigFromModel 能把 types.Model 完整映射到 ChatConfig,
// 避免后续新加字段时只在生产路径或测试路径其中一侧更新导致功能漏配。
func TestConfigFromModel(t *testing.T) {
m := &types.Model{
ID: "model-123",
Name: "gpt-4o-mini",
Source: types.ModelSourceRemote,
Parameters: types.ModelParameters{
BaseURL: "https://api.example.com/v1",
APIKey: "sk-xxx",
Provider: "openai",
ExtraConfig: map[string]string{
"deployment": "prod",
},
CustomHeaders: map[string]string{
"X-Gateway-Token": "abc",
"X-Trace-ID": "t-1",
},
},
}
cfg := ConfigFromModel(m, "app-id", "app-secret")
if cfg == nil {
t.Fatal("ConfigFromModel returned nil")
}
if cfg.ModelID != "model-123" || cfg.ModelName != "gpt-4o-mini" {
t.Errorf("model identity not propagated: %+v", cfg)
}
if cfg.BaseURL != "https://api.example.com/v1" || cfg.APIKey != "sk-xxx" {
t.Errorf("connection fields not propagated: %+v", cfg)
}
if cfg.Provider != "openai" {
t.Errorf("provider not propagated: %q", cfg.Provider)
}
if cfg.ExtraConfig["deployment"] != "prod" {
t.Errorf("ExtraConfig not propagated: %+v", cfg.ExtraConfig)
}
if cfg.CustomHeaders["X-Gateway-Token"] != "abc" ||
cfg.CustomHeaders["X-Trace-ID"] != "t-1" {
t.Errorf("CustomHeaders not propagated: %+v", cfg.CustomHeaders)
}
if cfg.AppID != "app-id" || cfg.AppSecret != "app-secret" {
t.Errorf("cloud credentials not propagated: %+v", cfg)
}
}
func TestConfigFromModel_Nil(t *testing.T) {
if got := ConfigFromModel(nil, "", ""); got != nil {
t.Fatalf("expected nil for nil model, got %+v", got)
}
}
+105
View File
@@ -0,0 +1,105 @@
package chat
import (
"encoding/base64"
"fmt"
"log"
"net/http"
"os"
"path/filepath"
"strings"
)
// resolveImageURLForLLM converts stored image paths to a format that LLM APIs can consume.
// - data: URIs and http(s):// URLs are returned as-is.
// - resource:// and provider-backed local paths are read through the
// application resolver and converted to base64 data URIs.
func resolveImageURLForLLM(imageURL string) string {
if strings.HasPrefix(imageURL, "data:") || strings.HasPrefix(imageURL, "http://") || strings.HasPrefix(imageURL, "https://") {
return imageURL
}
if isApplicationStoredImage(imageURL) {
data := readLocalStorageBytes(imageURL)
if data != nil {
mime := http.DetectContentType(data)
return fmt.Sprintf("data:%s;base64,%s", mime, base64.StdEncoding.EncodeToString(data))
}
}
return imageURL
}
// resolveImageURLForOllama converts stored image paths to raw bytes for the Ollama API.
func resolveImageURLForOllama(imageURL string) []byte {
if strings.HasPrefix(imageURL, "data:") {
idx := strings.Index(imageURL, ";base64,")
if idx < 0 {
return nil
}
decoded, err := base64.StdEncoding.DecodeString(imageURL[idx+8:])
if err != nil {
return nil
}
return decoded
}
if isApplicationStoredImage(imageURL) {
return readLocalStorageBytes(imageURL)
}
return nil
}
func isApplicationStoredImage(imageURL string) bool {
return strings.HasPrefix(imageURL, "resource://") ||
strings.HasPrefix(imageURL, "local://") ||
strings.HasPrefix(imageURL, "storage://")
}
// LocalImageResolver, when set by the application layer at startup, resolves a
// resource:// or provider storage URL to bytes using the owning tenant's
// storage config.
// Stored local:// URLs are relative to the storage base dir and do NOT encode
// the tenant's configured PathPrefix, so a plain env-based join would miss the
// prefix. When nil (e.g. in tests), callers fall back to the env-based
// LOCAL_STORAGE_BASE_DIR resolution below.
var LocalImageResolver func(storageURL string) ([]byte, bool)
// readLocalStorageBytes resolves a local:// storage path to disk bytes.
func readLocalStorageBytes(storagePath string) []byte {
if LocalImageResolver != nil {
if data, ok := LocalImageResolver(storagePath); ok {
return data
}
}
relPath := strings.TrimPrefix(storagePath, "local://")
baseDir := os.Getenv("LOCAL_STORAGE_BASE_DIR")
if baseDir == "" {
baseDir = "/data/files"
}
localPath := filepath.Join(baseDir, filepath.FromSlash(relPath))
data, err := os.ReadFile(localPath)
if err != nil {
log.Printf("[image-resolve] failed to read local file %s: %v", localPath, err)
return nil
}
return data
}
// isMultimodalNotSupportedError checks if an error indicates the model does not
// support multimodal/image input.
func isMultimodalNotSupportedError(err error) bool {
if err == nil {
return false
}
msg := strings.ToLower(err.Error())
return (strings.Contains(msg, "multimodal") || strings.Contains(msg, "image") || strings.Contains(msg, "vision")) &&
(strings.Contains(msg, "not support") || strings.Contains(msg, "unsupported") || strings.Contains(msg, "400"))
}
// stripImagesFromMessages returns a copy of messages with all image data removed.
func stripImagesFromMessages(messages []Message) []Message {
cleaned := make([]Message, len(messages))
for i, msg := range messages {
cleaned[i] = msg
cleaned[i].Images = nil
}
return cleaned
}
@@ -0,0 +1,224 @@
package chat
import (
"strings"
"unicode/utf8"
)
// jsonFieldExtractor extracts a specific string field value from streaming JSON fragments.
// It processes incremental JSON argument chunks from LLM tool calls.
//
// Example: for fieldName="answer", expected JSON format: {"answer":"...content..."}
// The extractor uses a simple state machine to skip the JSON prefix and extract the string value.
type jsonFieldExtractor struct {
fieldName string // the JSON field name to extract (e.g. "answer", "thought")
buffer string // accumulated full arguments string
valueStart int // byte offset where the field value starts (-1 if not found yet)
lastEmit int // byte offset of the last emitted position within the value
done bool // whether we've seen the closing quote
}
// newJSONFieldExtractor creates a new extractor instance for the given field name
func newJSONFieldExtractor(fieldName string) *jsonFieldExtractor {
return &jsonFieldExtractor{
fieldName: fieldName,
valueStart: -1,
lastEmit: 0,
}
}
// Feed processes a new argument delta and returns any new content to emit.
// Returns empty string if no new content is available yet.
func (e *jsonFieldExtractor) Feed(argsDelta string) string {
if e.done {
return ""
}
e.buffer += argsDelta
// If we haven't found the value start yet, try to find it
if e.valueStart < 0 {
idx := findFieldValueStart(e.buffer, e.fieldName)
if idx < 0 {
return "" // Haven't seen the value start yet
}
e.valueStart = idx
e.lastEmit = 0
}
// Extract new content from the value portion
valueContent := e.buffer[e.valueStart:]
// Find how far we can safely emit (stop before potential incomplete escape at the end)
safeEnd, finished := findSafeEnd(valueContent, e.lastEmit)
if safeEnd <= e.lastEmit {
if finished {
e.done = true
}
return ""
}
// Extract the new chunk and unescape JSON string escapes
rawChunk := valueContent[e.lastEmit:safeEnd]
unescaped := unescapeJSONString(rawChunk)
e.lastEmit = safeEnd
if finished {
e.done = true
}
return unescaped
}
// IsDone returns whether the extractor has finished (closing quote found)
func (e *jsonFieldExtractor) IsDone() bool {
return e.done
}
// findFieldValueStart finds the byte offset where the field's string value content begins
// (after the opening quote of the value). Returns -1 if not found.
func findFieldValueStart(buf string, fieldName string) int {
// Look for "fieldName" key followed by colon and opening quote
key := `"` + fieldName + `"`
idx := strings.Index(buf, key)
if idx < 0 {
return -1
}
// Skip past the key
pos := idx + len(key)
// Skip whitespace and colon
for pos < len(buf) {
ch := buf[pos]
if ch == ':' {
pos++
continue
}
if ch == ' ' || ch == '\t' || ch == '\n' || ch == '\r' {
pos++
continue
}
if ch == '"' {
// Found the opening quote of the value
return pos + 1
}
// Unexpected character
return -1
}
return -1 // Haven't seen the opening quote yet
}
// findSafeEnd finds the safe end position for emission within the value content.
// It scans from lastEmit forward, handling escape sequences.
// Returns (safeEnd, finished) where finished=true if the closing quote was found.
func findSafeEnd(value string, from int) (int, bool) {
i := from
for i < len(value) {
ch := value[i]
if ch == '\\' {
// Escape sequence - need at least 2 bytes
if i+1 >= len(value) {
// Incomplete escape at end, stop before it
return i, false
}
nextCh := value[i+1]
if nextCh == 'u' {
// Unicode escape \uXXXX - need 6 bytes total
if i+5 >= len(value) {
return i, false
}
i += 6
} else {
// Simple escape: \", \\, \n, \t, \r, \/, \b, \f
i += 2
}
} else if ch == '"' {
// Closing quote of the JSON string value
return i, true
} else {
// Regular character - handle multi-byte UTF-8
_, size := utf8.DecodeRuneInString(value[i:])
if size == 0 {
size = 1
}
i += size
}
}
return i, false
}
// unescapeJSONString converts JSON string escape sequences to their actual characters
func unescapeJSONString(s string) string {
if !strings.ContainsRune(s, '\\') {
return s
}
var b strings.Builder
b.Grow(len(s))
i := 0
for i < len(s) {
if s[i] == '\\' && i+1 < len(s) {
switch s[i+1] {
case '"':
b.WriteByte('"')
i += 2
case '\\':
b.WriteByte('\\')
i += 2
case '/':
b.WriteByte('/')
i += 2
case 'n':
b.WriteByte('\n')
i += 2
case 'r':
b.WriteByte('\r')
i += 2
case 't':
b.WriteByte('\t')
i += 2
case 'b':
b.WriteByte('\b')
i += 2
case 'f':
b.WriteByte('\f')
i += 2
case 'u':
// Unicode escape \uXXXX
if i+5 < len(s) {
// Parse hex digits
hexStr := s[i+2 : i+6]
var codepoint int
for _, h := range hexStr {
codepoint <<= 4
switch {
case h >= '0' && h <= '9':
codepoint += int(h - '0')
case h >= 'a' && h <= 'f':
codepoint += int(h-'a') + 10
case h >= 'A' && h <= 'F':
codepoint += int(h-'A') + 10
}
}
b.WriteRune(rune(codepoint))
i += 6
} else {
b.WriteByte(s[i])
i++
}
default:
b.WriteByte(s[i])
i++
}
} else {
b.WriteByte(s[i])
i++
}
}
return b.String()
}
@@ -0,0 +1,165 @@
package chat
import (
"testing"
)
func TestJSONFieldExtractor_Basic(t *testing.T) {
e := newJSONFieldExtractor("answer")
// Simulate streaming JSON: {"answer":"Hello world"}
got := ""
got += e.Feed(`{"answer":"`)
got += e.Feed(`Hello`)
got += e.Feed(` world`)
got += e.Feed(`"}`)
if got != "Hello world" {
t.Errorf("expected 'Hello world', got %q", got)
}
if !e.IsDone() {
t.Error("expected extractor to be done")
}
}
func TestJSONFieldExtractor_WithEscapes(t *testing.T) {
e := newJSONFieldExtractor("answer")
// Simulate: {"answer":"line1\nline2 and \"quoted\""}
got := ""
got += e.Feed(`{"answer":"line1\nline2`)
got += e.Feed(` and \"quoted`)
got += e.Feed(`\""}`)
expected := "line1\nline2 and \"quoted\""
if got != expected {
t.Errorf("expected %q, got %q", expected, got)
}
}
func TestJSONFieldExtractor_OneChunk(t *testing.T) {
e := newJSONFieldExtractor("answer")
got := e.Feed(`{"answer":"complete answer here"}`)
if got != "complete answer here" {
t.Errorf("expected 'complete answer here', got %q", got)
}
if !e.IsDone() {
t.Error("expected extractor to be done")
}
}
func TestJSONFieldExtractor_SmallChunks(t *testing.T) {
e := newJSONFieldExtractor("answer")
// Very small chunks
got := ""
chunks := []string{`{`, `"`, `a`, `n`, `s`, `w`, `e`, `r`, `"`, `:`, `"`, `H`, `i`, `"`, `}`}
for _, c := range chunks {
got += e.Feed(c)
}
if got != "Hi" {
t.Errorf("expected 'Hi', got %q", got)
}
}
func TestJSONFieldExtractor_Markdown(t *testing.T) {
e := newJSONFieldExtractor("answer")
got := ""
got += e.Feed(`{"answer":"# Title\n\n`)
got += e.Feed(`This is **bold** and `)
got += e.Feed(`*italic* text.\n\n`)
got += e.Feed(`- item 1\n- item 2`)
got += e.Feed(`"}`)
expected := "# Title\n\nThis is **bold** and *italic* text.\n\n- item 1\n- item 2"
if got != expected {
t.Errorf("expected %q, got %q", expected, got)
}
}
func TestJSONFieldExtractor_UnicodeEscape(t *testing.T) {
e := newJSONFieldExtractor("answer")
got := ""
got += e.Feed(`{"answer":"Hello \u4e16\u754c`)
got += e.Feed(`"}`)
expected := "Hello 世界"
if got != expected {
t.Errorf("expected %q, got %q", expected, got)
}
}
func TestJSONFieldExtractor_IncompleteEscapeAtBoundary(t *testing.T) {
e := newJSONFieldExtractor("answer")
// Escape sequence split across chunks
got := ""
got += e.Feed(`{"answer":"before\`)
got += e.Feed(`nafter"}`)
expected := "before\nafter"
if got != expected {
t.Errorf("expected %q, got %q", expected, got)
}
}
func TestJSONFieldExtractor_WhitespaceInJSON(t *testing.T) {
e := newJSONFieldExtractor("answer")
// Whitespace around colon
got := e.Feed(`{ "answer" : "content here" }`)
if got != "content here" {
t.Errorf("expected 'content here', got %q", got)
}
}
func TestJSONFieldExtractor_EmptyAnswer(t *testing.T) {
e := newJSONFieldExtractor("answer")
got := e.Feed(`{"answer":""}`)
if got != "" {
t.Errorf("expected empty string, got %q", got)
}
if !e.IsDone() {
t.Error("expected extractor to be done")
}
}
// Test extracting "thought" field (for thinking tool)
func TestJSONFieldExtractor_ThoughtField(t *testing.T) {
e := newJSONFieldExtractor("thought")
got := ""
got += e.Feed(`{"thought":"Let me analyze`)
got += e.Feed(` the problem step by step`)
got += e.Feed(`","next_thought_needed":true,"thought_number":1,"total_thoughts":3}`)
expected := "Let me analyze the problem step by step"
if got != expected {
t.Errorf("expected %q, got %q", expected, got)
}
if !e.IsDone() {
t.Error("expected extractor to be done")
}
}
func TestJSONFieldExtractor_ThoughtFieldWithEscapes(t *testing.T) {
e := newJSONFieldExtractor("thought")
got := ""
got += e.Feed(`{"thought":"Step 1:\n- Analyze the query\n- `)
got += e.Feed(`Search for \"relevant\" info`)
got += e.Feed(`","thought_number":1}`)
expected := "Step 1:\n- Analyze the query\n- Search for \"relevant\" info"
if got != expected {
t.Errorf("expected %q, got %q", expected, got)
}
}
+249
View File
@@ -0,0 +1,249 @@
package chat
import (
"context"
"time"
"github.com/Tencent/WeKnora/internal/tracing/langfuse"
"github.com/Tencent/WeKnora/internal/types"
)
// langfuseChat wraps a Chat implementation and emits a Langfuse generation
// observation for every Chat/ChatStream call, capturing prompt, response and
// token usage. The wrapper is only installed when the Langfuse manager is
// enabled, so there is no cost for deployments that don't use Langfuse.
type langfuseChat struct {
inner Chat
}
func (l *langfuseChat) GetModelName() string { return l.inner.GetModelName() }
func (l *langfuseChat) GetModelID() string { return l.inner.GetModelID() }
func (l *langfuseChat) Chat(ctx context.Context, messages []Message, opts *ChatOptions) (*types.ChatResponse, error) {
mgr := langfuse.GetManager()
if !mgr.Enabled() {
return l.inner.Chat(ctx, messages, opts)
}
purpose, prefixFingerprint := types.LLMCallMetadataFromContext(ctx)
genCtx, gen := mgr.StartGeneration(ctx, langfuse.GenerationOptions{
Name: "chat.completion",
Model: l.inner.GetModelName(),
Input: buildLangfuseMessages(messages),
ModelParameters: buildLangfuseModelParams(opts),
Metadata: map[string]interface{}{
"model_id": l.inner.GetModelID(),
"streaming": false,
"has_tools": opts != nil && len(opts.Tools) > 0,
"call_purpose": purpose,
"prompt_prefix_fingerprint": prefixFingerprint,
},
})
resp, err := l.inner.Chat(genCtx, messages, opts)
var usage *langfuse.TokenUsage
var output interface{}
if resp != nil {
usage = convertUsage(&resp.Usage)
output = buildLangfuseGenerationOutput(
resp.Content, resp.ReasoningContent, resp.FinishReason, resp.ToolCalls,
)
}
gen.Finish(output, usage, err)
return resp, err
}
func (l *langfuseChat) ChatStream(ctx context.Context, messages []Message, opts *ChatOptions) (<-chan types.StreamResponse, error) {
mgr := langfuse.GetManager()
if !mgr.Enabled() {
return l.inner.ChatStream(ctx, messages, opts)
}
purpose, prefixFingerprint := types.LLMCallMetadataFromContext(ctx)
genCtx, gen := mgr.StartGeneration(ctx, langfuse.GenerationOptions{
Name: "chat.completion.stream",
Model: l.inner.GetModelName(),
Input: buildLangfuseMessages(messages),
ModelParameters: buildLangfuseModelParams(opts),
Metadata: map[string]interface{}{
"model_id": l.inner.GetModelID(),
"streaming": true,
"has_tools": opts != nil && len(opts.Tools) > 0,
"call_purpose": purpose,
"prompt_prefix_fingerprint": prefixFingerprint,
},
})
ch, err := l.inner.ChatStream(genCtx, messages, opts)
if err != nil {
gen.Finish(nil, nil, err)
return ch, err
}
if ch == nil {
gen.Finish(nil, nil, nil)
return nil, nil
}
wrapped := make(chan types.StreamResponse)
go func() {
defer close(wrapped)
var contentBuf []byte
var reasoningBuf []byte
var usage *types.TokenUsage
var toolCalls []types.LLMToolCall
var finishReason string
var firstToken bool
for resp := range ch {
if resp.ResponseType == types.ResponseTypeThinking && resp.Content != "" {
if !firstToken {
gen.MarkCompletionStart(time.Now())
firstToken = true
}
reasoningBuf = append(reasoningBuf, resp.Content...)
}
if resp.ResponseType == types.ResponseTypeAnswer && resp.Content != "" {
if !firstToken {
gen.MarkCompletionStart(time.Now())
firstToken = true
}
contentBuf = append(contentBuf, resp.Content...)
}
if resp.Usage != nil {
usage = resp.Usage
}
if len(resp.ToolCalls) > 0 {
// The downstream model-context registry decodes arguments in
// place before tool execution. Snapshot the provider payload so
// the generation observation remains the exact model output and
// cannot be changed through the shared slice backing array.
toolCalls = snapshotLangfuseToolCalls(resp.ToolCalls)
}
if resp.FinishReason != "" {
finishReason = resp.FinishReason
}
wrapped <- resp
}
output := buildLangfuseGenerationOutput(
string(contentBuf), string(reasoningBuf), finishReason, toolCalls,
)
gen.Finish(output, convertUsage(usage), nil)
}()
return wrapped, nil
}
func snapshotLangfuseToolCalls(toolCalls []types.LLMToolCall) []types.LLMToolCall {
return append([]types.LLMToolCall(nil), toolCalls...)
}
func buildLangfuseMessages(messages []Message) []map[string]interface{} {
out := make([]map[string]interface{}, 0, len(messages))
for _, m := range messages {
entry := map[string]interface{}{
"role": m.Role,
}
if m.Content != "" {
entry["content"] = m.Content
}
if len(m.MultiContent) > 0 {
entry["content"] = m.MultiContent
}
if m.Name != "" {
entry["name"] = m.Name
}
if m.ToolCallID != "" {
entry["tool_call_id"] = m.ToolCallID
}
if len(m.ToolCalls) > 0 {
entry["tool_calls"] = m.ToolCalls
}
if m.ReasoningContent != "" {
entry["reasoning_content"] = m.ReasoningContent
}
out = append(out, entry)
}
return out
}
func buildLangfuseGenerationOutput(
content, reasoningContent, finishReason string,
toolCalls []types.LLMToolCall,
) map[string]interface{} {
output := map[string]interface{}{
"content": content,
"tool_calls": toolCalls,
"finish_reason": finishReason,
}
if reasoningContent != "" {
output["reasoning_content"] = reasoningContent
}
return output
}
func buildLangfuseModelParams(opts *ChatOptions) map[string]interface{} {
if opts == nil {
return nil
}
params := map[string]interface{}{}
if opts.Temperature != 0 {
params["temperature"] = opts.Temperature
}
if opts.TopP != 0 {
params["top_p"] = opts.TopP
}
if opts.MaxTokens > 0 {
params["max_tokens"] = opts.MaxTokens
}
if opts.MaxCompletionTokens > 0 {
params["max_completion_tokens"] = opts.MaxCompletionTokens
}
if opts.FrequencyPenalty != 0 {
params["frequency_penalty"] = opts.FrequencyPenalty
}
if opts.PresencePenalty != 0 {
params["presence_penalty"] = opts.PresencePenalty
}
if opts.Seed != 0 {
params["seed"] = opts.Seed
}
if opts.ToolChoice != "" {
params["tool_choice"] = opts.ToolChoice
}
if len(params) == 0 {
return nil
}
return params
}
func convertUsage(u *types.TokenUsage) *langfuse.TokenUsage {
if u == nil {
return nil
}
if u.PromptTokens == 0 && u.CompletionTokens == 0 && u.TotalTokens == 0 {
return nil
}
return &langfuse.TokenUsage{
Input: u.PromptTokens,
Output: u.CompletionTokens,
Total: u.TotalTokens,
CacheRead: u.CacheReadTokens,
CacheWrite: u.CacheWriteTokens,
CacheMiss: u.CacheMissTokens,
Unit: "TOKENS",
}
}
// wrapChatLangfuse wraps a Chat in a Langfuse-aware decorator when the
// manager is enabled. Called from NewChat after the debug wrapper so both
// sinks observe the same call.
func wrapChatLangfuse(c Chat, err error) (Chat, error) {
if err != nil || c == nil {
return c, err
}
if !langfuse.GetManager().Enabled() {
return c, nil
}
return &langfuseChat{inner: c}, nil
}
@@ -0,0 +1,71 @@
package chat
import (
"reflect"
"testing"
"github.com/Tencent/WeKnora/internal/types"
)
func TestBuildLangfuseGenerationOutput(t *testing.T) {
toolCalls := []types.LLMToolCall{{ID: "call_1", Type: "function"}}
got := buildLangfuseGenerationOutput("", "", "tool_calls", toolCalls)
want := map[string]interface{}{
"content": "",
"tool_calls": toolCalls,
"finish_reason": "tool_calls",
}
if !reflect.DeepEqual(got, want) {
t.Fatalf("output without reasoning = %#v; want %#v", got, want)
}
got = buildLangfuseGenerationOutput("answer", "thinking", "stop", nil)
want = map[string]interface{}{
"content": "answer",
"tool_calls": []types.LLMToolCall(nil),
"finish_reason": "stop",
"reasoning_content": "thinking",
}
if !reflect.DeepEqual(got, want) {
t.Fatalf("output with reasoning = %#v; want %#v", got, want)
}
}
func TestSnapshotLangfuseToolCallsKeepsModelArguments(t *testing.T) {
providerCalls := []types.LLMToolCall{{
ID: "call_1",
Function: types.FunctionCall{Name: "wiki_read_page", Arguments: `{"slugs":["res://0001"]}`},
}}
snapshot := snapshotLangfuseToolCalls(providerCalls)
providerCalls[0].Function.Arguments = `{"slugs":["summary/uuid"]}`
if got := snapshot[0].Function.Arguments; got != `{"slugs":["res://0001"]}` {
t.Fatalf("Langfuse snapshot was mutated to %s", got)
}
}
func TestBuildLangfuseMessagesReasoningContent(t *testing.T) {
msgs := buildLangfuseMessages([]Message{
{Role: "assistant", ReasoningContent: "chain of thought", ToolCalls: []ToolCall{{ID: "tc1"}}},
})
if len(msgs) != 1 {
t.Fatalf("len(messages) = %d; want 1", len(msgs))
}
if msgs[0]["reasoning_content"] != "chain of thought" {
t.Fatalf("reasoning_content = %v; want chain of thought", msgs[0]["reasoning_content"])
}
}
func TestConvertUsageIncludesPromptCacheCounters(t *testing.T) {
got := convertUsage(&types.TokenUsage{
PromptTokens: 1000, CompletionTokens: 50, TotalTokens: 1050,
CacheReadTokens: 800, CacheWriteTokens: 100, CacheMissTokens: 200,
})
if got == nil {
t.Fatal("convertUsage returned nil")
}
if got.CacheRead != 800 || got.CacheWrite != 100 || got.CacheMiss != 200 {
t.Fatalf("cache usage = read:%d write:%d miss:%d", got.CacheRead, got.CacheWrite, got.CacheMiss)
}
}
+215
View File
@@ -0,0 +1,215 @@
package chat
import (
"context"
"fmt"
"strings"
"time"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/types"
)
func buildLLMMessages(messages []Message) []logger.LLMMessage {
out := make([]logger.LLMMessage, 0, len(messages))
for _, m := range messages {
lm := logger.LLMMessage{
Role: m.Role,
Content: m.Content,
Name: m.Name,
ToolCallID: m.ToolCallID,
Images: m.Images,
}
if lm.Content == "" && len(m.MultiContent) > 0 {
var parts []string
for _, mc := range m.MultiContent {
switch mc.Type {
case "text":
parts = append(parts, mc.Text)
case "image_url":
if mc.ImageURL != nil {
parts = append(parts, fmt.Sprintf("[image_url: %s]", truncateForDebug(mc.ImageURL.URL, 120)))
}
}
}
lm.Content = strings.Join(parts, "\n")
}
for _, tc := range m.ToolCalls {
lm.ToolCalls = append(lm.ToolCalls, logger.LLMToolCallInfo{
ID: tc.ID,
FuncName: tc.Function.Name,
Arguments: tc.Function.Arguments,
})
}
out = append(out, lm)
}
return out
}
func buildOptionsSection(opts *ChatOptions) string {
if opts == nil {
return ""
}
var parts []string
parts = append(parts, fmt.Sprintf("Temperature=%.2f", opts.Temperature))
if opts.TopP > 0 {
parts = append(parts, fmt.Sprintf("TopP=%.2f", opts.TopP))
}
if opts.MaxTokens > 0 {
parts = append(parts, fmt.Sprintf("MaxTokens=%d", opts.MaxTokens))
}
if opts.MaxCompletionTokens > 0 {
parts = append(parts, fmt.Sprintf("MaxCompletionTokens=%d", opts.MaxCompletionTokens))
}
if opts.FrequencyPenalty > 0 {
parts = append(parts, fmt.Sprintf("FrequencyPenalty=%.2f", opts.FrequencyPenalty))
}
if opts.PresencePenalty > 0 {
parts = append(parts, fmt.Sprintf("PresencePenalty=%.2f", opts.PresencePenalty))
}
if opts.ToolChoice != "" {
parts = append(parts, fmt.Sprintf("ToolChoice=%s", opts.ToolChoice))
}
if len(opts.Format) > 0 {
parts = append(parts, "ResponseFormat=json_object")
}
return strings.Join(parts, ", ")
}
func buildToolsSection(opts *ChatOptions) string {
if opts == nil || len(opts.Tools) == 0 {
return ""
}
var b strings.Builder
for i, t := range opts.Tools {
if i > 0 {
b.WriteString("\n")
}
b.WriteString(fmt.Sprintf("- %s: %s", t.Function.Name, t.Function.Description))
}
return b.String()
}
func buildResponseToolCalls(tcs []types.LLMToolCall) []logger.LLMToolCallInfo {
if len(tcs) == 0 {
return nil
}
out := make([]logger.LLMToolCallInfo, 0, len(tcs))
for _, tc := range tcs {
out = append(out, logger.LLMToolCallInfo{
ID: tc.ID,
FuncName: tc.Function.Name,
Arguments: tc.Function.Arguments,
})
}
return out
}
func usageString(u types.TokenUsage) string {
return fmt.Sprintf("Prompt: %d, Completion: %d, Total: %d, CacheRead: %d, CacheWrite: %d, CacheMiss: %d, CacheStatus: %s",
u.PromptTokens, u.CompletionTokens, u.TotalTokens, u.CacheReadTokens,
u.CacheWriteTokens, u.CacheMissTokens, u.CacheStatus)
}
// logLLMDebugCall logs a complete non-stream LLM chat call.
func logLLMDebugCall(ctx context.Context, model string, messages []Message, opts *ChatOptions, resp *types.ChatResponse, callErr error, dur time.Duration) {
if !logger.LLMDebugEnabled() {
return
}
record := &logger.LLMCallRecord{
CallType: "Chat",
Model: model,
Duration: dur,
}
record.Sections = append(record.Sections, logger.RecordSection{
Title: "Messages",
Content: logger.FormatMessages(buildLLMMessages(messages)),
})
if s := buildOptionsSection(opts); s != "" {
record.Sections = append(record.Sections, logger.RecordSection{Title: "Options", Content: s})
}
if s := buildToolsSection(opts); s != "" {
record.Sections = append(record.Sections, logger.RecordSection{Title: "Tools", Content: s})
}
if resp != nil {
var respText strings.Builder
if resp.Content != "" {
respText.WriteString("[assistant]\n")
respText.WriteString(resp.Content)
respText.WriteString("\n")
}
tcs := buildResponseToolCalls(resp.ToolCalls)
if len(tcs) > 0 {
respText.WriteString(logger.FormatToolCalls(tcs))
}
if respText.Len() > 0 {
record.Sections = append(record.Sections, logger.RecordSection{Title: "Response", Content: respText.String()})
}
record.Sections = append(record.Sections, logger.RecordSection{Title: "Usage", Content: usageString(resp.Usage)})
}
if callErr != nil {
record.Error = callErr.Error()
}
logger.LLMDebugLog(ctx, record)
}
// logLLMDebugStream logs a complete stream LLM chat call after all chunks have been received.
func logLLMDebugStream(ctx context.Context, model string, messages []Message, opts *ChatOptions, fullContent string, toolCalls []types.LLMToolCall, usage *types.TokenUsage, callErr error, dur time.Duration) {
if !logger.LLMDebugEnabled() {
return
}
record := &logger.LLMCallRecord{
CallType: "Chat Stream",
Model: model,
Duration: dur,
}
record.Sections = append(record.Sections, logger.RecordSection{
Title: "Messages",
Content: logger.FormatMessages(buildLLMMessages(messages)),
})
if s := buildOptionsSection(opts); s != "" {
record.Sections = append(record.Sections, logger.RecordSection{Title: "Options", Content: s})
}
if s := buildToolsSection(opts); s != "" {
record.Sections = append(record.Sections, logger.RecordSection{Title: "Tools", Content: s})
}
var respText strings.Builder
if fullContent != "" {
respText.WriteString("[assistant]\n")
respText.WriteString(fullContent)
respText.WriteString("\n")
}
tcs := buildResponseToolCalls(toolCalls)
if len(tcs) > 0 {
respText.WriteString(logger.FormatToolCalls(tcs))
}
if respText.Len() > 0 {
record.Sections = append(record.Sections, logger.RecordSection{Title: "Response", Content: respText.String()})
}
if usage != nil {
record.Sections = append(record.Sections, logger.RecordSection{Title: "Usage", Content: usageString(*usage)})
}
if callErr != nil {
record.Error = callErr.Error()
}
logger.LLMDebugLog(ctx, record)
}
func truncateForDebug(s string, maxRunes int) string {
runes := []rune(s)
if len(runes) <= maxRunes {
return s
}
return string(runes[:maxRunes]) + fmt.Sprintf("...(%d chars)", len(runes))
}
+77
View File
@@ -0,0 +1,77 @@
package chat
import (
"context"
"fmt"
"strings"
"time"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/types"
)
// debugChat wraps a Chat implementation and logs all calls to the LLM debug log.
// It works with both RemoteAPIChat and OllamaChat (or any future Chat impl).
type debugChat struct {
inner Chat
}
func (d *debugChat) GetModelName() string { return d.inner.GetModelName() }
func (d *debugChat) GetModelID() string { return d.inner.GetModelID() }
func (d *debugChat) Chat(ctx context.Context, messages []Message, opts *ChatOptions) (*types.ChatResponse, error) {
callStart := time.Now()
resp, err := d.inner.Chat(ctx, messages, opts)
logLLMDebugCall(ctx, d.inner.GetModelName(), messages, opts, resp, err, time.Since(callStart))
return resp, err
}
func (d *debugChat) ChatStream(ctx context.Context, messages []Message, opts *ChatOptions) (<-chan types.StreamResponse, error) {
callStart := time.Now()
ch, err := d.inner.ChatStream(ctx, messages, opts)
if err != nil {
logLLMDebugStream(ctx, d.inner.GetModelName(), messages, opts, "", nil, nil, err, time.Since(callStart))
return ch, err
}
if ch == nil {
return nil, nil
}
wrapped := make(chan types.StreamResponse)
go func() {
defer close(wrapped)
var content strings.Builder
var usage *types.TokenUsage
var toolCalls []types.LLMToolCall
var streamErr error
for resp := range ch {
if resp.ResponseType == types.ResponseTypeAnswer && resp.Content != "" {
content.WriteString(resp.Content)
}
if resp.ResponseType == types.ResponseTypeError {
streamErr = fmt.Errorf("%s", resp.Content)
}
if resp.Usage != nil {
usage = resp.Usage
}
if len(resp.ToolCalls) > 0 {
toolCalls = resp.ToolCalls
}
wrapped <- resp
}
logLLMDebugStream(ctx, d.inner.GetModelName(), messages, opts,
content.String(), toolCalls, usage, streamErr, time.Since(callStart))
}()
return wrapped, nil
}
// wrapChatDebug wraps a Chat if LLM debug logging is enabled.
func wrapChatDebug(c Chat, err error) (Chat, error) {
if err != nil || !logger.LLMDebugEnabled() {
return c, err
}
return &debugChat{inner: c}, nil
}
+405
View File
@@ -0,0 +1,405 @@
package chat
import (
"context"
"encoding/json"
"fmt"
"io"
"strconv"
"strings"
"time"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/models/requestbudget"
"github.com/Tencent/WeKnora/internal/models/utils/ollama"
"github.com/Tencent/WeKnora/internal/types"
secutils "github.com/Tencent/WeKnora/internal/utils"
ollamaapi "github.com/ollama/ollama/api"
)
// OllamaChat 实现了基于 Ollama 的聊天
type OllamaChat struct {
modelName string
modelID string
ollamaService *ollama.OllamaService
}
// NewOllamaChat 创建 Ollama 聊天实例
func NewOllamaChat(config *ChatConfig, ollamaService *ollama.OllamaService) (*OllamaChat, error) {
return &OllamaChat{
modelName: config.ModelName,
modelID: config.ModelID,
ollamaService: ollamaService,
}, nil
}
// convertMessages 转换消息格式为Ollama API格式
func (c *OllamaChat) convertMessages(messages []Message) []ollamaapi.Message {
ollamaMessages := make([]ollamaapi.Message, 0, len(messages))
for _, msg := range messages {
msgOllama := ollamaapi.Message{
Role: msg.Role,
Content: msg.Content,
ToolCalls: c.toolCallFrom(msg.ToolCalls),
}
if msg.Role == "tool" {
msgOllama.ToolName = msg.Name
}
if len(msg.Images) > 0 && msg.Role == "user" {
for _, imgURL := range msg.Images {
if imgData := resolveImageForOllama(imgURL); imgData != nil {
msgOllama.Images = append(msgOllama.Images, imgData)
}
}
}
ollamaMessages = append(ollamaMessages, msgOllama)
}
return ollamaMessages
}
// resolveImageForOllama resolves an image URL into raw bytes for Ollama.
// Handles local serving paths (/files/...), data URIs, and remote HTTP URLs.
func resolveImageForOllama(imageURL string) ollamaapi.ImageData {
if data := resolveImageURLForOllama(imageURL); data != nil {
return data
}
if strings.HasPrefix(imageURL, "http://") || strings.HasPrefix(imageURL, "https://") {
if err := secutils.ValidateURLForSSRF(imageURL); err != nil {
return nil
}
client := secutils.NewSSRFSafeHTTPClient(secutils.SSRFSafeHTTPClientConfig{
Timeout: 30 * time.Second,
MaxRedirects: 5,
})
resp, err := client.Get(imageURL)
if err != nil {
return nil
}
defer resp.Body.Close()
data, err := io.ReadAll(io.LimitReader(resp.Body, 20*1024*1024))
if err != nil {
return nil
}
return data
}
return nil
}
// buildChatRequest 构建聊天请求参数
func (c *OllamaChat) buildChatRequest(messages []Message, opts *ChatOptions, isStream bool) *ollamaapi.ChatRequest {
// 设置流式标志
streamFlag := isStream
// 构建请求参数
chatReq := &ollamaapi.ChatRequest{
Model: c.modelName,
Messages: c.convertMessages(messages),
Stream: &streamFlag,
Options: make(map[string]interface{}),
}
// 添加可选参数
if opts != nil {
chatReq.Options["temperature"] = opts.Temperature
if opts.TopP > 0 {
chatReq.Options["top_p"] = opts.TopP
}
if opts.MaxTokens > 0 {
chatReq.Options["num_predict"] = opts.MaxTokens
}
if opts.Thinking != nil {
chatReq.Think = &ollamaapi.ThinkValue{
Value: *opts.Thinking,
}
}
if len(opts.Format) > 0 {
chatReq.Format = opts.Format
}
if len(opts.Tools) > 0 {
chatReq.Tools = c.toolFrom(opts.Tools)
}
}
return chatReq
}
// Chat 进行非流式聊天
func (c *OllamaChat) Chat(ctx context.Context, messages []Message, opts *ChatOptions) (*types.ChatResponse, error) {
if err := requestbudget.RequireSupported(ctx, false); err != nil {
return nil, err
}
// 确保模型可用
if err := c.ensureModelAvailable(ctx); err != nil {
return nil, err
}
// 构建请求参数
chatReq := c.buildChatRequest(messages, opts, false)
// 记录请求日志
logger.GetLogger(ctx).Infof("发送聊天请求到模型 %s", c.modelName)
var responseContent string
var toolCalls []types.LLMToolCall
var promptTokens, completionTokens int
// 使用 Ollama 客户端发送请求
err := c.ollamaService.Chat(ctx, chatReq, func(resp ollamaapi.ChatResponse) error {
responseContent = resp.Message.Content
// 当 Content 为空但 Thinking 有内容时(如推理模型未正确配置 thinking 参数),使用 Thinking 作为兜底
if responseContent == "" && resp.Message.Thinking != "" {
responseContent = resp.Message.Thinking
}
toolCalls = c.toolCallTo(resp.Message.ToolCalls)
// 获取token计数
if resp.EvalCount > 0 {
promptTokens = resp.PromptEvalCount
completionTokens = resp.EvalCount - promptTokens
}
return nil
})
if err != nil {
return nil, fmt.Errorf("聊天请求失败: %w", err)
}
usage := types.TokenUsage{
PromptTokens: promptTokens,
CompletionTokens: completionTokens,
TotalTokens: promptTokens + completionTokens,
}
usage.MarkPromptCacheUnsupported()
logUsage(ctx, c.modelName, &usage)
return &types.ChatResponse{
Content: responseContent,
ToolCalls: toolCalls,
Usage: usage,
}, nil
}
// ChatStream 进行流式聊天
func (c *OllamaChat) ChatStream(
ctx context.Context,
messages []Message,
opts *ChatOptions,
) (<-chan types.StreamResponse, error) {
if err := requestbudget.RequireSupported(ctx, false); err != nil {
return nil, err
}
// 确保模型可用
if err := c.ensureModelAvailable(ctx); err != nil {
return nil, err
}
// 构建请求参数
chatReq := c.buildChatRequest(messages, opts, true)
// 记录请求日志
logger.GetLogger(ctx).Infof("发送流式聊天请求到模型 %s", c.modelName)
// 创建流式响应通道
streamChan := make(chan types.StreamResponse)
// 启动goroutine处理流式响应
go func() {
defer close(streamChan)
var thinking thinkingEmitter
err := c.ollamaService.Chat(ctx, chatReq, func(resp ollamaapi.ChatResponse) error {
// 发送思考内容(支持 Qwen3、DeepSeek 等推理模型)
if resp.Message.Thinking != "" {
thinking.emit(streamChan, resp.Message.Thinking)
}
if resp.Message.Content != "" {
// 思考阶段结束后,发送思考完成事件
thinking.finish(streamChan)
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeAnswer,
Content: resp.Message.Content,
Done: false,
}
}
if len(resp.Message.ToolCalls) > 0 {
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeToolCall,
ToolCalls: c.toolCallTo(resp.Message.ToolCalls),
Done: false,
}
// Ollama returns tool calls as complete objects (not incremental deltas).
// Log this so we can trace non-streaming thought delivery.
for _, tc := range resp.Message.ToolCalls {
if tc.Function.Name == "thinking" {
argsBytes, _ := json.Marshal(tc.Function.Arguments)
logger.Warnf(ctx, "[Ollama Stream] Tool %q arrived non-incrementally (%d bytes args), "+
"thought will not be token-streamed to frontend",
tc.Function.Name, len(argsBytes))
}
}
for _, tc := range resp.Message.ToolCalls {
argsMap := tc.Function.Arguments.ToMap()
switch tc.Function.Name {
case "thinking":
if thought, ok := argsMap["thought"].(string); ok && thought != "" {
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeThinking,
Content: thought,
Done: false,
Data: map[string]interface{}{
"source": "thinking_tool",
"tool_call_id": tooli2s(tc.Function.Index),
},
}
}
}
}
}
if resp.Done {
var usage *types.TokenUsage
if resp.PromptEvalCount > 0 || resp.EvalCount > 0 {
usage = &types.TokenUsage{
PromptTokens: resp.PromptEvalCount,
CompletionTokens: resp.EvalCount,
TotalTokens: resp.PromptEvalCount + resp.EvalCount,
}
usage.MarkPromptCacheUnsupported()
}
logUsage(ctx, c.modelName, usage)
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeAnswer,
Done: true,
Usage: usage,
}
}
return nil
})
if err != nil {
logger.GetLogger(ctx).Errorf("流式聊天请求失败: %v", err)
// 发送错误响应
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeError,
Content: err.Error(),
Done: true,
}
}
}()
return streamChan, nil
}
// 确保模型可用
func (c *OllamaChat) ensureModelAvailable(ctx context.Context) error {
logger.GetLogger(ctx).Infof("确保模型 %s 可用", c.modelName)
return c.ollamaService.EnsureModelAvailable(ctx, c.modelName)
}
// GetModelName 获取模型名称
func (c *OllamaChat) GetModelName() string {
return c.modelName
}
// GetModelID 获取模型ID
func (c *OllamaChat) GetModelID() string {
return c.modelID
}
// toolFrom 将本模块的 Tool 转换为 Ollama 的 Tool
func (c *OllamaChat) toolFrom(tools []Tool) ollamaapi.Tools {
if len(tools) == 0 {
return nil
}
ollamaTools := make(ollamaapi.Tools, 0, len(tools))
for _, tool := range tools {
function := ollamaapi.ToolFunction{
Name: tool.Function.Name,
Description: tool.Function.Description,
}
if len(tool.Function.Parameters) > 0 {
_ = json.Unmarshal(tool.Function.Parameters, &function.Parameters)
}
ollamaTools = append(ollamaTools, ollamaapi.Tool{
Type: tool.Type,
Function: function,
})
}
return ollamaTools
}
// toolTo 将 Ollama 的 Tool 转换为本模块的 Tool
func (c *OllamaChat) toolTo(ollamaTools ollamaapi.Tools) []Tool {
if len(ollamaTools) == 0 {
return nil
}
tools := make([]Tool, 0, len(ollamaTools))
for _, tool := range ollamaTools {
paramsBytes, _ := json.Marshal(tool.Function.Parameters)
tools = append(tools, Tool{
Type: tool.Type,
Function: FunctionDef{
Name: tool.Function.Name,
Description: tool.Function.Description,
Parameters: paramsBytes,
},
})
}
return tools
}
// toolCallFrom 将本模块的 ToolCall 转换为 Ollama 的 ToolCall
func (c *OllamaChat) toolCallFrom(toolCalls []ToolCall) []ollamaapi.ToolCall {
if len(toolCalls) == 0 {
return nil
}
ollamaToolCalls := make([]ollamaapi.ToolCall, 0, len(toolCalls))
for _, tc := range toolCalls {
args := ollamaapi.NewToolCallFunctionArguments()
if tc.Function.Arguments != "" {
_ = args.UnmarshalJSON([]byte(tc.Function.Arguments))
}
ollamaToolCalls = append(ollamaToolCalls, ollamaapi.ToolCall{
Function: ollamaapi.ToolCallFunction{
Index: tools2i(tc.ID),
Name: tc.Function.Name,
Arguments: args,
},
})
}
return ollamaToolCalls
}
// toolCallTo 将 Ollama 的 ToolCall 转换为本模块的 ToolCall
func (c *OllamaChat) toolCallTo(ollamaToolCalls []ollamaapi.ToolCall) []types.LLMToolCall {
if len(ollamaToolCalls) == 0 {
return nil
}
toolCalls := make([]types.LLMToolCall, 0, len(ollamaToolCalls))
for _, tc := range ollamaToolCalls {
argsBytes, _ := json.Marshal(tc.Function.Arguments)
toolCalls = append(toolCalls, types.LLMToolCall{
ID: tooli2s(tc.Function.Index),
Type: "function",
Function: types.FunctionCall{
Name: tc.Function.Name,
Arguments: string(argsBytes),
},
})
}
return toolCalls
}
func tooli2s(i int) string {
return strconv.Itoa(i)
}
func tools2i(s string) int {
i, _ := strconv.Atoi(s)
return i
}
+34
View File
@@ -0,0 +1,34 @@
package chat
import (
"net/http"
"net/http/httptest"
"testing"
secutils "github.com/Tencent/WeKnora/internal/utils"
)
func TestResolveImageForOllamaRejectsInternalURL(t *testing.T) {
t.Setenv("SSRF_WHITELIST", "")
secutils.ResetSSRFWhitelistForTest()
t.Cleanup(secutils.ResetSSRFWhitelistForTest)
if data := resolveImageForOllama("http://169.254.169.254/latest/meta-data/"); data != nil {
t.Fatalf("resolveImageForOllama returned data for blocked internal URL")
}
}
func TestResolveImageForOllamaBlocksRedirectToInternalURL(t *testing.T) {
t.Setenv("SSRF_WHITELIST", "127.0.0.1")
secutils.ResetSSRFWhitelistForTest()
t.Cleanup(secutils.ResetSSRFWhitelistForTest)
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
http.Redirect(w, r, "http://169.254.169.254/latest/meta-data/", http.StatusFound)
}))
defer server.Close()
if data := resolveImageForOllama(server.URL); data != nil {
t.Fatalf("resolveImageForOllama returned data after redirect to blocked internal URL")
}
}
+242
View File
@@ -0,0 +1,242 @@
package chat
import (
"encoding/json"
"fmt"
"github.com/sashabaranov/go-openai"
)
// ConvertMessages 转换消息格式为 OpenAI 格式(导出供子类使用)
func (c *RemoteAPIChat) ConvertMessages(messages []Message) []openai.ChatCompletionMessage {
openaiMessages := make([]openai.ChatCompletionMessage, 0, len(messages))
for _, msg := range messages {
openaiMsg := openai.ChatCompletionMessage{
Role: msg.Role,
}
// 优先处理多内容消息(包含图片等)
if len(msg.MultiContent) > 0 {
openaiMsg.MultiContent = make([]openai.ChatMessagePart, 0, len(msg.MultiContent))
for _, part := range msg.MultiContent {
switch part.Type {
case "text":
openaiMsg.MultiContent = append(openaiMsg.MultiContent, openai.ChatMessagePart{
Type: openai.ChatMessagePartTypeText,
Text: part.Text,
})
case "image_url":
if part.ImageURL != nil {
openaiMsg.MultiContent = append(openaiMsg.MultiContent, openai.ChatMessagePart{
Type: openai.ChatMessagePartTypeImageURL,
ImageURL: &openai.ChatMessageImageURL{
URL: part.ImageURL.URL,
Detail: openai.ImageURLDetail(part.ImageURL.Detail),
},
})
}
}
}
} else if len(msg.Images) > 0 && msg.Role == "user" {
parts := make([]openai.ChatMessagePart, 0, len(msg.Images)+1)
for _, imgURL := range msg.Images {
resolved := resolveImageURLForLLM(imgURL)
parts = append(parts, openai.ChatMessagePart{
Type: openai.ChatMessagePartTypeImageURL,
ImageURL: &openai.ChatMessageImageURL{
URL: resolved,
Detail: openai.ImageURLDetailAuto,
},
})
}
parts = append(parts, openai.ChatMessagePart{
Type: openai.ChatMessagePartTypeText,
Text: msg.Content,
})
openaiMsg.MultiContent = parts
} else if msg.Content != "" {
openaiMsg.Content = msg.Content
}
if len(msg.ToolCalls) > 0 {
openaiMsg.ToolCalls = make([]openai.ToolCall, 0, len(msg.ToolCalls))
for _, tc := range msg.ToolCalls {
toolType := openai.ToolType(tc.Type)
openaiMsg.ToolCalls = append(openaiMsg.ToolCalls, openai.ToolCall{
ID: tc.ID,
Type: toolType,
Function: openai.FunctionCall{
Name: tc.Function.Name,
Arguments: tc.Function.Arguments,
},
})
}
}
if msg.Role == "tool" {
openaiMsg.ToolCallID = msg.ToolCallID
openaiMsg.Name = msg.Name
}
// Round-trip reasoning_content on assistant turns. MiMo and DeepSeek V3.2+
// thinking mode reject multi-turn requests where the prior assistant
// message lacks its reasoning_content with HTTP 400 ("The reasoning_content
// in the thinking mode must be passed back to the API."). Providers that
// don't recognize the field ignore it harmlessly.
if msg.Role == "assistant" && msg.ReasoningContent != "" {
openaiMsg.ReasoningContent = msg.ReasoningContent
}
openaiMessages = append(openaiMessages, openaiMsg)
}
return openaiMessages
}
// BuildChatCompletionRequest 构建标准聊天请求参数(导出供子类使用)。
//
// 这是一个不含任何 provider 特定逻辑的通用实现:所有采样参数(temperature /
// top_p / penalties)与 max_tokens 都按 opts 直接映射。供应商相关的特判
// (OpenAI o-series / GPT-5 改用 max_completion_tokens、Moonshot 固定温度等)
// 由对应的 providerAdapter.ShapeRequest 在事后施加,见 provider.go。
func (c *RemoteAPIChat) BuildChatCompletionRequest(
messages []Message, opts *ChatOptions, isStream bool,
) openai.ChatCompletionRequest {
req := openai.ChatCompletionRequest{
Model: c.modelName,
Messages: c.ConvertMessages(messages),
Stream: isStream,
}
if isStream {
req.StreamOptions = &openai.StreamOptions{IncludeUsage: true}
}
if opts == nil {
return req
}
req.Temperature = float32(opts.Temperature)
if opts.TopP > 0 {
req.TopP = float32(opts.TopP)
}
if opts.FrequencyPenalty > 0 {
req.FrequencyPenalty = float32(opts.FrequencyPenalty)
}
if opts.PresencePenalty > 0 {
req.PresencePenalty = float32(opts.PresencePenalty)
}
if opts.MaxTokens > 0 {
req.MaxTokens = opts.MaxTokens
}
if opts.MaxCompletionTokens > 0 {
req.MaxCompletionTokens = opts.MaxCompletionTokens
}
// 处理 Tools
if len(opts.Tools) > 0 {
req.Tools = make([]openai.Tool, 0, len(opts.Tools))
for _, tool := range opts.Tools {
toolType := openai.ToolType(tool.Type)
openaiTool := openai.Tool{
Type: toolType,
Function: &openai.FunctionDefinition{
Name: tool.Function.Name,
Description: tool.Function.Description,
},
}
if tool.Function.Parameters != nil {
openaiTool.Function.Parameters = tool.Function.Parameters
}
req.Tools = append(req.Tools, openaiTool)
}
}
// 处理 ParallelToolCalls
if opts.ParallelToolCalls != nil {
val := *opts.ParallelToolCalls
req.ParallelToolCalls = val
}
// 处理 ToolChoice(标准实现)
if opts.ToolChoice != "" {
switch opts.ToolChoice {
case "none", "required", "auto":
req.ToolChoice = opts.ToolChoice
default:
req.ToolChoice = openai.ToolChoice{
Type: "function",
Function: openai.ToolFunction{
Name: opts.ToolChoice,
},
}
}
}
if len(opts.Format) > 0 {
req.ResponseFormat = &openai.ChatCompletionResponseFormat{
Type: openai.ChatCompletionResponseFormatTypeJSONObject,
}
req.Messages[len(req.Messages)-1].Content += fmt.Sprintf("\nUse this JSON schema: %s", opts.Format)
}
return req
}
func (c *RemoteAPIChat) buildProviderOpenAIRequest(
body any,
openAIMessages []openai.ChatCompletionMessage,
messages []Message,
) (map[string]any, error) {
data, err := json.Marshal(body)
if err != nil {
return nil, fmt.Errorf("marshal provider request: %w", err)
}
var out map[string]any
if err := json.Unmarshal(data, &out); err != nil {
return nil, fmt.Errorf("unmarshal provider request: %w", err)
}
providerMessages := make([]map[string]any, 0, len(openAIMessages))
for i, msg := range openAIMessages {
msgData, err := json.Marshal(msg)
if err != nil {
return nil, fmt.Errorf("marshal provider message: %w", err)
}
var msgMap map[string]any
if err := json.Unmarshal(msgData, &msgMap); err != nil {
return nil, fmt.Errorf("unmarshal provider message: %w", err)
}
if i < len(messages) && len(messages[i].ToolCalls) > 0 && len(msg.ToolCalls) > 0 {
toolCalls := make([]map[string]any, 0, len(msg.ToolCalls))
for j, tc := range msg.ToolCalls {
tcData, err := json.Marshal(tc)
if err != nil {
return nil, fmt.Errorf("marshal provider tool call: %w", err)
}
var tcMap map[string]any
if err := json.Unmarshal(tcData, &tcMap); err != nil {
return nil, fmt.Errorf("unmarshal provider tool call: %w", err)
}
if j < len(messages[i].ToolCalls) {
c.adapter.InjectToolCallMetadata(tcMap, messages[i].ToolCalls[j].ProviderMetadata)
}
toolCalls = append(toolCalls, tcMap)
}
msgMap["tool_calls"] = toolCalls
}
providerMessages = append(providerMessages, msgMap)
}
out["messages"] = providerMessages
return out, nil
}
func (c *RemoteAPIChat) shapeProviderRequest(body any, req openai.ChatCompletionRequest, messages []Message) (any, error) {
if !c.adapter.ForceRawHTTP() {
return body, nil
}
return c.buildProviderOpenAIRequest(body, req.Messages, messages)
}
+631
View File
@@ -0,0 +1,631 @@
package chat
import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
"time"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/types"
"github.com/sashabaranov/go-openai"
)
// parseCompletionResponse 解析非流式响应
func (c *RemoteAPIChat) parseCompletionResponse(resp *openai.ChatCompletionResponse) (*types.ChatResponse, error) {
if len(resp.Choices) == 0 {
return nil, fmt.Errorf("no response from API")
}
choice := resp.Choices[0]
// 处理思考模型的输出:移除 <think></think> 标签包裹的思考过程
// 为设置了 Thinking=false 但模型仍返回思考内容的情况和部分不支持Thinking=false的思考模型(例如Miniax-M2.1)提供兜底策略
content := removeThinkingContent(choice.Message.Content)
usage := tokenUsageFromOpenAI(resp.Usage, c.provider)
response := &types.ChatResponse{
Content: content,
ReasoningContent: choice.Message.ReasoningContent,
FinishReason: string(choice.FinishReason),
Usage: usage,
}
if len(choice.Message.ToolCalls) > 0 {
response.ToolCalls = make([]types.LLMToolCall, 0, len(choice.Message.ToolCalls))
for _, tc := range choice.Message.ToolCalls {
response.ToolCalls = append(response.ToolCalls, types.LLMToolCall{
ID: tc.ID,
Type: string(tc.Type),
Function: types.FunctionCall{
Name: tc.Function.Name,
Arguments: tc.Function.Arguments,
},
})
}
}
return response, nil
}
func (c *RemoteAPIChat) applyCompletionToolCallMetadata(body []byte, result *types.ChatResponse) {
if result == nil || len(result.ToolCalls) == 0 {
return
}
var raw struct {
Choices []struct {
Message struct {
ToolCalls []json.RawMessage `json:"tool_calls,omitempty"`
} `json:"message"`
} `json:"choices"`
}
if err := json.Unmarshal(body, &raw); err != nil || len(raw.Choices) == 0 {
return
}
for i, rawToolCall := range raw.Choices[0].Message.ToolCalls {
var indexed struct {
Index *int `json:"index,omitempty"`
}
_ = json.Unmarshal(rawToolCall, &indexed)
idx := i
if indexed.Index != nil {
idx = *indexed.Index
}
if idx >= 0 && idx < len(result.ToolCalls) {
result.ToolCalls[idx].ProviderMetadata = c.adapter.ExtractToolCallMetadata(rawToolCall)
}
}
}
// removeThinkingContent 移除思考模型输出中的 <think></think> 思考过程
// 仅当内容以 <think> 开头时才处理
func removeThinkingContent(content string) string {
const thinkStartTag = "<think>"
const thinkEndTag = "</think>"
trimmed := strings.TrimSpace(content)
if !strings.HasPrefix(trimmed, thinkStartTag) {
return content
}
// 查找最后一个 </think> 标签(处理嵌套情况)
if lastEndIdx := strings.LastIndex(trimmed, thinkEndTag); lastEndIdx != -1 {
if result := strings.TrimSpace(trimmed[lastEndIdx+len(thinkEndTag):]); result != "" {
return result
}
return ""
}
return "" // 未找到 </think>,可能思考内容过长被截断,返回空字符串
}
// processStream 处理 OpenAI SDK 流式响应
func (c *RemoteAPIChat) processStream(
ctx context.Context,
stream *openai.ChatCompletionStream,
streamChan chan types.StreamResponse,
dumper *streamPacketDumper,
) {
defer close(streamChan)
defer stream.Close()
state := newStreamState()
for {
response, err := stream.Recv()
if err != nil {
if err == io.EOF {
logUsage(ctx, c.modelName, state.usage)
toolCalls := state.buildOrderedToolCalls()
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeAnswer,
Content: "",
Done: true,
ToolCalls: toolCalls,
Usage: state.usage,
FinishReason: state.lastFinishReason,
}
} else {
logger.Errorf(ctx, "Stream read error: %v (tool_calls_assembled=%d)",
err, len(state.toolCallMap))
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeError,
Content: err.Error(),
Done: true,
ToolCalls: state.buildOrderedToolCalls(),
Usage: state.usage,
FinishReason: types.FinishReasonIncomplete,
}
}
return
}
if dumper != nil {
dumper.WritePacket(response)
}
if response.Usage != nil {
usage := tokenUsageFromOpenAI(*response.Usage, c.provider)
state.usage = &usage
}
if len(response.Choices) > 0 {
c.processStreamDelta(ctx, &response.Choices[0], state, streamChan, response.Choices[0].Delta.ReasoningContent)
}
}
}
// processRawHTTPStream 处理原始 HTTP 流式响应
func (c *RemoteAPIChat) processRawHTTPStream(
ctx context.Context,
resp *http.Response,
streamChan chan types.StreamResponse,
dumper *streamPacketDumper,
) {
defer close(streamChan)
defer resp.Body.Close()
state := newStreamState()
reader := NewSSEReader(resp.Body)
for {
event, err := reader.ReadEvent()
if err != nil {
if err == io.EOF {
logUsage(ctx, c.modelName, state.usage)
toolCalls := state.buildOrderedToolCalls()
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeAnswer,
Content: "",
Done: true,
ToolCalls: toolCalls,
Usage: state.usage,
}
} else {
logger.Errorf(ctx, "Stream read error: %v (tool_calls_assembled=%d)",
err, len(state.toolCallMap))
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeError,
Content: err.Error(),
Done: true,
ToolCalls: state.buildOrderedToolCalls(),
Usage: state.usage,
FinishReason: types.FinishReasonIncomplete,
}
}
return
}
if event == nil {
continue
}
if event.Done {
logUsage(ctx, c.modelName, state.usage)
toolCalls := state.buildOrderedToolCalls()
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeAnswer,
Content: "",
Done: true,
ToolCalls: toolCalls,
Usage: state.usage,
}
return
}
if event.Data == nil {
continue
}
if dumper != nil {
// 保留上游 SSE data 行的原始 JSON,不经过中间结构体裁剪。
raw := make([]byte, len(event.Data))
copy(raw, event.Data)
dumper.WritePacketRaw(raw)
}
// 使用局部结构体进行一次性解析,同时捕捉标准字段和 vLLM 的 reasoning 字段,避免性能损失
var streamResp struct {
openai.ChatCompletionStreamResponse
Choices []struct {
Index int `json:"index"`
Delta struct {
openai.ChatCompletionStreamChoiceDelta
Reasoning string `json:"reasoning,omitempty"`
} `json:"delta"`
FinishReason openai.FinishReason `json:"finish_reason"`
} `json:"choices"`
}
if err := json.Unmarshal(event.Data, &streamResp); err != nil {
logger.Errorf(ctx, "Failed to parse stream response: %v", err)
continue
}
if streamResp.Usage != nil {
usage := tokenUsageFromOpenAI(*streamResp.Usage, c.provider)
applyRawPromptCacheUsage(event.Data, &usage)
state.usage = &usage
}
if len(streamResp.Choices) > 0 {
choice := streamResp.Choices[0]
// 统一获取逻辑(支持标准和 vLLM 两种路径)
reasoning := choice.Delta.Reasoning
if reasoning == "" {
reasoning = choice.Delta.ReasoningContent
}
// 构造一个标准 SDK 兼容的 choice 对象传给下游,保证现有逻辑完全不动
sdkChoice := openai.ChatCompletionStreamChoice{
Index: choice.Index,
Delta: choice.Delta.ChatCompletionStreamChoiceDelta,
FinishReason: choice.FinishReason,
}
c.applyStreamToolCallMetadata(event.Data, state)
c.processStreamDelta(ctx, &sdkChoice, state, streamChan, reasoning)
}
}
}
func (c *RemoteAPIChat) applyStreamToolCallMetadata(data []byte, state *streamState) {
if state == nil {
return
}
var raw struct {
Choices []struct {
Delta struct {
ToolCalls []json.RawMessage `json:"tool_calls,omitempty"`
} `json:"delta"`
} `json:"choices"`
}
if err := json.Unmarshal(data, &raw); err != nil || len(raw.Choices) == 0 {
return
}
for _, rawToolCall := range raw.Choices[0].Delta.ToolCalls {
metadata := c.adapter.ExtractToolCallMetadata(rawToolCall)
if len(metadata) == 0 {
continue
}
var indexed struct {
Index *int `json:"index,omitempty"`
}
_ = json.Unmarshal(rawToolCall, &indexed)
idx := 0
if indexed.Index != nil {
idx = *indexed.Index
}
state.setToolCallProviderMetadata(idx, metadata)
}
}
// streamState 流式处理状态
type streamState struct {
thinkingEmitter
toolCallMap map[int]*types.LLMToolCall
lastFunctionName map[int]string
nameNotified map[int]bool
fieldExtractors map[int]*jsonFieldExtractor // per tool-call-index extractors for streaming field extraction
fileProgress map[int]*sandboxFileProgress
usage *types.TokenUsage // captured from the final stream chunk when include_usage is enabled
lastFinishReason string // last observed finish_reason for EOF handler fallback
// Diagnostic flags (fire-once) used to log earliest signals of tool_call
// presence/absence at the OpenAI-protocol level. These are independent of
// the higher-level ResponseTypeToolCall marker (which only fires once
// function name has stabilized) and let us distinguish between
// (A) no tool_calls field ever observed (true natural-stop), and
// (B) tool_calls field observed but marker not yet emitted.
firstToolCallSeen bool // true once any delta carried tool_calls
noToolCallStopLogged bool // true once we logged "stop without tool_calls"
firstContentSeen bool // true once delta.Content first appeared
firstReasoningSeen bool // true once reasoning_content first appeared
streamStartedAt time.Time
}
func newStreamState() *streamState {
return &streamState{
toolCallMap: make(map[int]*types.LLMToolCall),
lastFunctionName: make(map[int]string),
nameNotified: make(map[int]bool),
fieldExtractors: make(map[int]*jsonFieldExtractor),
fileProgress: make(map[int]*sandboxFileProgress),
streamStartedAt: time.Now(),
}
}
// elapsedMs returns the milliseconds elapsed since the stream state was
// initialized. Used to attach time-since-stream-start to fire-once diagnostic
// logs so a single grep can reveal the temporal layout of a single stream
// (TTFC / TTFT / first-tool-call / natural-stop confirmation, etc).
func (s *streamState) elapsedMs() int64 {
if s.streamStartedAt.IsZero() {
return 0
}
return time.Since(s.streamStartedAt).Milliseconds()
}
func (s *streamState) buildOrderedToolCalls() []types.LLMToolCall {
if len(s.toolCallMap) == 0 {
return nil
}
result := make([]types.LLMToolCall, 0, len(s.toolCallMap))
for i := 0; i < len(s.toolCallMap); i++ {
if tc, ok := s.toolCallMap[i]; ok && tc != nil {
result = append(result, *tc)
}
}
if len(result) == 0 {
return nil
}
return result
}
func (s *streamState) setToolCallProviderMetadata(index int, metadata types.ToolCallMetadata) {
if len(metadata) == 0 {
return
}
toolCallEntry, exists := s.toolCallMap[index]
if !exists || toolCallEntry == nil {
toolCallEntry = &types.LLMToolCall{
Type: "function",
Function: types.FunctionCall{
Name: "",
Arguments: "",
},
}
s.toolCallMap[index] = toolCallEntry
}
toolCallEntry.ProviderMetadata = metadata
}
// processStreamDelta 处理流式响应的单个 delta
func (c *RemoteAPIChat) processStreamDelta(
ctx context.Context,
choice *openai.ChatCompletionStreamChoice,
state *streamState,
streamChan chan types.StreamResponse,
reasoningContent string,
) {
delta := choice.Delta
isDone := string(choice.FinishReason) != ""
// Track finish_reason for EOF handler fallback
if isDone {
state.lastFinishReason = string(choice.FinishReason)
}
// 处理 tool calls
if len(delta.ToolCalls) > 0 {
c.processToolCallsDelta(ctx, delta.ToolCalls, state, streamChan)
}
// Earliest reliable "no tool_calls" signal at the OpenAI-protocol level:
// finish_reason=stop arrived AND we never observed a tool_calls field on
// any prior delta. Logged once per stream so callers can grep for the
// natural-stop entry point without waiting for the higher-level summary.
if isDone &&
string(choice.FinishReason) == "stop" &&
!state.firstToolCallSeen &&
!state.noToolCallStopLogged {
logger.Infof(ctx, "[LLM Stream] Natural-stop at OpenAI layer "+
"(finish=stop, tool_calls field never observed, thinking_seen=%t, "+
"first_content_seen=%t, elapsed_ms=%d)",
state.active, state.firstContentSeen, state.elapsedMs())
state.noToolCallStopLogged = true
}
// 发送思考内容(ReasoningContent,支持 DeepSeek 等模型)
if reasoningContent != "" {
// Earliest reasoning_content signal at the OpenAI-protocol level. Fired
// once per stream so we can distinguish "model emitted thinking before
// answer" vs "model never produced thinking" when triaging logs.
if !state.firstReasoningSeen {
state.firstReasoningSeen = true
logger.Infof(ctx, "[LLM Stream] First reasoning_content at OpenAI layer "+
"(len=%d, preview=%q, elapsed_ms=%d)",
len(reasoningContent), truncateForDebug(reasoningContent, 80), state.elapsedMs())
}
state.emit(streamChan, reasoningContent)
}
// 发送回答内容
if delta.Content != "" {
// Earliest delta.Content signal at the OpenAI-protocol level. Fired once
// per stream so we can measure TTFC (time-to-first-content) and tell
// "answer started before any tool_call" from "tool_call came first".
if !state.firstContentSeen {
state.firstContentSeen = true
logger.Infof(ctx, "[LLM Stream] First delta.Content at OpenAI layer "+
"(len=%d, preview=%q, tool_call_seen=%t, thinking_seen=%t, elapsed_ms=%d)",
len(delta.Content), truncateForDebug(delta.Content, 80),
state.firstToolCallSeen, state.firstReasoningSeen, state.elapsedMs())
}
// If we had thinking content and this is the first answer chunk,
// send a thinking done event first.
state.finish(streamChan)
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeAnswer,
Content: delta.Content,
Done: isDone,
ToolCalls: state.buildOrderedToolCalls(),
FinishReason: string(choice.FinishReason),
}
}
if isDone && len(state.toolCallMap) > 0 {
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeAnswer,
Content: "",
Done: true,
ToolCalls: state.buildOrderedToolCalls(),
FinishReason: string(choice.FinishReason),
}
}
// Ensure thinking done is sent when stream finishes without any answer content
// (e.g., model only produced reasoning then hit finish_reason with empty content).
if isDone {
state.finish(streamChan)
}
// Catch-all: isDone but none of the above branches sent a response with
// FinishReason (empty content, no tool calls, no thinking). This prevents
// the finish_reason from being lost in the streaming pipeline.
if isDone && delta.Content == "" && len(state.toolCallMap) == 0 {
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeAnswer,
Done: true,
FinishReason: string(choice.FinishReason),
}
}
}
// processToolCallsDelta 处理 tool calls 的增量更新
func (c *RemoteAPIChat) processToolCallsDelta(
ctx context.Context,
toolCalls []openai.ToolCall,
state *streamState,
streamChan chan types.StreamResponse,
) {
// Earliest signal at the OpenAI-protocol level that this stream will
// produce at least one tool call. Fires *before* the function name has
// stabilized, i.e. earlier than the higher-level ResponseTypeToolCall
// marker downstream consumers see. Useful for distinguishing
// "tool_calls field arrived but marker not yet emitted" from
// "tool_calls field truly absent" when triaging stream behavior.
if !state.firstToolCallSeen && len(toolCalls) > 0 {
state.firstToolCallSeen = true
var firstID, firstName string
for _, tc := range toolCalls {
if tc.ID != "" {
firstID = tc.ID
}
if tc.Function.Name != "" {
firstName = tc.Function.Name
}
if firstID != "" || firstName != "" {
break
}
}
logger.Infof(ctx, "[LLM Stream] First tool_calls delta at OpenAI layer "+
"(count=%d, first_id=%q, first_name=%q, "+
"first_content_seen=%t, thinking_seen=%t, elapsed_ms=%d)",
len(toolCalls), firstID, firstName,
state.firstContentSeen, state.firstReasoningSeen, state.elapsedMs())
}
for _, tc := range toolCalls {
var toolCallIndex int
if tc.Index != nil {
toolCallIndex = *tc.Index
}
toolCallEntry, exists := state.toolCallMap[toolCallIndex]
if !exists || toolCallEntry == nil {
toolCallEntry = &types.LLMToolCall{
Type: string(tc.Type),
Function: types.FunctionCall{
Name: "",
Arguments: "",
},
}
state.toolCallMap[toolCallIndex] = toolCallEntry
}
if tc.ID != "" {
toolCallEntry.ID = tc.ID
}
if tc.Type != "" {
toolCallEntry.Type = string(tc.Type)
}
if tc.Function.Name != "" {
// 防御性校验:解决部分供应商(如vLLM Ascend等)在每个流 Chunk 中重复发送完整工具名的问题。
// 如果当前已存名字与新收到名字一致,则视为冗余重复,不进行叠加。
if toolCallEntry.Function.Name != tc.Function.Name {
toolCallEntry.Function.Name += tc.Function.Name
}
}
argsUpdated := false
if tc.Function.Arguments != "" {
toolCallEntry.Function.Arguments += tc.Function.Arguments
argsUpdated = true
}
currName := toolCallEntry.Function.Name
var progressArgs map[string]any
if isSandboxMutationTool(currName) && argsUpdated && tc.Function.Arguments != "" {
prog := state.fileProgress[toolCallIndex]
if prog == nil {
prog = newSandboxFileProgress(currName)
state.fileProgress[toolCallIndex] = prog
}
if payload, ok := prog.Feed(tc.Function.Arguments); ok {
progressArgs = payload
}
}
if currName != "" &&
currName == state.lastFunctionName[toolCallIndex] &&
argsUpdated &&
!state.nameNotified[toolCallIndex] &&
toolCallEntry.ID != "" {
data := map[string]interface{}{
"tool_name": currName,
"tool_call_id": toolCallEntry.ID,
}
if progressArgs != nil {
data["arguments"] = progressArgs
}
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeToolCall,
Content: "",
Done: false,
Data: data,
}
state.nameNotified[toolCallIndex] = true
progressArgs = nil
} else if progressArgs != nil && toolCallEntry.ID != "" && currName != "" {
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeToolCall,
Content: "",
Done: false,
Data: map[string]interface{}{
"tool_name": currName,
"tool_call_id": toolCallEntry.ID,
"arguments": progressArgs,
},
}
}
state.lastFunctionName[toolCallIndex] = currName
// Stream thinking tool's thought field as thinking-type chunks
if toolCallEntry.Function.Name == "thinking" && argsUpdated {
extractor, exists := state.fieldExtractors[toolCallIndex]
if !exists {
extractor = newJSONFieldExtractor("thought")
state.fieldExtractors[toolCallIndex] = extractor
}
thoughtChunk := extractor.Feed(tc.Function.Arguments)
if thoughtChunk != "" {
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeThinking,
Content: thoughtChunk,
Done: false,
Data: map[string]interface{}{
"source": "thinking_tool",
"tool_call_id": toolCallEntry.ID,
},
}
}
}
}
}
+374
View File
@@ -0,0 +1,374 @@
package chat
import (
"context"
"crypto/sha256"
"encoding/hex"
"encoding/json"
"fmt"
"net/http"
"strings"
"github.com/Tencent/WeKnora/internal/models/provider"
"github.com/Tencent/WeKnora/internal/types"
"github.com/sashabaranov/go-openai"
)
// FingerprintPromptPrefix returns a short, non-reversible identifier suitable
// for logs and cache routing. Raw prompts must never be used as metric labels.
func FingerprintPromptPrefix(parts ...string) string {
h := sha256.New()
for _, part := range parts {
_, _ = h.Write([]byte(part))
_, _ = h.Write([]byte{0})
}
return hex.EncodeToString(h.Sum(nil))[:16]
}
// PromptPrefixFingerprint hashes the stable portion common to normal chat and
// agent requests: leading system messages plus the deterministic tool schema.
// Dynamic conversation/user messages intentionally do not participate.
func PromptPrefixFingerprint(messages []Message, opts *ChatOptions) string {
type stablePrefix struct {
System []Message `json:"system,omitempty"`
Tools []Tool `json:"tools,omitempty"`
}
prefix := stablePrefix{}
for _, message := range messages {
if message.Role != "system" {
break
}
prefix.System = append(prefix.System, message)
}
if opts != nil {
prefix.Tools = opts.Tools
}
data, _ := json.Marshal(prefix)
return FingerprintPromptPrefix(string(data))
}
// BuildPromptCacheKey derives an opaque process-local coordination key.
// Tenant and model identifiers are hashed rather than retained in memory.
func BuildPromptCacheKey(tenantID uint64, modelID, purpose, prefixFingerprint string) string {
return "wk-" + FingerprintPromptPrefix(
fmt.Sprintf("%d", tenantID), modelID, purpose, prefixFingerprint,
)
}
func providerCacheAccountingStatus(name provider.ProviderName) types.PromptCacheStatus {
switch name {
case provider.ProviderOpenAI,
provider.ProviderAzureOpenAI,
provider.ProviderDeepSeek,
provider.ProviderAliyun,
provider.ProviderAnthropic:
return types.PromptCacheStatusUnreported
default:
return types.PromptCacheStatusUnsupported
}
}
func tokenUsageFromOpenAI(usage openai.Usage, providerName provider.ProviderName) types.TokenUsage {
u := types.TokenUsage{
PromptTokens: usage.PromptTokens,
CompletionTokens: usage.CompletionTokens,
TotalTokens: usage.TotalTokens,
}
if usage.PromptTokensDetails != nil {
read := usage.PromptTokensDetails.CachedTokens
u.SetPromptCacheUsage(read, 0, max(0, usage.PromptTokens-read), true)
return u
}
if providerCacheAccountingStatus(providerName) == types.PromptCacheStatusUnsupported {
u.MarkPromptCacheUnsupported()
} else {
u.SetPromptCacheUsage(0, 0, 0, false)
}
return u
}
// cachedTokens is retained as the nil-safe primitive used by older callers
// and focused tests; normalization happens in tokenUsageFromOpenAI.
func cachedTokens(details *openai.PromptTokensDetails) int {
if details == nil {
return 0
}
return details.CachedTokens
}
type rawPromptCacheUsage struct {
Usage struct {
PromptTokens int `json:"prompt_tokens"`
PromptCacheHit *int `json:"prompt_cache_hit_tokens"`
PromptCacheMiss *int `json:"prompt_cache_miss_tokens"`
CacheReadInput *int `json:"cache_read_input_tokens"`
CacheCreationInput *int `json:"cache_creation_input_tokens"`
PromptTokensDetails *struct {
CachedTokens *int `json:"cached_tokens"`
CacheWriteTokens *int `json:"cache_write_tokens"`
} `json:"prompt_tokens_details"`
} `json:"usage"`
}
// applyRawPromptCacheUsage captures native fields discarded by the pinned
// OpenAI-compatible SDK (notably DeepSeek hit/miss counters).
func applyRawPromptCacheUsage(data []byte, usage *types.TokenUsage) {
if usage == nil || len(data) == 0 {
return
}
var raw rawPromptCacheUsage
if json.Unmarshal(data, &raw) != nil {
return
}
if raw.Usage.PromptCacheHit != nil || raw.Usage.PromptCacheMiss != nil {
read := valueOrZero(raw.Usage.PromptCacheHit)
miss := valueOrZero(raw.Usage.PromptCacheMiss)
usage.SetPromptCacheUsage(read, 0, miss, true)
return
}
if raw.Usage.CacheReadInput != nil || raw.Usage.CacheCreationInput != nil {
read := valueOrZero(raw.Usage.CacheReadInput)
write := valueOrZero(raw.Usage.CacheCreationInput)
usage.SetPromptCacheUsage(read, write, max(0, usage.PromptTokens-read), true)
return
}
if details := raw.Usage.PromptTokensDetails; details != nil {
read := valueOrZero(details.CachedTokens)
write := valueOrZero(details.CacheWriteTokens)
usage.SetPromptCacheUsage(read, write, max(0, usage.PromptTokens-read), true)
}
}
func valueOrZero(value *int) int {
if value == nil {
return 0
}
return *value
}
// CacheRetention is the prompt-cache TTL preference. Empty means short (the
// default 5-minute provider cache). Compaction/summarization uses none so a
// different prompt prefix does not occupy the session's cache slot.
type CacheRetention string
const (
CacheRetentionNone CacheRetention = "none"
CacheRetentionShort CacheRetention = "short"
CacheRetentionLong CacheRetention = "long"
)
const openAIPromptCacheKeyMaxLength = 64
func clampPromptCacheKey(key string) string {
if key == "" {
return ""
}
runes := []rune(key)
if len(runes) <= openAIPromptCacheKeyMaxLength {
return key
}
return string(runes[:openAIPromptCacheKeyMaxLength])
}
func resolveCacheRetention(opts *ChatOptions) CacheRetention {
if opts != nil && opts.CacheRetention != "" {
return opts.CacheRetention
}
return CacheRetentionShort
}
func promptCacheSessionID(ctx context.Context, opts *ChatOptions) string {
if opts != nil && opts.PromptCacheKey != "" {
return clampPromptCacheKey(opts.PromptCacheKey)
}
if sessionID, ok := types.SessionIDFromContext(ctx); ok {
return clampPromptCacheKey(sessionID)
}
return ""
}
type promptCachePolicy struct {
sendKey bool
sendCacheControl bool
sendAffinity bool
}
func promptCachePolicyFor(name provider.ProviderName, baseURL string) promptCachePolicy {
switch name {
case provider.ProviderOpenAI, provider.ProviderAzureOpenAI, provider.ProviderOpenRouter:
return promptCachePolicy{sendKey: true, sendAffinity: true}
case provider.ProviderAliyun:
return promptCachePolicy{sendCacheControl: true}
case provider.ProviderAnthropic:
return promptCachePolicy{sendCacheControl: true}
}
if strings.Contains(baseURL, "api.openai.com") {
return promptCachePolicy{sendKey: true, sendAffinity: true}
}
return promptCachePolicy{}
}
type cacheControlMarker struct {
Type string `json:"type"`
TTL string `json:"ttl,omitempty"`
}
func cacheControlFor(retention CacheRetention, longTTL string) *cacheControlMarker {
if retention == CacheRetentionNone {
return nil
}
marker := &cacheControlMarker{Type: "ephemeral"}
if retention == CacheRetentionLong && longTTL != "" {
marker.TTL = longTTL
}
return marker
}
// applyPromptCacheToJSONBody injects provider cache routing and breakpoints
// into an already-shaped OpenAI-compatible request object. Returns the (possibly
// rewritten) body and whether the caller must send it via raw HTTP because the
// SDK struct cannot carry these fields.
func applyPromptCacheToJSONBody(
body any,
policy promptCachePolicy,
sessionID string,
retention CacheRetention,
) (any, bool, error) {
if retention == CacheRetentionNone {
return body, false, nil
}
if !policy.sendKey && !policy.sendCacheControl {
return body, false, nil
}
data, err := json.Marshal(body)
if err != nil {
return nil, false, err
}
var payload map[string]any
if err := json.Unmarshal(data, &payload); err != nil {
return nil, false, err
}
rewritten := false
if policy.sendKey && sessionID != "" {
payload["prompt_cache_key"] = sessionID
if retention == CacheRetentionLong {
payload["prompt_cache_retention"] = "24h"
}
rewritten = true
}
if policy.sendCacheControl {
marker := cacheControlFor(retention, "1h")
if marker != nil {
applyCacheControlBreakpoints(payload, marker)
rewritten = true
}
}
if !rewritten {
return body, false, nil
}
return payload, true, nil
}
func applyCacheControlBreakpoints(payload map[string]any, marker *cacheControlMarker) {
if marker == nil {
return
}
applyCacheControlToInstructionMessages(payload["messages"], marker)
applyCacheControlToLastTool(payload["tools"], marker)
applyCacheControlToLastConversationMessage(payload["messages"], marker)
}
func applyCacheControlToInstructionMessages(raw any, marker *cacheControlMarker) {
messages, ok := raw.([]any)
if !ok {
return
}
for _, item := range messages {
msg, ok := item.(map[string]any)
if !ok {
continue
}
role, _ := msg["role"].(string)
if role == "system" || role == "developer" {
addCacheControlToMessageContent(msg, marker)
return
}
}
}
func applyCacheControlToLastConversationMessage(raw any, marker *cacheControlMarker) {
messages, ok := raw.([]any)
if !ok {
return
}
for i := len(messages) - 1; i >= 0; i-- {
msg, ok := messages[i].(map[string]any)
if !ok {
continue
}
role, _ := msg["role"].(string)
if role == "user" || role == "assistant" || role == "tool" {
if addCacheControlToMessageContent(msg, marker) {
return
}
}
}
}
func applyCacheControlToLastTool(raw any, marker *cacheControlMarker) {
tools, ok := raw.([]any)
if !ok || len(tools) == 0 {
return
}
last, ok := tools[len(tools)-1].(map[string]any)
if !ok {
return
}
last["cache_control"] = marker
}
func addCacheControlToMessageContent(msg map[string]any, marker *cacheControlMarker) bool {
content, ok := msg["content"]
if !ok || content == nil {
return false
}
if text, ok := content.(string); ok {
if text == "" {
return false
}
msg["content"] = []any{
map[string]any{
"type": "text",
"text": text,
"cache_control": marker,
},
}
return true
}
parts, ok := content.([]any)
if !ok {
return false
}
for i := len(parts) - 1; i >= 0; i-- {
part, ok := parts[i].(map[string]any)
if !ok {
continue
}
if partType, _ := part["type"].(string); partType == "text" || partType == "tool_result" {
part["cache_control"] = marker
return true
}
}
return false
}
func attachPromptCacheHeaders(req *http.Request, policy promptCachePolicy, sessionID string) {
if req == nil || !policy.sendAffinity || sessionID == "" {
return
}
req.Header.Set("session_id", sessionID)
req.Header.Set("x-client-request-id", sessionID)
req.Header.Set("x-session-affinity", sessionID)
}
@@ -0,0 +1,117 @@
package chat
import (
"context"
"encoding/json"
"net/http"
"strings"
"testing"
"github.com/Tencent/WeKnora/internal/models/provider"
"github.com/Tencent/WeKnora/internal/types"
"github.com/sashabaranov/go-openai"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestClampPromptCacheKey(t *testing.T) {
assert.Equal(t, "", clampPromptCacheKey(""))
assert.Equal(t, "sess-1", clampPromptCacheKey("sess-1"))
long := strings.Repeat("a", 80)
got := clampPromptCacheKey(long)
assert.Equal(t, 64, len([]rune(got)))
assert.Equal(t, strings.Repeat("a", 64), got)
}
func TestApplyPromptCacheToJSONBody_OpenAIKey(t *testing.T) {
req := openai.ChatCompletionRequest{
Model: "gpt-4o",
Messages: []openai.ChatCompletionMessage{
{Role: "system", Content: "sys"},
{Role: "user", Content: "hi"},
},
}
policy := promptCachePolicyFor(provider.ProviderOpenAI, "")
body, forceRaw, err := applyPromptCacheToJSONBody(req, policy, "sess-abc", CacheRetentionShort)
require.NoError(t, err)
require.True(t, forceRaw)
payload := body.(map[string]any)
assert.Equal(t, "sess-abc", payload["prompt_cache_key"])
assert.NotContains(t, payload, "prompt_cache_retention")
}
func TestApplyPromptCacheToJSONBody_OpenAILongRetention(t *testing.T) {
req := openai.ChatCompletionRequest{Model: "gpt-4o"}
policy := promptCachePolicyFor(provider.ProviderOpenAI, "")
body, forceRaw, err := applyPromptCacheToJSONBody(req, policy, "sess-abc", CacheRetentionLong)
require.NoError(t, err)
require.True(t, forceRaw)
payload := body.(map[string]any)
assert.Equal(t, "24h", payload["prompt_cache_retention"])
}
func TestApplyPromptCacheToJSONBody_NoneLeavesBodyUntouched(t *testing.T) {
req := &openai.ChatCompletionRequest{Model: "gpt-4o"}
policy := promptCachePolicyFor(provider.ProviderOpenAI, "")
body, forceRaw, err := applyPromptCacheToJSONBody(req, policy, "sess-abc", CacheRetentionNone)
require.NoError(t, err)
assert.False(t, forceRaw)
assert.Equal(t, req, body)
}
func TestApplyPromptCacheToJSONBody_AliyunCacheControlBreakpoints(t *testing.T) {
req := openai.ChatCompletionRequest{
Model: "qwen-plus",
Messages: []openai.ChatCompletionMessage{
{Role: "system", Content: "stable system"},
{Role: "user", Content: "turn question"},
},
Tools: []openai.Tool{{
Type: openai.ToolTypeFunction,
Function: &openai.FunctionDefinition{Name: "search", Description: "d", Parameters: json.RawMessage(`{}`)},
}},
}
policy := promptCachePolicyFor(provider.ProviderAliyun, "")
body, forceRaw, err := applyPromptCacheToJSONBody(req, policy, "", CacheRetentionShort)
require.NoError(t, err)
require.True(t, forceRaw)
wire, err := json.Marshal(body)
require.NoError(t, err)
var payload map[string]any
require.NoError(t, json.Unmarshal(wire, &payload))
messages := payload["messages"].([]any)
sys := messages[0].(map[string]any)
sysParts := sys["content"].([]any)
sysCache := sysParts[0].(map[string]any)["cache_control"].(map[string]any)
assert.Equal(t, "ephemeral", sysCache["type"])
user := messages[1].(map[string]any)
userParts := user["content"].([]any)
userCache := userParts[0].(map[string]any)["cache_control"].(map[string]any)
assert.Equal(t, "ephemeral", userCache["type"])
tools := payload["tools"].([]any)
lastTool := tools[len(tools)-1].(map[string]any)
toolCache := lastTool["cache_control"].(map[string]any)
assert.Equal(t, "ephemeral", toolCache["type"])
}
func TestBuildOutbound_OpenAIPromptCacheKeyFromSession(t *testing.T) {
c := newOutboundChat(t, string(provider.ProviderOpenAI), "gpt-4o", nil)
ctx := types.WithSessionID(context.Background(), "sess-live")
body, _, useRaw, err := c.buildOutbound(ctx, []Message{{Role: "user", Content: "hi"}}, &ChatOptions{}, false)
require.NoError(t, err)
require.True(t, useRaw)
payload := body.(map[string]any)
assert.Equal(t, "sess-live", payload["prompt_cache_key"])
}
func TestAttachPromptCacheHeaders(t *testing.T) {
req, err := http.NewRequest(http.MethodPost, "https://api.openai.com/v1/chat/completions", nil)
require.NoError(t, err)
attachPromptCacheHeaders(req, promptCachePolicy{sendAffinity: true}, "sess-1")
assert.Equal(t, "sess-1", req.Header.Get("session_id"))
assert.Equal(t, "sess-1", req.Header.Get("x-client-request-id"))
assert.Equal(t, "sess-1", req.Header.Get("x-session-affinity"))
}
+305
View File
@@ -0,0 +1,305 @@
package chat
import (
"encoding/json"
"net/http"
"strings"
"github.com/Tencent/WeKnora/internal/models/provider"
modelutils "github.com/Tencent/WeKnora/internal/models/utils"
"github.com/Tencent/WeKnora/internal/types"
"github.com/google/uuid"
"github.com/sashabaranov/go-openai"
)
// authCreds carries the credentials a providerAdapter needs to authenticate a
// raw HTTP request. APIKey covers the common Bearer / api-key cases; AppID and
// AppSecret are only used by signing providers (WeKnoraCloud).
type authCreds struct {
APIKey string
AppID string
AppSecret string
}
// providerAdapter captures everything provider-specific about an
// OpenAI-compatible chat backend. Every method has a sensible default on
// baseProvider, so a new provider is added by embedding baseProvider and
// overriding only the one or two methods that actually differ.
type providerAdapter interface {
// Name is the provider this adapter handles.
Name() provider.ProviderName
// Matches reports whether this adapter applies to the given model name.
// Used for sub-provider routing (e.g. Qwen thinking models within Aliyun,
// reasoning models within OpenAI). Default: true.
Matches(model string) bool
// Thinking is how this provider encodes ChatOptions.Thinking. Default: none.
Thinking() ThinkingStrategy
// ShapeRequest applies in-place parameter quirks to the standard request
// (stripping unsupported fields, pinning temperature, …). Default: noop.
ShapeRequest(req *openai.ChatCompletionRequest, opts *ChatOptions, isStream bool)
// TransformMessages rewrites the converted messages (e.g. downgrading
// multi-content to plain text). Default: identity.
TransformMessages(msgs []openai.ChatCompletionMessage) []openai.ChatCompletionMessage
// Endpoint overrides the request URL. Empty string means the standard
// "<baseURL>/chat/completions" handled by the caller. Default: "".
Endpoint(baseURL, modelID string, isStream bool) string
// Auth sets authentication headers on a raw HTTP request. Default: Bearer.
Auth(req *http.Request, creds authCreds, body []byte)
// ForceRawHTTP forces the raw HTTP path even when the body is standard
// (needed by providers that must sign the exact request bytes). Default: false.
ForceRawHTTP() bool
// ExtractToolCallMetadata captures provider-specific state from a raw
// OpenAI-compatible tool_call object. Default: nil.
ExtractToolCallMetadata(raw json.RawMessage) types.ToolCallMetadata
// InjectToolCallMetadata writes provider-specific state back into an outbound
// OpenAI-compatible tool_call object. Default: noop.
InjectToolCallMetadata(toolCall map[string]any, metadata types.ToolCallMetadata)
}
// baseProvider supplies the default behavior for every providerAdapter method.
// It is also the fallback returned by resolveProvider for unknown providers:
// Bearer auth, standard endpoint, no thinking, no request shaping.
type baseProvider struct{}
func (baseProvider) Name() provider.ProviderName { return "" }
func (baseProvider) Matches(string) bool { return true }
func (baseProvider) Thinking() ThinkingStrategy { return noThinking{} }
func (baseProvider) ShapeRequest(*openai.ChatCompletionRequest, *ChatOptions, bool) {}
func (baseProvider) TransformMessages(msgs []openai.ChatCompletionMessage) []openai.ChatCompletionMessage {
return msgs
}
func (baseProvider) Endpoint(string, string, bool) string { return "" }
func (baseProvider) Auth(req *http.Request, creds authCreds, _ []byte) {
req.Header.Set("Authorization", "Bearer "+creds.APIKey)
}
func (baseProvider) ForceRawHTTP() bool { return false }
func (baseProvider) ExtractToolCallMetadata(json.RawMessage) types.ToolCallMetadata {
return nil
}
func (baseProvider) InjectToolCallMetadata(map[string]any, types.ToolCallMetadata) {}
// --- WeKnoraCloud: custom endpoint + request signing + multi-content downgrade ---
type weKnoraCloudProvider struct{ baseProvider }
func (weKnoraCloudProvider) Name() provider.ProviderName { return provider.ProviderWeKnoraCloud }
func (weKnoraCloudProvider) Endpoint(baseURL, _ string, _ bool) string {
return strings.TrimRight(baseURL, "/") + "/api/v1/chat/completions"
}
func (weKnoraCloudProvider) ForceRawHTTP() bool { return true }
func (weKnoraCloudProvider) Auth(req *http.Request, creds authCreds, body []byte) {
requestID := uuid.NewString()
headers := modelutils.Sign(creds.AppID, creds.AppSecret, requestID, string(body))
for k, v := range headers {
req.Header.Set(k, v)
}
}
// TransformMessages downgrades MultiContent to plain text while preserving
// tool_calls / tool_call_id / name so the function-calling protocol keeps working.
func (weKnoraCloudProvider) TransformMessages(messages []openai.ChatCompletionMessage) []openai.ChatCompletionMessage {
result := make([]openai.ChatCompletionMessage, 0, len(messages))
for _, m := range messages {
msg := m
if msg.Content == "" && len(msg.MultiContent) > 0 {
var textParts []string
for _, part := range msg.MultiContent {
if part.Type == openai.ChatMessagePartTypeText && part.Text != "" {
textParts = append(textParts, part.Text)
}
}
msg.Content = strings.Join(textParts, "\n")
msg.MultiContent = nil
}
result = append(result, msg)
}
return result
}
// --- Aliyun Qwen thinking models: enable_thinking (always sent, forced off non-stream) ---
type qwenThinkingProvider struct{ baseProvider }
func (qwenThinkingProvider) Name() provider.ProviderName { return provider.ProviderAliyun }
func (qwenThinkingProvider) Matches(model string) bool { return provider.IsQwenThinkingModel(model) }
func (qwenThinkingProvider) Thinking() ThinkingStrategy {
return enableThinking{alwaysSend: true, disableOnNonStream: true}
}
// --- LKEAP: thinking via { "thinking": { "type": ... } }, only for DeepSeek V3.x ---
// R1 series enables chain-of-thought by default and is left untouched (falls
// back to baseProvider). See https://cloud.tencent.com/document/product/1772/115963
type lkeapProvider struct{ baseProvider }
func (lkeapProvider) Name() provider.ProviderName { return provider.ProviderLKEAP }
func (lkeapProvider) Matches(model string) bool {
return strings.Contains(strings.ToLower(model), "deepseek-v3")
}
func (lkeapProvider) Thinking() ThinkingStrategy { return thinkingTypeField{} }
// --- DeepSeek: does not support tool_choice ---
type deepseekProvider struct{ baseProvider }
func (deepseekProvider) Name() provider.ProviderName { return provider.ProviderDeepSeek }
func (deepseekProvider) Thinking() ThinkingStrategy { return thinkingTypeField{} }
// Native DeepSeek cache counters are not represented by go-openai v1.41.2;
// use the raw path so prompt_cache_hit_tokens/miss_tokens remain observable.
func (deepseekProvider) ForceRawHTTP() bool { return true }
func (deepseekProvider) ShapeRequest(req *openai.ChatCompletionRequest, opts *ChatOptions, _ bool) {
if opts != nil && opts.ToolChoice != "" {
req.ToolChoice = nil
}
}
// --- Generic (vLLM) / NVIDIA / LiteLLM: thinking via chat_template_kwargs ---
type genericProvider struct{ baseProvider }
func (genericProvider) Name() provider.ProviderName { return provider.ProviderGeneric }
func (genericProvider) Thinking() ThinkingStrategy { return chatTemplateKwargs{} }
type nvidiaProvider struct{ baseProvider }
func (nvidiaProvider) Name() provider.ProviderName { return provider.ProviderNvidia }
func (nvidiaProvider) Thinking() ThinkingStrategy { return chatTemplateKwargs{} }
type liteLLMProvider struct{ baseProvider }
func (liteLLMProvider) Name() provider.ProviderName { return provider.ProviderLiteLLM }
func (liteLLMProvider) Thinking() ThinkingStrategy { return chatTemplateKwargs{} }
// --- Gemini OpenAI compatibility: tool thought signatures live in extra_content ---
type geminiProvider struct{ baseProvider }
func (geminiProvider) Name() provider.ProviderName { return provider.ProviderGemini }
func (geminiProvider) ForceRawHTTP() bool { return true }
func (geminiProvider) ExtractToolCallMetadata(raw json.RawMessage) types.ToolCallMetadata {
var tc struct {
ExtraContent map[string]json.RawMessage `json:"extra_content,omitempty"`
}
if err := json.Unmarshal(raw, &tc); err != nil {
return nil
}
google, ok := tc.ExtraContent["google"]
if !ok || len(google) == 0 {
return nil
}
return types.ToolCallMetadata{"google": google}
}
func (geminiProvider) InjectToolCallMetadata(toolCall map[string]any, metadata types.ToolCallMetadata) {
if len(metadata) == 0 {
return
}
google, ok := metadata["google"]
if !ok || len(google) == 0 {
return
}
var googleValue any
if err := json.Unmarshal(google, &googleValue); err != nil {
return
}
toolCall["extra_content"] = map[string]any{"google": googleValue}
}
// --- Volcengine (火山引擎 Ark): thinking via { "thinking": { "type": ... } } ---
type volcengineProvider struct{ baseProvider }
func (volcengineProvider) Name() provider.ProviderName { return provider.ProviderVolcengine }
func (volcengineProvider) Thinking() ThinkingStrategy { return thinkingTypeField{} }
// --- Azure OpenAI: api-key auth (reasoning variant also strips sampling params) ---
type azureProvider struct{ baseProvider }
func (azureProvider) Name() provider.ProviderName { return provider.ProviderAzureOpenAI }
func (azureProvider) Auth(req *http.Request, creds authCreds, _ []byte) {
req.Header.Set("api-key", creds.APIKey)
}
type azureReasoningProvider struct{ azureProvider }
func (azureReasoningProvider) Matches(model string) bool {
return provider.IsOpenAIReasoningOrGPT5Model(model)
}
func (azureReasoningProvider) ShapeRequest(req *openai.ChatCompletionRequest, _ *ChatOptions, _ bool) {
shapeOpenAIReasoning(req)
}
// --- OpenAI reasoning / GPT-5: no sampling params, must use max_completion_tokens ---
type openAIReasoningProvider struct{ baseProvider }
func (openAIReasoningProvider) Name() provider.ProviderName { return provider.ProviderOpenAI }
func (openAIReasoningProvider) Matches(model string) bool {
return provider.IsOpenAIReasoningOrGPT5Model(model)
}
func (openAIReasoningProvider) ShapeRequest(req *openai.ChatCompletionRequest, _ *ChatOptions, _ bool) {
shapeOpenAIReasoning(req)
}
// --- Moonshot: v1 models accept only temperature=1 ---
type moonshotProvider struct{ baseProvider }
func (moonshotProvider) Name() provider.ProviderName { return provider.ProviderMoonshot }
func (moonshotProvider) Matches(model string) bool {
return provider.IsMoonshotFixedTempModel(model)
}
func (moonshotProvider) ShapeRequest(req *openai.ChatCompletionRequest, _ *ChatOptions, _ bool) {
// Pin temperature to 1 and drop the other sampling params, matching the
// pre-refactor behavior where these fields were never set for this model.
req.Temperature = 1
req.TopP = 0
req.FrequencyPenalty = 0
req.PresencePenalty = 0
}
// shapeOpenAIReasoning strips sampling params (unsupported by o-series / GPT-5)
// and migrates max_tokens to max_completion_tokens. See issue #1283.
func shapeOpenAIReasoning(req *openai.ChatCompletionRequest) {
req.Temperature = 0
req.TopP = 0
req.FrequencyPenalty = 0
req.PresencePenalty = 0
if req.MaxCompletionTokens == 0 && req.MaxTokens > 0 {
req.MaxCompletionTokens = req.MaxTokens
}
req.MaxTokens = 0
}
// providerRegistry is ordered: more specific adapters (those with a real
// Matches predicate) must precede the generic catch-all for the same provider.
var providerRegistry = []providerAdapter{
weKnoraCloudProvider{},
qwenThinkingProvider{},
lkeapProvider{},
deepseekProvider{},
genericProvider{},
liteLLMProvider{},
geminiProvider{},
volcengineProvider{},
nvidiaProvider{},
azureReasoningProvider{},
azureProvider{},
openAIReasoningProvider{},
moonshotProvider{},
}
// resolveProvider returns the adapter handling the given provider+model, or
// baseProvider{} (Bearer auth, standard endpoint, no thinking) when none matches.
func resolveProvider(name provider.ProviderName, model string) providerAdapter {
for _, p := range providerRegistry {
if p.Name() == name && p.Matches(model) {
return p
}
}
return baseProvider{}
}
+229
View File
@@ -0,0 +1,229 @@
package chat
import (
"context"
"encoding/json"
"testing"
"github.com/Tencent/WeKnora/internal/models/provider"
"github.com/Tencent/WeKnora/internal/types"
"github.com/sashabaranov/go-openai"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// TestResolveProvider pins the provider+model routing table, including the
// sub-model matchers (reasoning models, Qwen thinking, LKEAP DeepSeek V3) and
// the baseProvider fallback for everything else.
func TestResolveProvider(t *testing.T) {
cases := []struct {
name string
prov provider.ProviderName
model string
want providerAdapter
}{
{"deepseek", provider.ProviderDeepSeek, "deepseek-chat", deepseekProvider{}},
{"lkeap v3", provider.ProviderLKEAP, "deepseek-v3.1", lkeapProvider{}},
{"lkeap r1 falls back", provider.ProviderLKEAP, "deepseek-r1", baseProvider{}},
{"qwen thinking", provider.ProviderAliyun, "qwen3-32b", qwenThinkingProvider{}},
{"generic", provider.ProviderGeneric, "anything", genericProvider{}},
{"litellm", provider.ProviderLiteLLM, "anything", liteLLMProvider{}},
{"gemini", provider.ProviderGemini, "gemini-3-flash-preview", geminiProvider{}},
{"nvidia", provider.ProviderNvidia, "anything", nvidiaProvider{}},
{"volcengine", provider.ProviderVolcengine, "doubao", volcengineProvider{}},
{"openai non-reasoning falls back", provider.ProviderOpenAI, "gpt-4o", baseProvider{}},
{"openai reasoning", provider.ProviderOpenAI, "gpt-5", openAIReasoningProvider{}},
{"azure non-reasoning", provider.ProviderAzureOpenAI, "gpt-4", azureProvider{}},
{"azure reasoning", provider.ProviderAzureOpenAI, "gpt-5-mini", azureReasoningProvider{}},
{"moonshot fixed temp", provider.ProviderMoonshot, "moonshot-v1-8k", moonshotProvider{}},
{"moonshot other falls back", provider.ProviderMoonshot, "kimi-latest", baseProvider{}},
{"weknora cloud", provider.ProviderWeKnoraCloud, "anything", weKnoraCloudProvider{}},
{"unknown falls back", provider.ProviderName("nope"), "x", baseProvider{}},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
assert.IsType(t, tc.want, resolveProvider(tc.prov, tc.model))
})
}
}
func newOutboundChat(t *testing.T, providerName, model string, extra map[string]string) *RemoteAPIChat {
t.Helper()
c, err := NewRemoteAPIChat(&ChatConfig{
Source: types.ModelSourceRemote,
ModelName: model,
APIKey: "k",
ModelID: model,
Provider: providerName,
ExtraConfig: extra,
})
require.NoError(t, err)
return c
}
// TestBuildOutbound_Thinking is the characterization suite for the merged
// thinking-control path: it asserts that buildOutbound produces the same wire
// formats the pre-refactor provider customizers did.
func TestBuildOutbound_Thinking(t *testing.T) {
msgs := []Message{{Role: "user", Content: "hi"}}
t.Run("deepseek honors explicit thinking controls", func(t *testing.T) {
for _, stream := range []bool{false, true} {
for _, enabled := range []bool{false, true} {
c := newOutboundChat(t, string(provider.ProviderDeepSeek), "deepseek-v4-flash", nil)
body, _, useRaw, err := c.buildOutbound(context.Background(), msgs, &ChatOptions{
Thinking: ptrBool(enabled), MaxTokens: 4096,
}, stream)
require.NoError(t, err)
require.True(t, useRaw)
var request map[string]any
require.NoError(t, json.Unmarshal([]byte(mustJSON(t, body)), &request))
want := "disabled"
if enabled {
want = "enabled"
}
assert.Equal(t, map[string]any{"type": want}, request["thinking"])
assert.EqualValues(t, 4096, request["max_tokens"])
assert.NotContains(t, request, "chat_template_kwargs")
}
}
})
t.Run("deepseek preserves provider default when thinking is unset", func(t *testing.T) {
c := newOutboundChat(t, string(provider.ProviderDeepSeek), "deepseek-v4-flash", nil)
body, _, _, err := c.buildOutbound(context.Background(), msgs, &ChatOptions{MaxTokens: 4096}, false)
require.NoError(t, err)
assert.NotContains(t, mustJSON(t, body), `"thinking"`)
})
t.Run("generic explicit thinking_type overrides legacy kwargs", func(t *testing.T) {
c := newOutboundChat(t, string(provider.ProviderGeneric), "deepseek-v4-flash",
map[string]string{ExtraConfigThinkingControl: "thinking_type"})
body, _, useRaw, err := c.buildOutbound(context.Background(), msgs, &ChatOptions{Thinking: ptrBool(false)}, true)
require.NoError(t, err)
require.True(t, useRaw)
js := mustJSON(t, body)
assert.Contains(t, js, `"thinking"`)
assert.Contains(t, js, `"disabled"`)
assert.NotContains(t, js, "chat_template_kwargs")
})
t.Run("generic legacy chat_template_kwargs", func(t *testing.T) {
c := newOutboundChat(t, string(provider.ProviderGeneric), "qwen", nil)
body, _, useRaw, err := c.buildOutbound(context.Background(), msgs, &ChatOptions{Thinking: ptrBool(false)}, true)
require.NoError(t, err)
require.True(t, useRaw)
assert.Contains(t, mustJSON(t, body), "chat_template_kwargs")
})
t.Run("none keeps the standard SDK request", func(t *testing.T) {
c := newOutboundChat(t, string(provider.ProviderGeneric), "x",
map[string]string{ExtraConfigThinkingControl: "none"})
body, _, useRaw, err := c.buildOutbound(context.Background(), msgs, &ChatOptions{Thinking: ptrBool(false)}, true)
require.NoError(t, err)
assert.False(t, useRaw)
_, ok := body.(*openai.ChatCompletionRequest)
assert.True(t, ok)
})
t.Run("qwen non-stream forces disabled", func(t *testing.T) {
c := newOutboundChat(t, string(provider.ProviderAliyun), "qwen3-32b", nil)
body, _, useRaw, err := c.buildOutbound(context.Background(), msgs, &ChatOptions{Thinking: ptrBool(true)}, false)
require.NoError(t, err)
require.True(t, useRaw)
assert.Contains(t, mustJSON(t, body), `"enable_thinking":false`)
})
t.Run("qwen stream honors requested true", func(t *testing.T) {
c := newOutboundChat(t, string(provider.ProviderAliyun), "qwen3-32b", nil)
body, _, _, err := c.buildOutbound(context.Background(), msgs, &ChatOptions{Thinking: ptrBool(true)}, true)
require.NoError(t, err)
assert.Contains(t, mustJSON(t, body), `"enable_thinking":true`)
})
t.Run("volcengine thinking enabled", func(t *testing.T) {
c := newOutboundChat(t, string(provider.ProviderVolcengine), "doubao", nil)
body, _, useRaw, err := c.buildOutbound(context.Background(), msgs, &ChatOptions{Thinking: ptrBool(true)}, true)
require.NoError(t, err)
require.True(t, useRaw)
js := mustJSON(t, body)
assert.Contains(t, js, `"thinking"`)
assert.Contains(t, js, `"enabled"`)
})
t.Run("lkeap deepseek-v3 emits thinking type", func(t *testing.T) {
c := newOutboundChat(t, string(provider.ProviderLKEAP), "deepseek-v3.1", nil)
body, _, useRaw, err := c.buildOutbound(context.Background(), msgs, &ChatOptions{Thinking: ptrBool(false)}, true)
require.NoError(t, err)
require.True(t, useRaw)
assert.Contains(t, mustJSON(t, body), `"thinking"`)
})
t.Run("lkeap r1 left untouched", func(t *testing.T) {
c := newOutboundChat(t, string(provider.ProviderLKEAP), "deepseek-r1", nil)
body, _, useRaw, err := c.buildOutbound(context.Background(), msgs, &ChatOptions{Thinking: ptrBool(false)}, true)
require.NoError(t, err)
assert.False(t, useRaw)
_, ok := body.(*openai.ChatCompletionRequest)
assert.True(t, ok)
})
}
// TestBuildOutbound_ShapeRequest covers the param-shaping providers that used
// to live inline in BuildChatCompletionRequest.
func TestBuildOutbound_ShapeRequest(t *testing.T) {
msgs := []Message{{Role: "user", Content: "hi"}}
t.Run("deepseek strips tool_choice", func(t *testing.T) {
c := newOutboundChat(t, string(provider.ProviderDeepSeek), "deepseek-chat", nil)
body, _, useRaw, err := c.buildOutbound(context.Background(), msgs, &ChatOptions{ToolChoice: "auto"}, false)
require.NoError(t, err)
assert.True(t, useRaw)
request, ok := body.(map[string]any)
require.True(t, ok)
assert.NotContains(t, request, "tool_choice")
})
t.Run("moonshot pins temperature to 1", func(t *testing.T) {
c := newOutboundChat(t, string(provider.ProviderMoonshot), "moonshot-v1-8k", nil)
body, _, _, err := c.buildOutbound(context.Background(), msgs, &ChatOptions{Temperature: 0.7, TopP: 0.9}, false)
require.NoError(t, err)
req := body.(*openai.ChatCompletionRequest)
assert.EqualValues(t, 1, req.Temperature)
assert.EqualValues(t, 0, req.TopP)
})
}
func TestBuildOutbound_GeminiProviderMetadata(t *testing.T) {
c := newOutboundChat(t, string(provider.ProviderGemini), "gemini-3-flash-preview", nil)
messages := []Message{
{Role: "user", Content: "find docs"},
{
Role: "assistant",
ToolCalls: []ToolCall{{
ID: "call_1",
Type: "function",
ProviderMetadata: types.ToolCallMetadata{"google": json.RawMessage(`{"thought_signature":"gemini-signature"}`)},
Function: FunctionCall{
Name: "wiki_search",
Arguments: `{"query":"MACS"}`,
},
}},
},
}
body, _, useRaw, err := c.buildOutbound(context.Background(), messages, &ChatOptions{}, false)
require.NoError(t, err)
require.True(t, useRaw)
js := mustJSON(t, body)
assert.Contains(t, js, `"extra_content"`)
assert.Contains(t, js, `"thought_signature":"gemini-signature"`)
}
func mustJSON(t *testing.T, v any) string {
t.Helper()
b, err := json.Marshal(v)
require.NoError(t, err)
return string(b)
}
+440
View File
@@ -0,0 +1,440 @@
package chat
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/models/provider"
"github.com/Tencent/WeKnora/internal/models/requestbudget"
"github.com/Tencent/WeKnora/internal/types"
secutils "github.com/Tencent/WeKnora/internal/utils"
"github.com/sashabaranov/go-openai"
)
// RemoteAPIChat 实现了基于 OpenAI 兼容 API 的聊天。
// 它本身只负责通用的请求/响应/流式处理;所有 provider 特定行为都委托给
// providerAdapter(见 provider.go),thinking 编码委托给 ThinkingStrategy
// (见 thinking.go)。
type RemoteAPIChat struct {
modelName string
client *openai.Client
modelID string
baseURL string
apiKey string
provider provider.ProviderName
appID string
appSecret string
// customHeaders 为用户在模型配置中指定的自定义 HTTP 请求头(类似 OpenAI Python SDK 的 extra_headers)。
customHeaders map[string]string
// adapter 承载所有 provider 特定行为(thinking / 参数特判 / endpoint / 鉴权 / 消息变换)。
adapter providerAdapter
// thinkingOverride 来自 extra_config.thinking_control,非 nil 时覆盖 adapter.Thinking()。
thinkingOverride ThinkingStrategy
}
// NewRemoteAPIChat 创建远程 API 聊天实例
func NewRemoteAPIChat(chatConfig *ChatConfig) (*RemoteAPIChat, error) {
if chatConfig.BaseURL != "" {
if err := secutils.ValidateURLForSSRF(chatConfig.BaseURL); err != nil {
return nil, fmt.Errorf("baseURL SSRF check failed: %w", err)
}
}
apiKey := chatConfig.APIKey
providerName := provider.ProviderName(chatConfig.Provider)
if providerName == "" {
providerName = provider.DetectProvider(chatConfig.BaseURL)
}
var config openai.ClientConfig
if providerName == provider.ProviderAzureOpenAI {
config = openai.DefaultAzureConfig(apiKey, chatConfig.BaseURL)
config.AzureModelMapperFunc = func(model string) string {
return model
}
if chatConfig.ExtraConfig != nil {
if v, ok := chatConfig.ExtraConfig["api_version"]; ok {
config.APIVersion = v
}
}
} else {
config = openai.DefaultConfig(apiKey)
if baseURL := chatConfig.BaseURL; baseURL != "" {
config.BaseURL = baseURL
} else if providerName == provider.ProviderDeepSeek {
config.BaseURL = provider.DeepSeekBaseURL
}
}
// The SDK must use the same SSRF-safe transport as the raw HTTP paths.
// Constructor-time URL validation alone cannot prevent DNS rebinding or a
// later redirect to an internal address.
sdkHTTPClient := rawHTTPClient
// 如果指定了 CustomHeaders,则给 SDK 使用的 HTTPClient 挂一层 RoundTripper,
// 在每个请求上自动注入这些 header(raw HTTP 路径会在发送前单独处理)。
if len(chatConfig.CustomHeaders) > 0 {
sdkHTTPClient = secutils.WrapHTTPClientWithHeaders(sdkHTTPClient, chatConfig.CustomHeaders)
}
config.HTTPClient = sdkHTTPClient
modelName := chatConfig.ModelName
if chatConfig.ExtraConfig != nil {
if override := strings.TrimSpace(chatConfig.ExtraConfig["remote_model_name"]); override != "" {
modelName = override
}
}
if providerName == provider.ProviderWeKnoraCloud {
if chatConfig.AppID == "" {
return nil, fmt.Errorf("WeKnoraCloud provider: AppID is required")
}
if chatConfig.AppSecret == "" {
return nil, fmt.Errorf("WeKnoraCloud provider: AppSecret is required")
}
}
return &RemoteAPIChat{
modelName: modelName,
client: openai.NewClientWithConfig(config),
modelID: chatConfig.ModelID,
baseURL: strings.TrimRight(config.BaseURL, "/"),
apiKey: apiKey,
provider: providerName,
appID: chatConfig.AppID,
appSecret: chatConfig.AppSecret,
customHeaders: chatConfig.CustomHeaders,
adapter: resolveProvider(providerName, modelName),
thinkingOverride: parseThinkingOverride(chatConfig.ExtraConfig),
}, nil
}
// authCreds bundles the credentials passed to the adapter's Auth method.
func (c *RemoteAPIChat) authCreds() authCreds {
return authCreds{APIKey: c.apiKey, AppID: c.appID, AppSecret: c.appSecret}
}
// shapedRequest builds the standard request and applies the adapter's message
// transform and parameter shaping (but not thinking, which may wrap the body).
func (c *RemoteAPIChat) shapedRequest(messages []Message, opts *ChatOptions, isStream bool) openai.ChatCompletionRequest {
req := c.BuildChatCompletionRequest(messages, opts, isStream)
req.Messages = c.adapter.TransformMessages(req.Messages)
c.adapter.ShapeRequest(&req, opts, isStream)
return req
}
// buildOutbound assembles the final outbound request: the body to send, the
// endpoint override (empty for the standard endpoint), and whether the raw HTTP
// path is required. This is the single place that composes adapter + thinking,
// replacing the former buildRequestCustomizer plumbing.
func (c *RemoteAPIChat) buildOutbound(
ctx context.Context, messages []Message, opts *ChatOptions, isStream bool,
) (body any, endpoint string, useRawHTTP bool, err error) {
req := c.shapedRequest(messages, opts, isStream)
thinking := c.thinkingOverride
if thinking == nil {
thinking = c.adapter.Thinking()
}
customBody, useRaw := thinking.Apply(&req, opts, isStream)
body = &req
if customBody != nil {
body = customBody
}
body, err = c.shapeProviderRequest(body, req, messages)
if err != nil {
return nil, "", false, err
}
retention := resolveCacheRetention(opts)
policy := promptCachePolicyFor(c.provider, c.baseURL)
sessionID := promptCacheSessionID(ctx, opts)
cachedBody, forceRaw, err := applyPromptCacheToJSONBody(body, policy, sessionID, retention)
if err != nil {
return nil, "", false, err
}
body = cachedBody
endpoint = c.adapter.Endpoint(c.baseURL, c.modelID, isStream)
useRawHTTP = useRaw || c.adapter.ForceRawHTTP() || endpoint != "" || forceRaw
return body, endpoint, useRawHTTP, nil
}
// logRequest 记录请求日志
func (c *RemoteAPIChat) logRequest(ctx context.Context, req any, isStream bool) {
if jsonData, err := json.MarshalIndent(req, "", " "); err == nil {
logger.Infof(ctx, "[LLM Request] model=%s, stream=%v, request:\n%s",
c.modelName, isStream, secutils.CompactImageDataURLForLog(string(jsonData)))
}
}
// Chat 进行非流式聊天
func (c *RemoteAPIChat) Chat(ctx context.Context, messages []Message, opts *ChatOptions) (*types.ChatResponse, error) {
ctx = requestbudget.WithModel(ctx, c.modelID)
if err := requestbudget.Check(ctx); err != nil {
return nil, err
}
// 仅在调用方未设置 deadline 时附加一个兜底超时,防止 hung 请求永久阻塞 worker;
// 调用方若显式设置了更短或更长的 deadline,都会被原样尊重。
timeoutCtx, cancel := withLLMTimeout(ctx, defaultChatTimeout)
defer cancel()
body, endpoint, useRawHTTP, err := c.buildOutbound(timeoutCtx, messages, opts, false)
if err != nil {
return nil, err
}
if useRawHTTP {
return c.chatWithRawHTTP(timeoutCtx, endpoint, body, opts)
}
req := *(body.(*openai.ChatCompletionRequest))
c.logRequest(timeoutCtx, req, false)
resp, err := c.client.CreateChatCompletion(timeoutCtx, req)
if err != nil {
if isMultimodalNotSupportedError(err) {
logger.Warnf(timeoutCtx, "[LLM Request] Model %s does not support multimodal, retrying without images", c.modelName)
cleaned := stripImagesFromMessages(messages)
req = c.shapedRequest(cleaned, opts, false)
resp, err = c.client.CreateChatCompletion(timeoutCtx, req)
}
if err != nil {
return nil, fmt.Errorf("create chat completion: %w", err)
}
}
result, err := c.parseCompletionResponse(&resp)
if err != nil {
return nil, err
}
logUsage(timeoutCtx, c.modelName, &result.Usage)
return result, nil
}
// chatWithRawHTTP 使用原始 HTTP 请求进行聊天(供自定义请求使用)
func (c *RemoteAPIChat) chatWithRawHTTP(ctx context.Context, endpoint string, customReq any, opts *ChatOptions) (*types.ChatResponse, error) {
ctx = requestbudget.WithModel(ctx, c.modelID)
jsonData, err := json.Marshal(customReq)
if err != nil {
return nil, fmt.Errorf("marshal request: %w", err)
}
if endpoint == "" {
endpoint = c.baseURL + "/chat/completions"
}
if err := secutils.ValidateURLForSSRF(endpoint); err != nil {
return nil, fmt.Errorf("endpoint SSRF check failed: %w", err)
}
logger.Infof(ctx, "[LLM Request] Remote HTTP, endpoint=%s, model=%s, raw HTTP request:\n%s",
endpoint, c.modelName, secutils.CompactImageDataURLForLog(string(jsonData)))
httpReq, err := http.NewRequestWithContext(ctx, "POST", endpoint, bytes.NewBuffer(jsonData))
if err != nil {
return nil, fmt.Errorf("create request: %w", err)
}
httpReq.Header.Set("Content-Type", "application/json")
c.adapter.Auth(httpReq, c.authCreds(), jsonData)
// 注入用户自定义 header(保留头会在工具内部自动跳过)
secutils.ApplyCustomHeaders(httpReq, c.customHeaders)
attachPromptCacheHeaders(httpReq, promptCachePolicyFor(c.provider, c.baseURL), promptCacheSessionID(ctx, opts))
logger.Infof(ctx, "[LLM Request] Remote HTTP, endpoint=%s, model=%s",
endpoint, c.modelName)
resp, err := rawHTTPClient.Do(httpReq)
if err != nil {
return nil, fmt.Errorf("send request: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("read response: %w", err)
}
var chatResp openai.ChatCompletionResponse
if err := json.Unmarshal(body, &chatResp); err != nil {
return nil, fmt.Errorf("decode response: %w", err)
}
result, err := c.parseCompletionResponse(&chatResp)
if err != nil {
return nil, err
}
c.applyCompletionToolCallMetadata(body, result)
applyRawPromptCacheUsage(body, &result.Usage)
logUsage(ctx, c.modelName, &result.Usage)
return result, nil
}
// ChatStream 进行流式聊天
func (c *RemoteAPIChat) ChatStream(ctx context.Context, messages []Message, opts *ChatOptions) (<-chan types.StreamResponse, error) {
ctx = requestbudget.WithModel(ctx, c.modelID)
if err := requestbudget.Check(ctx); err != nil {
return nil, err
}
// 仅在调用方未设置 deadline 时附加兜底超时;流式调用默认超时更长,
// 因为带思考/推理的模型可能数十秒甚至几分钟才产出首 token。
timeoutCtx, cancel := withLLMTimeout(ctx, defaultStreamTimeout)
body, endpoint, useRawHTTP, err := c.buildOutbound(timeoutCtx, messages, opts, true)
if err != nil {
cancel()
return nil, err
}
if useRawHTTP {
ch, err := c.chatStreamWithRawHTTP(timeoutCtx, endpoint, body, opts)
return wrapStreamCancel(ch, err, cancel)
}
req := *(body.(*openai.ChatCompletionRequest))
c.logRequest(timeoutCtx, req, true)
streamDumper := newStreamPacketDumper(c.modelName, &req)
if streamDumper != nil {
logger.Infof(timeoutCtx, "[LLM Stream Raw Dump] writing packets to %s", streamDumper.Path())
}
streamChan := make(chan types.StreamResponse)
stream, err := c.client.CreateChatCompletionStream(timeoutCtx, req)
if err != nil {
if isMultimodalNotSupportedError(err) {
logger.Warnf(timeoutCtx, "[LLM Stream] Model %s does not support multimodal, retrying without images", c.modelName)
cleaned := stripImagesFromMessages(messages)
req = c.shapedRequest(cleaned, opts, true)
stream, err = c.client.CreateChatCompletionStream(timeoutCtx, req)
}
if err != nil {
cancel()
close(streamChan)
return nil, fmt.Errorf("create chat completion stream: %w", err)
}
}
go func() {
defer cancel()
if streamDumper != nil {
defer streamDumper.Close()
}
c.processStream(timeoutCtx, stream, streamChan, streamDumper)
}()
return streamChan, nil
}
// wrapStreamCancel 在子 channel 关闭后执行 cancel,避免 timeout context 泄漏。
// 当底层调用直接返回 error 时,立即调用 cancel 并将 error 透出。
func wrapStreamCancel(in <-chan types.StreamResponse, err error, cancel context.CancelFunc) (<-chan types.StreamResponse, error) {
if err != nil {
cancel()
return nil, err
}
out := make(chan types.StreamResponse)
go func() {
defer cancel()
defer close(out)
for v := range in {
out <- v
}
}()
return out, nil
}
// chatStreamWithRawHTTP 使用原始 HTTP 请求进行流式聊天
func (c *RemoteAPIChat) chatStreamWithRawHTTP(ctx context.Context, endpoint string, customReq any, opts *ChatOptions) (<-chan types.StreamResponse, error) {
ctx = requestbudget.WithModel(ctx, c.modelID)
jsonData, err := json.Marshal(customReq)
if err != nil {
return nil, fmt.Errorf("marshal request: %w", err)
}
if endpoint == "" {
endpoint = c.baseURL + "/chat/completions"
}
if err := secutils.ValidateURLForSSRF(endpoint); err != nil {
return nil, fmt.Errorf("endpoint SSRF check failed: %w", err)
}
if prettyJSON, pErr := json.MarshalIndent(customReq, "", " "); pErr == nil {
logger.Infof(ctx, "[LLM Stream Request] endpoint=%s, model=%s, stream=true, request:\n%s",
endpoint, c.modelName, secutils.CompactImageDataURLForLog(string(prettyJSON)))
} else {
logger.Infof(ctx, "[LLM Stream] endpoint=%s, model=%s", endpoint, c.modelName)
}
httpReq, err := http.NewRequestWithContext(ctx, "POST", endpoint, bytes.NewBuffer(jsonData))
if err != nil {
return nil, fmt.Errorf("create request: %w", err)
}
httpReq.Header.Set("Content-Type", "application/json")
c.adapter.Auth(httpReq, c.authCreds(), jsonData)
httpReq.Header.Set("Accept", "text/event-stream")
// 注入用户自定义 header(保留头会在工具内部自动跳过)
secutils.ApplyCustomHeaders(httpReq, c.customHeaders)
attachPromptCacheHeaders(httpReq, promptCachePolicyFor(c.provider, c.baseURL), promptCacheSessionID(ctx, opts))
resp, err := rawHTTPClient.Do(httpReq)
if err != nil {
return nil, fmt.Errorf("send request: %w", err)
}
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
resp.Body.Close()
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
streamChan := make(chan types.StreamResponse)
streamDumper := newStreamPacketDumper(c.modelName, customReq)
if streamDumper != nil {
logger.Infof(ctx, "[LLM Stream Raw Dump] writing packets to %s", streamDumper.Path())
}
go func() {
if streamDumper != nil {
defer streamDumper.Close()
}
c.processRawHTTPStream(ctx, resp, streamChan, streamDumper)
}()
return streamChan, nil
}
// GetModelName 获取模型名称
func (c *RemoteAPIChat) GetModelName() string {
return c.modelName
}
// GetModelID 获取模型ID
func (c *RemoteAPIChat) GetModelID() string {
return c.modelID
}
// GetProvider 获取 provider 名称
func (c *RemoteAPIChat) GetProvider() provider.ProviderName {
return c.provider
}
// GetBaseURL 获取 baseURL
func (c *RemoteAPIChat) GetBaseURL() string {
return c.baseURL
}
// GetAPIKey 获取 apiKey
func (c *RemoteAPIChat) GetAPIKey() string {
return c.apiKey
}
+575
View File
@@ -0,0 +1,575 @@
package chat
import (
"context"
"encoding/json"
"os"
"testing"
"time"
"github.com/Tencent/WeKnora/internal/types"
"github.com/sashabaranov/go-openai"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func newTestRemoteChat(t *testing.T) *RemoteAPIChat {
t.Helper()
chat, err := NewRemoteAPIChat(&ChatConfig{
Source: types.ModelSourceRemote,
BaseURL: "",
ModelName: "test-model",
APIKey: "test-key",
ModelID: "test-model",
})
require.NoError(t, err)
return chat
}
func TestBuildChatCompletionRequest_ParallelToolCalls(t *testing.T) {
chat := newTestRemoteChat(t)
messages := []Message{{Role: "user", Content: "hello"}}
t.Run("nil ParallelToolCalls leaves default", func(t *testing.T) {
opts := &ChatOptions{Temperature: 0.7}
req := chat.BuildChatCompletionRequest(messages, opts, false)
assert.Nil(t, req.ParallelToolCalls, "should be nil when not set")
})
t.Run("ParallelToolCalls true is propagated", func(t *testing.T) {
ptc := true
opts := &ChatOptions{
Temperature: 0.7,
ParallelToolCalls: &ptc,
Tools: []Tool{{
Type: "function",
Function: FunctionDef{
Name: "mcp_weather_getforecast",
Description: "Get weather",
Parameters: json.RawMessage(`{"type":"object"}`),
},
}},
}
req := chat.BuildChatCompletionRequest(messages, opts, true)
assert.NotNil(t, req.ParallelToolCalls)
val, ok := req.ParallelToolCalls.(bool)
if ok {
assert.Equal(t, true, val)
} else {
assert.Equal(t, true, req.ParallelToolCalls)
}
assert.Len(t, req.Tools, 1)
assert.Equal(t, "mcp_weather_getforecast", req.Tools[0].Function.Name)
})
t.Run("ParallelToolCalls false is propagated", func(t *testing.T) {
ptc := false
opts := &ChatOptions{
Temperature: 0.7,
ParallelToolCalls: &ptc,
}
req := chat.BuildChatCompletionRequest(messages, opts, false)
assert.NotNil(t, req.ParallelToolCalls)
val, ok := req.ParallelToolCalls.(bool)
if ok {
assert.Equal(t, false, val)
} else {
assert.Equal(t, false, req.ParallelToolCalls)
}
})
}
func TestBuildChatCompletionRequest_MCPToolsFormat(t *testing.T) {
chat := newTestRemoteChat(t)
messages := []Message{{Role: "user", Content: "查询乙醇的理化性质"}}
mcpTools := []Tool{
{
Type: "function",
Function: FunctionDef{
Name: "mcp_hazardous_chemicals_gethazardouschemicals",
Description: "[MCP Service: hazardous_chemicals (external)] Get hazardous chemicals list",
Parameters: json.RawMessage(`{"type":"object","properties":{}}`),
},
},
{
Type: "function",
Function: FunctionDef{
Name: "mcp_hazardous_chemicals_gethazardouschemicalbybizid",
Description: "[MCP Service: hazardous_chemicals (external)] Get hazardous chemical by biz ID",
Parameters: json.RawMessage(`{"type":"object","properties":{"bizId":{"type":"string"}},"required":["bizId"]}`),
},
},
}
ptc := true
opts := &ChatOptions{
Temperature: 0.7,
Tools: mcpTools,
ParallelToolCalls: &ptc,
}
req := chat.BuildChatCompletionRequest(messages, opts, true)
assert.Len(t, req.Tools, 2)
assert.Equal(t, "mcp_hazardous_chemicals_gethazardouschemicals", req.Tools[0].Function.Name)
assert.Equal(t, "mcp_hazardous_chemicals_gethazardouschemicalbybizid", req.Tools[1].Function.Name)
assert.Equal(t, true, req.ParallelToolCalls)
assert.True(t, req.Stream)
for _, tool := range req.Tools {
name := tool.Function.Name
assert.NotContains(t, name, "ed606721", "tool name must use service name, not UUID")
assert.Regexp(t, `^[a-zA-Z0-9_-]+$`, name, "tool name must match OpenAI pattern")
assert.LessOrEqual(t, len(name), 64, "tool name must be <= 64 chars")
}
}
// TestBuildChatCompletionRequest_GPT5MaxCompletionTokens 验证 GPT-5 / o-series
// 模型的 MaxTokens 自动迁移到 MaxCompletionTokens,且采样参数被剔除。
// 见 issue #1283:Azure OpenAI 的 gpt-5 系列模型不再支持 max_tokens 字段。
func TestBuildChatCompletionRequest_GPT5MaxCompletionTokens(t *testing.T) {
build := func(t *testing.T, providerName, modelName string) *RemoteAPIChat {
t.Helper()
c, err := NewRemoteAPIChat(&ChatConfig{
Source: types.ModelSourceRemote,
BaseURL: "https://example.openai.azure.com",
ModelName: modelName,
APIKey: "test-key",
ModelID: modelName,
Provider: providerName,
ExtraConfig: map[string]string{
"api_version": "2025-04-01-preview",
},
})
require.NoError(t, err)
return c
}
messages := []Message{{Role: "user", Content: "test"}}
cases := []struct {
name string
provider string
model string
shouldRewriteMaxT bool
}{
{"AzureOpenAI gpt-5.2", "azure_openai", "gpt-5.2", true},
{"AzureOpenAI gpt-5-mini", "azure_openai", "gpt-5-mini", true},
{"OpenAI gpt-5", "openai", "gpt-5", true},
{"OpenAI o1-mini", "openai", "o1-mini", true},
{"OpenAI o3", "openai", "o3", true},
{"OpenAI o4-mini", "openai", "o4-mini", true},
{"OpenAI gpt-4o (unchanged)", "openai", "gpt-4o", false},
{"AzureOpenAI gpt-4 (unchanged)", "azure_openai", "gpt-4", false},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
c := build(t, tc.provider, tc.model)
opts := &ChatOptions{
Temperature: 0.7,
TopP: 0.9,
MaxTokens: 128,
FrequencyPenalty: 0.1,
PresencePenalty: 0.2,
}
req := c.shapedRequest(messages, opts, false)
if tc.shouldRewriteMaxT {
assert.Equal(t, 0, req.MaxTokens, "MaxTokens must NOT be sent for GPT-5/o-series")
assert.Equal(t, 128, req.MaxCompletionTokens, "MaxCompletionTokens should be populated from MaxTokens")
assert.EqualValues(t, 0, req.Temperature, "temperature must be omitted")
assert.EqualValues(t, 0, req.TopP, "top_p must be omitted")
assert.EqualValues(t, 0, req.FrequencyPenalty, "frequency_penalty must be omitted")
assert.EqualValues(t, 0, req.PresencePenalty, "presence_penalty must be omitted")
} else {
assert.Equal(t, 128, req.MaxTokens)
assert.Equal(t, 0, req.MaxCompletionTokens)
assert.InDelta(t, 0.7, req.Temperature, 1e-6)
}
})
}
t.Run("MaxCompletionTokens takes precedence over MaxTokens", func(t *testing.T) {
c := build(t, "openai", "gpt-5.2")
opts := &ChatOptions{
MaxTokens: 128,
MaxCompletionTokens: 2048,
}
req := c.shapedRequest(messages, opts, false)
assert.Equal(t, 0, req.MaxTokens)
assert.Equal(t, 2048, req.MaxCompletionTokens)
})
}
func TestBuildChatCompletionRequest_ToolChoice(t *testing.T) {
chat := newTestRemoteChat(t)
messages := []Message{{Role: "user", Content: "test"}}
t.Run("auto tool choice", func(t *testing.T) {
opts := &ChatOptions{ToolChoice: "auto"}
req := chat.BuildChatCompletionRequest(messages, opts, false)
assert.Equal(t, "auto", req.ToolChoice)
})
t.Run("specific tool choice", func(t *testing.T) {
opts := &ChatOptions{ToolChoice: "mcp_svc_tool"}
req := chat.BuildChatCompletionRequest(messages, opts, false)
assert.NotNil(t, req.ToolChoice)
})
}
// TestConvertMessages_ReasoningContentRoundTrip verifies that assistant
// reasoning_content is propagated through ConvertMessages so that providers
// like MiMo / DeepSeek thinking-mode can read it back from prior turns.
// See issue #1302: MiMo rejects multi-turn requests with HTTP 400
// "The reasoning_content in the thinking mode must be passed back to the API."
// when this field is dropped.
func TestConvertMessages_ReasoningContentRoundTrip(t *testing.T) {
c := newTestRemoteChat(t)
t.Run("assistant reasoning_content propagated", func(t *testing.T) {
messages := []Message{
{Role: "user", Content: "hi"},
{
Role: "assistant",
Content: "the answer",
ReasoningContent: "let me think about this carefully",
},
{Role: "user", Content: "follow-up"},
}
out := c.ConvertMessages(messages)
require.Len(t, out, 3)
assert.Equal(t, "let me think about this carefully", out[1].ReasoningContent,
"assistant reasoning_content must be retained for multi-turn replay")
assert.Empty(t, out[0].ReasoningContent, "user message must not carry reasoning_content")
assert.Empty(t, out[2].ReasoningContent, "user message must not carry reasoning_content")
})
t.Run("non-assistant role drops reasoning_content even if set", func(t *testing.T) {
messages := []Message{
{Role: "user", Content: "hi", ReasoningContent: "should be dropped"},
}
out := c.ConvertMessages(messages)
require.Len(t, out, 1)
assert.Empty(t, out[0].ReasoningContent, "non-assistant roles must never carry reasoning_content upstream")
})
t.Run("empty assistant reasoning_content stays empty", func(t *testing.T) {
messages := []Message{
{Role: "assistant", Content: "no thinking"},
}
out := c.ConvertMessages(messages)
require.Len(t, out, 1)
assert.Empty(t, out[0].ReasoningContent)
})
}
func TestApplyCompletionToolCallMetadata(t *testing.T) {
c := newTestRemoteChat(t)
c.adapter = geminiProvider{}
resp := &types.ChatResponse{
ToolCalls: []types.LLMToolCall{{
ID: "call_1",
Type: "function",
Function: types.FunctionCall{
Name: "wiki_search",
Arguments: `{"query":"MACS"}`,
},
}},
}
body := []byte(`{
"choices":[{
"message":{
"tool_calls":[{
"id":"call_1",
"type":"function",
"function":{"name":"wiki_search","arguments":"{\"query\":\"MACS\"}"},
"extra_content":{"google":{"thought_signature":"sig-from-gemini"}}
}]
}
}]
}`)
c.applyCompletionToolCallMetadata(body, resp)
require.Len(t, resp.ToolCalls, 1)
assert.JSONEq(t, `{"thought_signature":"sig-from-gemini"}`,
string(resp.ToolCalls[0].ProviderMetadata["google"]))
}
func TestApplyStreamToolCallMetadata(t *testing.T) {
c := newTestRemoteChat(t)
c.adapter = geminiProvider{}
state := newStreamState()
body := []byte(`{
"choices":[{
"delta":{
"tool_calls":[{
"index":0,
"id":"call_1",
"type":"function",
"function":{"name":"wiki_search","arguments":"{\"query\":\"MACS\"}"},
"extra_content":{"google":{"thought_signature":"stream-sig-from-gemini"}}
}]
}
}]
}`)
c.applyStreamToolCallMetadata(body, state)
toolCalls := state.buildOrderedToolCalls()
require.Len(t, toolCalls, 1)
assert.JSONEq(t, `{"thought_signature":"stream-sig-from-gemini"}`,
string(toolCalls[0].ProviderMetadata["google"]))
}
// TestRemoteAPIChat 综合测试 Remote API Chat 的所有功能
func TestRemoteAPIChat(t *testing.T) {
// 获取环境变量
deepseekAPIKey := os.Getenv("DEEPSEEK_API_KEY")
aliyunAPIKey := os.Getenv("ALIYUN_API_KEY")
// 定义测试配置
testConfigs := []struct {
name string
apiKey string
config *ChatConfig
skipMsg string
}{
{
name: "DeepSeek API",
apiKey: deepseekAPIKey,
config: &ChatConfig{
Source: types.ModelSourceRemote,
BaseURL: "https://api.deepseek.com/v1",
ModelName: "deepseek-chat",
APIKey: deepseekAPIKey,
ModelID: "deepseek-chat",
},
skipMsg: "DEEPSEEK_API_KEY environment variable not set",
},
{
name: "Aliyun DeepSeek",
apiKey: aliyunAPIKey,
config: &ChatConfig{
Source: types.ModelSourceRemote,
BaseURL: "https://dashscope.aliyuncs.com/compatible-mode/v1",
ModelName: "deepseek-v3.1",
APIKey: aliyunAPIKey,
ModelID: "deepseek-v3.1",
},
skipMsg: "ALIYUN_API_KEY environment variable not set",
},
{
name: "Aliyun Qwen3-32b",
apiKey: aliyunAPIKey,
config: &ChatConfig{
Source: types.ModelSourceRemote,
BaseURL: "https://dashscope.aliyuncs.com/compatible-mode/v1",
ModelName: "qwen3-32b",
APIKey: aliyunAPIKey,
ModelID: "qwen3-32b",
},
skipMsg: "ALIYUN_API_KEY environment variable not set",
},
{
name: "Aliyun Qwen-max",
apiKey: aliyunAPIKey,
config: &ChatConfig{
Source: types.ModelSourceRemote,
BaseURL: "https://dashscope.aliyuncs.com/compatible-mode/v1",
ModelName: "qwen-max",
APIKey: aliyunAPIKey,
ModelID: "qwen-max",
},
skipMsg: "ALIYUN_API_KEY environment variable not set",
},
}
// 测试消息
testMessages := []Message{
{
Role: "user",
Content: "test",
},
}
// 测试选项
testOptions := &ChatOptions{
Temperature: 0.7,
MaxTokens: 100,
}
// 创建上下文
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
// 遍历所有配置进行测试
for _, tc := range testConfigs {
t.Run(tc.name, func(t *testing.T) {
// 检查 API Key
if tc.apiKey == "" {
t.Skip(tc.skipMsg)
}
// 创建聊天实例
chat, err := NewRemoteAPIChat(tc.config)
require.NoError(t, err)
assert.Equal(t, tc.config.ModelName, chat.GetModelName())
assert.Equal(t, tc.config.ModelID, chat.GetModelID())
// 测试基本聊天功能
t.Run("Basic Chat", func(t *testing.T) {
response, err := chat.Chat(ctx, testMessages, testOptions)
require.NoError(t, err)
require.NotNil(t, response, "response should not be nil")
assert.NotEmpty(t, response.Content)
assert.Greater(t, response.Usage.TotalTokens, 0)
assert.Greater(t, response.Usage.PromptTokens, 0)
assert.Greater(t, response.Usage.CompletionTokens, 0)
t.Logf("%s Response: %s", tc.name, response.Content)
t.Logf("Usage: Prompt=%d, Completion=%d, Total=%d",
response.Usage.PromptTokens,
response.Usage.CompletionTokens,
response.Usage.TotalTokens)
})
})
}
}
// TestCachedTokensHelper covers the nil-safety contract of the cachedTokens
// helper. Some providers omit PromptTokensDetails entirely; the helper must
// return zero rather than panic.
func TestCachedTokensHelper(t *testing.T) {
assert.Equal(t, 0, cachedTokens(nil), "nil details must return zero")
assert.Equal(t, 0, cachedTokens(&openai.PromptTokensDetails{}),
"empty details must return zero")
assert.Equal(t, 1234, cachedTokens(&openai.PromptTokensDetails{CachedTokens: 1234}),
"populated cached_tokens must round-trip")
}
func TestParseCompletionResponse_ReasoningRemainsSeparate(t *testing.T) {
c := newTestRemoteChat(t)
for _, tc := range []struct {
name, content, reasoning, want string
}{
{"answer with reasoning", "answer", "private reasoning", "answer"},
{"reasoning only", "", "private reasoning", ""},
{"ordinary answer", "answer", "", "answer"},
{"incomplete inline thinking", "<think>private incomplete reasoning", "", ""},
} {
t.Run(tc.name, func(t *testing.T) {
got, err := c.parseCompletionResponse(&openai.ChatCompletionResponse{
Choices: []openai.ChatCompletionChoice{{
Message: openai.ChatCompletionMessage{Role: "assistant", Content: tc.content, ReasoningContent: tc.reasoning},
FinishReason: openai.FinishReasonLength,
}},
})
require.NoError(t, err)
assert.Equal(t, tc.want, got.Content)
assert.Equal(t, tc.reasoning, got.ReasoningContent)
assert.Equal(t, "length", got.FinishReason)
})
}
}
// TestParseCompletionResponse_CachedTokens verifies that
// prompt_tokens_details.cached_tokens from an OpenAI-compatible response is
// propagated into TokenUsage.CachedTokens. This is the field Qwen explicit
// caching populates on a cache hit.
func TestParseCompletionResponse_CachedTokens(t *testing.T) {
c := newTestRemoteChat(t)
t.Run("cached_tokens populated from prompt_tokens_details", func(t *testing.T) {
resp := &openai.ChatCompletionResponse{
Choices: []openai.ChatCompletionChoice{
{
Message: openai.ChatCompletionMessage{Role: "assistant", Content: "hi"},
FinishReason: openai.FinishReasonStop,
},
},
Usage: openai.Usage{
PromptTokens: 6929,
CompletionTokens: 42,
TotalTokens: 6971,
PromptTokensDetails: &openai.PromptTokensDetails{
CachedTokens: 6900,
},
},
}
got, err := c.parseCompletionResponse(resp)
require.NoError(t, err)
require.NotNil(t, got)
assert.Equal(t, 6929, got.Usage.PromptTokens)
assert.Equal(t, 42, got.Usage.CompletionTokens)
assert.Equal(t, 6971, got.Usage.TotalTokens)
assert.Equal(t, 6900, got.Usage.CachedTokens,
"cached_tokens must mirror prompt_tokens_details.cached_tokens")
assert.Equal(t, 6900, got.Usage.CacheReadTokens)
assert.Equal(t, 29, got.Usage.CacheMissTokens)
assert.True(t, got.Usage.CacheReported)
assert.Equal(t, types.PromptCacheStatusHit, got.Usage.CacheStatus)
})
t.Run("missing prompt_tokens_details yields zero cached_tokens", func(t *testing.T) {
resp := &openai.ChatCompletionResponse{
Choices: []openai.ChatCompletionChoice{
{
Message: openai.ChatCompletionMessage{Role: "assistant", Content: "hi"},
FinishReason: openai.FinishReasonStop,
},
},
Usage: openai.Usage{
PromptTokens: 100,
CompletionTokens: 10,
TotalTokens: 110,
// PromptTokensDetails intentionally nil — providers like Ollama
// and older OpenAI-compat backends omit this block entirely.
},
}
got, err := c.parseCompletionResponse(resp)
require.NoError(t, err)
require.NotNil(t, got)
assert.Equal(t, 0, got.Usage.CachedTokens,
"missing details must surface as zero, not panic")
assert.False(t, got.Usage.CacheReported)
assert.Equal(t, types.PromptCacheStatusUnsupported, got.Usage.CacheStatus)
})
}
func TestApplyRawPromptCacheUsage_DeepSeekNativeFields(t *testing.T) {
usage := types.TokenUsage{PromptTokens: 4096, CompletionTokens: 10, TotalTokens: 4106}
applyRawPromptCacheUsage([]byte(`{"usage":{"prompt_tokens":4096,"prompt_cache_hit_tokens":3072,"prompt_cache_miss_tokens":1024}}`), &usage)
assert.Equal(t, 3072, usage.CacheReadTokens)
assert.Equal(t, 1024, usage.CacheMissTokens)
assert.True(t, usage.CacheReported)
assert.Equal(t, types.PromptCacheStatusHit, usage.CacheStatus)
}
// TestTokenUsage_CachedTokensJSONOmitempty ensures the new CachedTokens field
// stays out of serialized payloads when it is zero. This keeps logs and API
// responses unchanged for providers that never report cache hits.
func TestTokenUsage_CachedTokensJSONOmitempty(t *testing.T) {
t.Run("zero is omitted", func(t *testing.T) {
u := types.TokenUsage{PromptTokens: 10, CompletionTokens: 5, TotalTokens: 15}
b, err := json.Marshal(u)
require.NoError(t, err)
assert.NotContains(t, string(b), "cached_tokens")
})
t.Run("non-zero is emitted", func(t *testing.T) {
u := types.TokenUsage{PromptTokens: 10, CompletionTokens: 5, TotalTokens: 15, CachedTokens: 7}
b, err := json.Marshal(u)
require.NoError(t, err)
assert.Contains(t, string(b), `"cached_tokens":7`)
})
}
+44
View File
@@ -0,0 +1,44 @@
package chat
import (
"context"
"github.com/Tencent/WeKnora/internal/models/requestbudget"
"github.com/Tencent/WeKnora/internal/types"
)
// Keep the supported set explicit: adding a factory adapter cannot silently
// bypass a controlled task's actual-send budget.
type requestBudgetChat struct {
inner Chat
supported bool
}
func wrapChatRequestBudget(inner Chat, err error) (Chat, error) {
if err != nil || inner == nil {
return inner, err
}
supported := false
switch inner.(type) {
case *RemoteAPIChat, *AnthropicChat:
supported = true
}
return &requestBudgetChat{inner: inner, supported: supported}, nil
}
func (w *requestBudgetChat) GetModelName() string { return w.inner.GetModelName() }
func (w *requestBudgetChat) GetModelID() string { return w.inner.GetModelID() }
func (w *requestBudgetChat) Chat(ctx context.Context, messages []Message, opts *ChatOptions) (*types.ChatResponse, error) {
if err := requestbudget.RequireSupported(ctx, w.supported); err != nil {
return nil, err
}
return w.inner.Chat(ctx, messages, opts)
}
func (w *requestBudgetChat) ChatStream(ctx context.Context, messages []Message, opts *ChatOptions) (<-chan types.StreamResponse, error) {
if err := requestbudget.RequireSupported(ctx, w.supported); err != nil {
return nil, err
}
return w.inner.ChatStream(ctx, messages, opts)
}
+253
View File
@@ -0,0 +1,253 @@
package chat
import (
"context"
"errors"
"fmt"
"net/http"
"net/http/httptest"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/Tencent/WeKnora/internal/models/requestbudget"
"github.com/Tencent/WeKnora/internal/types"
secutils "github.com/Tencent/WeKnora/internal/utils"
)
type chatBudgetGuard struct {
mu sync.Mutex
limit int
models, phases []string
checkErr error
}
func (g *chatBudgetGuard) Key() string { return "local-test" }
func (g *chatBudgetGuard) Check(ctx context.Context) error {
if g.checkErr != nil {
return g.checkErr
}
return ctx.Err()
}
func (g *chatBudgetGuard) Reserve(ctx context.Context, model, phase string) error {
g.mu.Lock()
defer g.mu.Unlock()
if err := ctx.Err(); err != nil {
return err
}
if len(g.models) >= g.limit {
return requestbudget.ErrStopped
}
g.models, g.phases = append(g.models, model), append(g.phases, phase)
return nil
}
func chatBudgetContext(g *chatBudgetGuard) context.Context {
return requestbudget.WithModel(requestbudget.WithPhase(requestbudget.WithGuard(context.Background(), g), "wiki_generate"), "caller-forged-model")
}
func TestRequestBudgetChatSDKRawAndStreaming(t *testing.T) {
t.Setenv("SSRF_WHITELIST", "127.0.0.1")
secutils.ResetSSRFWhitelistForTest()
t.Cleanup(secutils.ResetSSRFWhitelistForTest)
for _, raw := range []bool{false, true} {
for _, stream := range []bool{false, true} {
for _, limit := range []int{0, 1} {
t.Run(fmt.Sprintf("raw=%t/stream=%t/limit=%d", raw, stream, limit), func(t *testing.T) {
var sends atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
sends.Add(1)
if stream {
w.Header().Set("Content-Type", "text/event-stream")
fmt.Fprint(w, "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"ok\"}}]}\n\ndata: [DONE]\n\n")
} else {
w.Header().Set("Content-Type", "application/json")
fmt.Fprint(w, `{"choices":[{"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}]}`)
}
}))
defer server.Close()
c, err := NewRemoteAPIChat(&ChatConfig{BaseURL: server.URL, ModelID: "concrete-chat", ModelName: "gpt-4o", APIKey: "local-only", Provider: "openai"})
if err != nil {
t.Fatal(err)
}
opts := &ChatOptions{CacheRetention: CacheRetentionNone}
if raw {
opts.CacheRetention, opts.PromptCacheKey = CacheRetentionShort, "force-raw"
}
g := &chatBudgetGuard{limit: limit}
ctx, cancel := context.WithTimeout(chatBudgetContext(g), 3*time.Second)
defer cancel()
messages := []Message{{Role: "user", Content: "local test"}}
_, _, actualRaw, err := c.buildOutbound(ctx, messages, opts, stream)
if err != nil || actualRaw != raw {
t.Fatalf("test route mismatch: %t %v", actualRaw, err)
}
if stream {
var responses <-chan types.StreamResponse
responses, err = c.ChatStream(ctx, messages, opts)
if err == nil {
for range responses {
}
}
} else {
_, err = c.Chat(ctx, messages, opts)
}
if limit == 0 && !errors.Is(err, requestbudget.ErrStopped) {
t.Fatalf("expected stop: %v", err)
}
if limit == 1 && err != nil {
t.Fatal(err)
}
if int(sends.Load()) != limit || len(g.models) != limit {
t.Fatalf("unaccounted sends: %d models=%v", sends.Load(), g.models)
}
if limit == 1 && (g.models[0] != "concrete-chat" || g.phases[0] != "wiki_generate") {
t.Fatalf("identity/phase lost: %v %v", g.models, g.phases)
}
})
}
}
}
}
func TestRequestBudgetMultimodalFallbackCountsEachSend(t *testing.T) {
t.Setenv("SSRF_WHITELIST", "127.0.0.1")
secutils.ResetSSRFWhitelistForTest()
t.Cleanup(secutils.ResetSSRFWhitelistForTest)
for _, stream := range []bool{false, true} {
for _, limit := range []int{1, 2} {
t.Run(fmt.Sprintf("stream=%t/limit=%d", stream, limit), func(t *testing.T) {
var sends atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if sends.Add(1) == 1 {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusBadRequest)
fmt.Fprint(w, `{"error":{"message":"image input not supported","type":"invalid_request_error"}}`)
return
}
if stream {
w.Header().Set("Content-Type", "text/event-stream")
fmt.Fprint(w, "data: [DONE]\n\n")
} else {
w.Header().Set("Content-Type", "application/json")
fmt.Fprint(w, `{"choices":[{"message":{"role":"assistant","content":"ok"}}]}`)
}
}))
defer server.Close()
c, err := NewRemoteAPIChat(&ChatConfig{BaseURL: server.URL, ModelID: "concrete-chat", ModelName: "gpt-4o", APIKey: "local-only", Provider: "openai"})
if err != nil {
t.Fatal(err)
}
g := &chatBudgetGuard{limit: limit}
ctx, cancel := context.WithTimeout(chatBudgetContext(g), 3*time.Second)
defer cancel()
messages := []Message{{Role: "user", Content: "local test", Images: []string{"data:image/png;base64,AA=="}}}
opts := &ChatOptions{CacheRetention: CacheRetentionNone}
if stream {
var responses <-chan types.StreamResponse
responses, err = c.ChatStream(ctx, messages, opts)
if err == nil {
for range responses {
}
}
} else {
_, err = c.Chat(ctx, messages, opts)
}
if limit == 1 && !errors.Is(err, requestbudget.ErrStopped) {
t.Fatalf("fallback escaped budget: %v", err)
}
if limit == 2 && err != nil {
t.Fatal(err)
}
if int(sends.Load()) != limit || len(g.models) != limit {
t.Fatalf("fallback sends=%d reservations=%d", sends.Load(), len(g.models))
}
for i := range g.models {
if g.models[i] != "concrete-chat" || g.phases[i] != "wiki_generate" {
t.Fatal("fallback context lost")
}
}
})
}
}
}
func TestRequestBudgetChatFactoryFailsClosed(t *testing.T) {
g := &chatBudgetGuard{limit: 1}
ctx := chatBudgetContext(g)
c, err := NewChat(&ChatConfig{Source: types.ModelSourceLocal, ModelID: "local-model"}, nil)
if err != nil {
t.Fatal(err)
}
if _, err := c.Chat(ctx, nil, nil); !errors.Is(err, requestbudget.ErrStopped) {
t.Fatalf("local adapter invoked: %v", err)
}
if _, err := c.ChatStream(ctx, nil, nil); !errors.Is(err, requestbudget.ErrStopped) {
t.Fatalf("local stream invoked: %v", err)
}
if len(g.models) != 0 {
t.Fatal("unsupported adapter reserved a send")
}
unknown, err := wrapChatRequestBudget(&fakeChat{id: "unknown"}, nil)
if err != nil {
t.Fatal(err)
}
if _, err := unknown.Chat(ctx, nil, nil); !errors.Is(err, requestbudget.ErrStopped) {
t.Fatalf("unknown adapter allowed: %v", err)
}
}
func TestRequestBudgetAnthropicFactory(t *testing.T) {
t.Setenv("SSRF_WHITELIST", "127.0.0.1")
secutils.ResetSSRFWhitelistForTest()
t.Cleanup(secutils.ResetSSRFWhitelistForTest)
for _, stream := range []bool{false, true} {
for _, limit := range []int{0, 1} {
t.Run(fmt.Sprintf("stream=%t/limit=%d", stream, limit), func(t *testing.T) {
var sends atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
sends.Add(1)
if stream {
w.Header().Set("Content-Type", "text/event-stream")
fmt.Fprint(w, "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n")
} else {
w.Header().Set("Content-Type", "application/json")
fmt.Fprint(w, `{"type":"message","role":"assistant","content":[{"type":"text","text":"ok"}],"stop_reason":"end_turn"}`)
}
}))
defer server.Close()
c, err := NewChat(&ChatConfig{Source: types.ModelSourceRemote, BaseURL: server.URL, ModelID: "concrete-chat", ModelName: "claude-test", APIKey: "local-only", Provider: "anthropic"}, nil)
if err != nil {
t.Fatal(err)
}
g := &chatBudgetGuard{limit: limit}
ctx, cancel := context.WithTimeout(chatBudgetContext(g), 3*time.Second)
defer cancel()
messages := []Message{{Role: "user", Content: "local test"}}
if stream {
var responses <-chan types.StreamResponse
responses, err = c.ChatStream(ctx, messages, nil)
if err == nil {
for range responses {
}
}
} else {
_, err = c.Chat(ctx, messages, nil)
}
if limit == 0 && !errors.Is(err, requestbudget.ErrStopped) {
t.Fatalf("guard not applied: %v", err)
}
if limit == 1 && err != nil {
t.Fatal(err)
}
if int(sends.Load()) != limit || len(g.models) != limit {
t.Fatalf("sends=%d models=%v", sends.Load(), g.models)
}
if limit == 1 && (g.models[0] != "concrete-chat" || g.phases[0] != "wiki_generate") {
t.Fatalf("context lost: %v %v", g.models, g.phases)
}
})
}
}
}
@@ -0,0 +1,293 @@
package chat
import (
"encoding/json"
"strings"
"time"
)
// Live write/edit progress while the model is still emitting tool-call JSON:
// path as soon as it appears, then a running +N / -M as content lands. The
// SSE payload is only those stats plus a short preview — never the whole file.
const (
sandboxFilePreviewMaxLines = 10
writeSandboxFileName = "write_sandbox_file"
editSandboxFileName = "edit_sandbox_file"
)
// sandboxProgressMinInterval caps how often we emit when only byte counts
// change (a long line with no newline). Line-count changes always emit.
var sandboxProgressMinInterval = 120 * time.Millisecond
func isSandboxMutationTool(name string) bool {
return name == writeSandboxFileName || name == editSandboxFileName
}
type sandboxFileProgress struct {
toolName string
pathEx *jsonFieldExtractor
contentEx *jsonFieldExtractor
path string
content string
buf string
added int
removed int
bytes int
preview string
lastPath string
lastAdded int
lastRemoved int
lastBytes int
lastPreview string
lastEmit time.Time
emitted bool
lastParseAt time.Time
lastParseLen int
}
func newSandboxFileProgress(toolName string) *sandboxFileProgress {
return &sandboxFileProgress{toolName: toolName}
}
// Feed consumes one arguments delta. ok is true when the UI should be sent
// an updated stats payload.
func (p *sandboxFileProgress) Feed(delta string) (map[string]any, bool) {
if p == nil || delta == "" {
return nil, false
}
switch p.toolName {
case writeSandboxFileName:
p.feedWrite(delta)
case editSandboxFileName:
p.feedEdit(delta)
default:
return nil, false
}
if !p.shouldEmit() {
return nil, false
}
p.markEmitted()
return p.payload(), true
}
func (p *sandboxFileProgress) feedWrite(delta string) {
if p.pathEx == nil {
p.pathEx = newJSONFieldExtractor("path")
}
if p.contentEx == nil {
p.contentEx = newJSONFieldExtractor("content")
}
if chunk := p.pathEx.Feed(delta); chunk != "" {
p.path += chunk
}
if chunk := p.contentEx.Feed(delta); chunk != "" {
p.content += chunk
}
p.added = countContentLines(p.content)
p.removed = 0
p.bytes = len(p.content)
p.preview = sandboxContentPreview(p.content)
}
func (p *sandboxFileProgress) feedEdit(delta string) {
p.buf += delta
if p.pathEx == nil {
p.pathEx = newJSONFieldExtractor("path")
}
if chunk := p.pathEx.Feed(delta); chunk != "" {
p.path += chunk
}
now := time.Now()
if p.lastParseLen > 0 &&
len(p.buf)-p.lastParseLen < 512 &&
now.Sub(p.lastParseAt) < sandboxProgressMinInterval {
return
}
p.lastParseAt = now
p.lastParseLen = len(p.buf)
obj := parsePartialJSONObject(p.buf)
if obj == nil {
return
}
if path, _ := obj["path"].(string); path != "" && p.path == "" {
p.path = path
}
p.added, p.removed = editArgsLineStats(obj)
p.bytes = len(p.buf)
p.preview = ""
}
func (p *sandboxFileProgress) shouldEmit() bool {
if p.path == "" && p.added == 0 && p.removed == 0 && p.bytes == 0 {
return false
}
pathChanged := p.path != p.lastPath
linesChanged := p.added != p.lastAdded || p.removed != p.lastRemoved
if !p.emitted || pathChanged || linesChanged {
return true
}
if p.bytes == p.lastBytes && p.preview == p.lastPreview {
return false
}
return time.Since(p.lastEmit) >= sandboxProgressMinInterval
}
func (p *sandboxFileProgress) markEmitted() {
p.emitted = true
p.lastPath = p.path
p.lastAdded = p.added
p.lastRemoved = p.removed
p.lastBytes = p.bytes
p.lastPreview = p.preview
p.lastEmit = time.Now()
}
func (p *sandboxFileProgress) payload() map[string]any {
out := map[string]any{
"added_lines": p.added,
"removed_lines": p.removed,
"bytes": p.bytes,
}
if p.path != "" {
out["path"] = p.path
}
if p.preview != "" {
out["preview"] = p.preview
}
return out
}
func countContentLines(s string) int {
if s == "" {
return 0
}
n := strings.Count(s, "\n")
if s[len(s)-1] != '\n' {
n++
}
return n
}
func sandboxContentPreview(s string) string {
if s == "" {
return ""
}
lines := strings.Split(s, "\n")
end := len(lines)
if end > 0 && lines[end-1] == "" {
end--
}
lines = lines[:end]
if len(lines) > sandboxFilePreviewMaxLines {
lines = lines[:sandboxFilePreviewMaxLines]
}
return strings.Join(lines, "\n")
}
func parsePartialJSONObject(s string) map[string]any {
var m map[string]any
if json.Unmarshal([]byte(s), &m) == nil {
return m
}
closed := closePartialJSON(s)
if closed == s {
return nil
}
if json.Unmarshal([]byte(closed), &m) == nil {
return m
}
return nil
}
// closePartialJSON closes an unterminated string and missing braces so a
// streaming prefix can be unmarshalled. Structural characters are ASCII, so
// byte iteration is enough.
func closePartialJSON(s string) string {
inString := false
escaped := false
var stack []byte
for i := 0; i < len(s); i++ {
c := s[i]
if escaped {
escaped = false
continue
}
if inString {
if c == '\\' {
escaped = true
continue
}
if c == '"' {
inString = false
}
continue
}
switch c {
case '"':
inString = true
case '{':
stack = append(stack, '}')
case '[':
stack = append(stack, ']')
case '}', ']':
if n := len(stack); n > 0 && stack[n-1] == c {
stack = stack[:n-1]
}
}
}
if !inString && !escaped && len(stack) == 0 {
return s
}
var b strings.Builder
b.Grow(len(s) + len(stack) + 2)
b.WriteString(s)
if escaped {
b.WriteByte('\\')
}
if inString || escaped {
b.WriteByte('"')
}
for i := len(stack) - 1; i >= 0; i-- {
b.WriteByte(stack[i])
}
return b.String()
}
func editArgsLineStats(obj map[string]any) (added, removed int) {
edits, _ := obj["edits"].([]any)
for _, raw := range edits {
m, ok := raw.(map[string]any)
if !ok {
continue
}
oldS := stringFromMap(m, "old_string", "oldText")
newS := stringFromMap(m, "new_string", "newText")
if oldS == "" && newS == "" {
continue
}
removed += countContentLines(oldS)
added += countContentLines(newS)
}
if added == 0 && removed == 0 {
oldS := stringFromMap(obj, "old_string", "oldText")
newS := stringFromMap(obj, "new_string", "newText")
removed = countContentLines(oldS)
added = countContentLines(newS)
}
return added, removed
}
func stringFromMap(m map[string]any, keys ...string) string {
for _, key := range keys {
if s, ok := m[key].(string); ok {
return s
}
}
return ""
}
@@ -0,0 +1,132 @@
package chat
import (
"encoding/json"
"strings"
"testing"
"time"
)
func TestCountContentLines(t *testing.T) {
cases := []struct {
in string
want int
}{
{"", 0},
{"hello", 1},
{"hello\n", 1},
{"hello\nworld", 2},
{"hello\nworld\n", 2},
{"a\n\nb", 3},
}
for _, tc := range cases {
if got := countContentLines(tc.in); got != tc.want {
t.Errorf("countContentLines(%q) = %d, want %d", tc.in, got, tc.want)
}
}
}
func TestSandboxFileProgressWriteStreamsLineCount(t *testing.T) {
p := newSandboxFileProgress(writeSandboxFileName)
if _, ok := p.Feed(`{"path":"`); ok {
t.Fatal("path still open, should not emit yet")
}
payload, ok := p.Feed(`/workspace/output/deck.py","content":"`)
if !ok {
t.Fatal("expected emit once path is complete")
}
if payload["path"] != "/workspace/output/deck.py" {
t.Fatalf("path = %v", payload["path"])
}
if payload["added_lines"] != 0 {
t.Fatalf("no content yet, added_lines = %v", payload["added_lines"])
}
payload, ok = p.Feed(`line1\nline2\nline3`)
if !ok {
t.Fatal("expected emit when line count grows")
}
if payload["added_lines"] != 3 {
t.Fatalf("added_lines = %v, want 3", payload["added_lines"])
}
preview, _ := payload["preview"].(string)
if !strings.Contains(preview, "line1") || !strings.Contains(preview, "line3") {
t.Fatalf("preview = %q", preview)
}
if _, exists := payload["content"]; exists {
t.Fatal("progress payload must not carry the file body")
}
}
func TestSandboxFileProgressWritePreviewCapsAtTenLines(t *testing.T) {
p := newSandboxFileProgress(writeSandboxFileName)
var body strings.Builder
body.WriteString(`{"path":"/workspace/a.py","content":"`)
for i := 0; i < 15; i++ {
if i > 0 {
body.WriteString(`\n`)
}
body.WriteString("x")
}
body.WriteString(`"}`)
payload, ok := p.Feed(body.String())
if !ok {
t.Fatal("expected emit")
}
if payload["added_lines"] != 15 {
t.Fatalf("added_lines = %v, want 15", payload["added_lines"])
}
preview, _ := payload["preview"].(string)
if strings.Count(preview, "\n")+1 != sandboxFilePreviewMaxLines {
t.Fatalf("preview lines = %q", preview)
}
}
func TestSandboxFileProgressEditCountsOldAndNew(t *testing.T) {
p := newSandboxFileProgress(editSandboxFileName)
payload, ok := p.Feed(
`{"path":"/workspace/a.py","edits":[{"old_string":"a\nb\n","new_string":"a\nb\nc\nd\n"}]}`,
)
if !ok {
t.Fatal("expected emit for complete edit JSON")
}
if payload["path"] != "/workspace/a.py" {
t.Fatalf("path = %v", payload["path"])
}
if payload["removed_lines"] != 2 {
t.Fatalf("removed_lines = %v, want 2", payload["removed_lines"])
}
if payload["added_lines"] != 4 {
t.Fatalf("added_lines = %v, want 4", payload["added_lines"])
}
}
func TestSandboxFileProgressEditPartialJSON(t *testing.T) {
p := newSandboxFileProgress(editSandboxFileName)
sandboxProgressMinInterval = 0
t.Cleanup(func() { sandboxProgressMinInterval = 120 * time.Millisecond })
payload, ok := p.Feed(`{"path":"/workspace/a.py","edits":[{"old_string":"foo","new_string":"foo\nbar\nbaz`)
if !ok {
t.Fatal("expected emit for closable partial JSON")
}
if payload["path"] != "/workspace/a.py" {
t.Fatalf("path = %v", payload["path"])
}
added, _ := payload["added_lines"].(int)
if added < 2 {
t.Fatalf("added_lines = %v, want at least 2 from the partial new_string", payload["added_lines"])
}
}
func TestClosePartialJSON(t *testing.T) {
closed := closePartialJSON(`{"path":"/x","content":"hello`)
var m map[string]any
if err := json.Unmarshal([]byte(closed), &m); err != nil {
t.Fatalf("closed JSON should parse: %v (%s)", err, closed)
}
if m["content"] != "hello" {
t.Fatalf("content = %v", m["content"])
}
}
+64
View File
@@ -0,0 +1,64 @@
package chat
import (
"bufio"
"io"
"strings"
)
// SSEEvent 表示一个 Server-Sent Events 事件
type SSEEvent struct {
Data []byte
Done bool
}
// SSEReader 用于读取 SSE 流
type SSEReader struct {
scanner *bufio.Scanner
}
// NewSSEReader 创建 SSE 读取器
func NewSSEReader(reader io.Reader) *SSEReader {
scanner := bufio.NewScanner(reader)
// 设置更大的缓冲区以处理长行(思维链内容可能很长)
buf := make([]byte, 1024*1024)
scanner.Buffer(buf, 1024*1024)
return &SSEReader{scanner: scanner}
}
// ReadEvent 读取下一个 SSE 事件
func (r *SSEReader) ReadEvent() (*SSEEvent, error) {
for r.scanner.Scan() {
line := r.scanner.Text()
// 空行,跳过
if line == "" {
continue
}
// 检查是否为结束标记
if line == "data: [DONE]" {
return &SSEEvent{Done: true}, nil
}
// 解析 data 行
if strings.HasPrefix(line, "data: ") {
jsonStr := line[6:]
return &SSEEvent{Data: []byte(jsonStr)}, nil
}
// 增强兼容性(data:后面没有紧接着一个空格)
if strings.HasPrefix(line, "data:") {
jsonStr := line[5:]
return &SSEEvent{Data: []byte(jsonStr)}, nil
}
// 其他行(如 event:, id: 等)跳过
}
if err := r.scanner.Err(); err != nil {
return nil, err
}
return nil, io.EOF
}
+36
View File
@@ -0,0 +1,36 @@
package chat
import "github.com/Tencent/WeKnora/internal/types"
// thinkingEmitter owns the "reasoning then answer" hand-off that every
// streaming Chat implementation shares: thinking chunks are forwarded as they
// arrive, and exactly one thinking-done marker is emitted before the first
// answer token (or when the stream ends without one). Centralizing the
// bookkeeping keeps the OpenAI-compatible and Ollama stream loops in sync.
type thinkingEmitter struct {
active bool
}
// emit forwards a reasoning chunk and records that a thinking-done marker is
// still owed.
func (e *thinkingEmitter) emit(ch chan types.StreamResponse, content string) {
e.active = true
ch <- types.StreamResponse{
ResponseType: types.ResponseTypeThinking,
Content: content,
Done: false,
}
}
// finish emits the single thinking-done marker if one is owed. Safe to call
// multiple times; only the first call after an emit sends anything.
func (e *thinkingEmitter) finish(ch chan types.StreamResponse) {
if !e.active {
return
}
e.active = false
ch <- types.StreamResponse{
ResponseType: types.ResponseTypeThinking,
Done: true,
}
}
+214
View File
@@ -0,0 +1,214 @@
package chat
import (
"encoding/json"
"fmt"
"os"
"path/filepath"
"strings"
"sync"
"time"
)
// streamRawDumpDir returns the directory for per-stream raw packet dumps.
// Enabled when WEKNORA_LLM_STREAM_RAW_DUMP_DIR is set, or when
// WEKNORA_LLM_STREAM_RAW_DUMP=1 (defaults to ~/.weknora/investigate/llm-stream).
func streamRawDumpDir() string {
if dir := strings.TrimSpace(os.Getenv("WEKNORA_LLM_STREAM_RAW_DUMP_DIR")); dir != "" {
return dir
}
v := strings.TrimSpace(os.Getenv("WEKNORA_LLM_STREAM_RAW_DUMP"))
if v == "1" || strings.EqualFold(v, "true") || strings.EqualFold(v, "yes") {
home, err := os.UserHomeDir()
if err != nil {
return ""
}
return filepath.Join(home, ".weknora", "investigate", "llm-stream")
}
return ""
}
// streamPacketDumper writes one stream session to a dedicated JSONL file:
// line 1 = request wrapper; following lines = raw provider chunk JSON.
type streamPacketDumper struct {
mu sync.Mutex
file *os.File
path string
model string
seq int
}
func newStreamPacketDumper(modelName string, request any) *streamPacketDumper {
dir := streamRawDumpDir()
if dir == "" {
return nil
}
if err := os.MkdirAll(dir, 0o755); err != nil {
return nil
}
safeModel := strings.Map(func(r rune) rune {
switch {
case r >= 'a' && r <= 'z', r >= 'A' && r <= 'Z', r >= '0' && r <= '9', r == '-', r == '_':
return r
default:
return '_'
}
}, modelName)
if safeModel == "" {
safeModel = "model"
}
name := fmt.Sprintf("llm_stream_%s_%s.jsonl", safeModel, time.Now().Format("20060102T150405.000000000"))
// Defense-in-depth against path traversal: the model name is externally
// controlled, so force the file to live directly under dir and verify the
// cleaned path cannot escape it.
name = filepath.Base(filepath.Clean("/" + name))
dir = filepath.Clean(dir)
path := filepath.Join(dir, name)
if rel, err := filepath.Rel(dir, path); err != nil || rel != name {
return nil
}
f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o644)
if err != nil {
return nil
}
d := &streamPacketDumper{file: f, path: path, model: modelName}
_ = d.writeRequest(request)
return d
}
func (d *streamPacketDumper) writeRequest(request any) error {
line, err := json.Marshal(map[string]any{
"type": "request",
"model": d.model,
"timestamp": time.Now().UTC().Format(time.RFC3339Nano),
"data": request,
})
if err != nil {
return err
}
return d.writeLine(line)
}
func (d *streamPacketDumper) writeLine(line []byte) error {
d.mu.Lock()
defer d.mu.Unlock()
if _, err := d.file.Write(line); err != nil {
return err
}
_, err := d.file.Write([]byte{'\n'})
return err
}
// WritePacketRaw appends one provider chunk as a single JSONL line (valid JSON written as-is).
func (d *streamPacketDumper) WritePacketRaw(raw []byte) {
if d == nil || d.file == nil || len(raw) == 0 {
return
}
raw = bytesTrimSpace(raw)
if len(raw) == 0 {
return
}
d.mu.Lock()
defer d.mu.Unlock()
d.seq++
if json.Valid(raw) {
_, _ = d.file.Write(raw)
_, _ = d.file.Write([]byte{'\n'})
return
}
line, _ := json.Marshal(map[string]any{
"type": "packet",
"seq": d.seq,
"timestamp": time.Now().UTC().Format(time.RFC3339Nano),
"data_raw": string(raw),
})
_, _ = d.file.Write(line)
_, _ = d.file.Write([]byte{'\n'})
}
// WriteError appends a terminal error record (stream read failure, API error, etc.).
func (d *streamPacketDumper) WriteError(message string) {
if d == nil || d.file == nil || strings.TrimSpace(message) == "" {
return
}
line, err := json.Marshal(map[string]any{
"type": "error",
"model": d.model,
"timestamp": time.Now().UTC().Format(time.RFC3339Nano),
"message": message,
})
if err != nil {
return
}
_ = d.writeLine(line)
}
// WriteHTTPError appends a non-2xx HTTP response before any SSE packets.
func (d *streamPacketDumper) WriteHTTPError(statusCode int, body []byte) {
if d == nil || d.file == nil {
return
}
entry := map[string]any{
"type": "http_error",
"model": d.model,
"timestamp": time.Now().UTC().Format(time.RFC3339Nano),
"status_code": statusCode,
}
if len(body) > 0 {
trimmed := bytesTrimSpace(body)
if json.Valid(trimmed) {
var parsed any
if json.Unmarshal(trimmed, &parsed) == nil {
entry["body"] = parsed
} else {
entry["body_raw"] = string(trimmed)
}
} else {
entry["body_raw"] = string(trimmed)
}
}
line, err := json.Marshal(entry)
if err != nil {
return
}
_ = d.writeLine(line)
}
// WritePacket marshals v as one JSON object per line (SDK stream Recv path).
func (d *streamPacketDumper) WritePacket(v any) {
if d == nil || v == nil {
return
}
line, err := json.Marshal(v)
if err != nil {
return
}
d.WritePacketRaw(line)
}
func (d *streamPacketDumper) Path() string {
if d == nil {
return ""
}
return d.path
}
func (d *streamPacketDumper) Close() {
if d == nil || d.file == nil {
return
}
_ = d.file.Close()
d.file = nil
}
func bytesTrimSpace(b []byte) []byte {
return []byte(strings.TrimSpace(string(b)))
}
+172
View File
@@ -0,0 +1,172 @@
package chat
import (
"strings"
"github.com/Tencent/WeKnora/internal/models/provider"
"github.com/sashabaranov/go-openai"
)
// ExtraConfigThinkingControl is the model parameters.extra_config key for
// selecting how ChatOptions.Thinking is translated to provider HTTP fields.
// The accepted values mirror the strings the frontend writes (see
// ModelEditorDialog.vue): "none", "enable_thinking", "thinking_type",
// "chat_template_kwargs".
const ExtraConfigThinkingControl = "thinking_control"
// Wire-format request bodies used by providers that express extended-thinking
// through a non-standard top-level field. They embed the standard OpenAI
// request so all other fields are marshalled unchanged.
// QwenChatCompletionRequest adds Aliyun Qwen's `enable_thinking` boolean.
type QwenChatCompletionRequest struct {
openai.ChatCompletionRequest
EnableThinking *bool `json:"enable_thinking,omitempty"`
}
// ThinkingConfig is the `{ "type": "enabled"|"disabled" }` block used by
// LKEAP / Volcengine style providers.
type ThinkingConfig struct {
Type string `json:"type"`
}
// ThinkingChatCompletionRequest adds the `thinking` object for providers that
// use the `{ "thinking": { "type": ... } }` wire format.
type ThinkingChatCompletionRequest struct {
openai.ChatCompletionRequest
Thinking *ThinkingConfig `json:"thinking,omitempty"`
}
// ThinkingStrategy encodes how ChatOptions.Thinking is mapped onto a provider's
// HTTP request. Apply returns (customBody, useRawHTTP):
// - (nil, false) means "send the standard OpenAI request unchanged" (the
// caller keeps using the SDK path).
// - a non-nil customBody must be sent verbatim over raw HTTP because it
// carries fields the OpenAI SDK would strip.
//
// When opts.Thinking is nil most strategies emit nothing, deferring to the
// model's own default; the exception is enableThinking{alwaysSend: true}
// (Aliyun Qwen), which must always pin the field.
type ThinkingStrategy interface {
Apply(req *openai.ChatCompletionRequest, opts *ChatOptions, isStream bool) (customBody any, useRawHTTP bool)
}
// noThinking sends no thinking-related fields at all.
type noThinking struct{}
func (noThinking) Apply(*openai.ChatCompletionRequest, *ChatOptions, bool) (any, bool) {
return nil, false
}
// enableThinking encodes thinking via Qwen's `enable_thinking` boolean.
//
// - alwaysSend: pin the field even when opts.Thinking is nil (Aliyun Qwen
// thinking models require it on every request; default value is false).
// - disableOnNonStream: force enable_thinking=false for non-stream requests
// (Qwen3 rejects thinking in non-stream mode).
type enableThinking struct {
alwaysSend bool
disableOnNonStream bool
}
func (s enableThinking) Apply(req *openai.ChatCompletionRequest, opts *ChatOptions, isStream bool) (any, bool) {
thinking := false
switch {
case opts != nil && opts.Thinking != nil:
thinking = *opts.Thinking
case !s.alwaysSend:
return nil, false
}
if s.disableOnNonStream && !isStream {
thinking = false
}
qwenReq := QwenChatCompletionRequest{ChatCompletionRequest: *req}
qwenReq.EnableThinking = &thinking
return qwenReq, true
}
// thinkingTypeField encodes thinking via the `{ "thinking": { "type": ... } }`
// object (LKEAP / Volcengine). Emits nothing when opts.Thinking is unset.
type thinkingTypeField struct{}
func (thinkingTypeField) Apply(req *openai.ChatCompletionRequest, opts *ChatOptions, _ bool) (any, bool) {
if opts == nil || opts.Thinking == nil {
return nil, false
}
r := ThinkingChatCompletionRequest{ChatCompletionRequest: *req}
thinkingType := "disabled"
if *opts.Thinking {
thinkingType = "enabled"
}
r.Thinking = &ThinkingConfig{Type: thinkingType}
return r, true
}
// chatTemplateKwargs encodes thinking via the standard request's
// `chat_template_kwargs.enable_thinking` (vLLM / NVIDIA / generic local
// deployments). Emits nothing when opts.Thinking is unset.
type chatTemplateKwargs struct{}
func (chatTemplateKwargs) Apply(req *openai.ChatCompletionRequest, opts *ChatOptions, _ bool) (any, bool) {
if opts == nil || opts.Thinking == nil {
return nil, false
}
req.ChatTemplateKwargs = map[string]interface{}{
"enable_thinking": *opts.Thinking,
}
return req, true
}
// parseThinkingOverride reads extra_config.thinking_control and returns the
// strategy it selects, or nil when unset (the provider adapter's default
// strategy then applies). An unrecognized non-empty value falls back to
// chat_template_kwargs, preserving the legacy default-mode behavior.
func parseThinkingOverride(extraConfig map[string]string) ThinkingStrategy {
if extraConfig == nil {
return nil
}
switch strings.ToLower(strings.TrimSpace(extraConfig[ExtraConfigThinkingControl])) {
case "":
return nil
case "none":
return noThinking{}
case "enable_thinking":
return enableThinking{}
case "thinking_type":
return thinkingTypeField{}
default:
// "chat_template_kwargs" and any unknown non-empty value.
return chatTemplateKwargs{}
}
}
// EffectiveThinkingControl reports the provider field that will carry
// ChatOptions.Thinking. It intentionally shares the same adapter/override
// resolution as the real request path so diagnostics do not guess from the
// frontend selection.
func EffectiveThinkingControl(config *ChatConfig) string {
if config == nil {
return "none"
}
if override := parseThinkingOverride(config.ExtraConfig); override != nil {
return thinkingStrategyName(override)
}
providerName := provider.ProviderName(config.Provider)
if providerName == "" {
providerName = provider.DetectProvider(config.BaseURL)
}
return thinkingStrategyName(resolveProvider(providerName, config.ModelName).Thinking())
}
func thinkingStrategyName(strategy ThinkingStrategy) string {
switch strategy.(type) {
case enableThinking:
return "enable_thinking"
case thinkingTypeField:
return "thinking_type"
case chatTemplateKwargs:
return "chat_template_kwargs"
default:
return "none"
}
}
+144
View File
@@ -0,0 +1,144 @@
package chat
import (
"encoding/json"
"testing"
"github.com/sashabaranov/go-openai"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func ptrBool(b bool) *bool { return &b }
// TestThinkingStrategy_NilThinking verifies the strategies that defer to the
// model default emit nothing when ChatOptions.Thinking is unset.
func TestThinkingStrategy_NilThinking(t *testing.T) {
req := openai.ChatCompletionRequest{Model: "test"}
strategies := []ThinkingStrategy{
noThinking{},
enableThinking{}, // not alwaysSend
thinkingTypeField{},
chatTemplateKwargs{},
}
for _, s := range strategies {
custom, raw := s.Apply(&req, nil, true)
assert.Nil(t, custom, "%T", s)
assert.False(t, raw, "%T", s)
}
}
// TestEnableThinking_QwenSemantics pins the Aliyun Qwen behavior: thinking is
// always sent, defaults to false, and is forced off on non-stream requests.
func TestEnableThinking_QwenSemantics(t *testing.T) {
s := enableThinking{alwaysSend: true, disableOnNonStream: true}
req := openai.ChatCompletionRequest{Model: "qwen3-32b"}
t.Run("non-stream forces false even when requested true", func(t *testing.T) {
custom, raw := s.Apply(&req, &ChatOptions{Thinking: ptrBool(true)}, false)
require.True(t, raw)
qwen, ok := custom.(QwenChatCompletionRequest)
require.True(t, ok)
require.NotNil(t, qwen.EnableThinking)
assert.False(t, *qwen.EnableThinking)
})
t.Run("stream honors requested true", func(t *testing.T) {
custom, raw := s.Apply(&req, &ChatOptions{Thinking: ptrBool(true)}, true)
require.True(t, raw)
qwen := custom.(QwenChatCompletionRequest)
require.NotNil(t, qwen.EnableThinking)
assert.True(t, *qwen.EnableThinking)
})
t.Run("stream defaults to false when unset", func(t *testing.T) {
custom, raw := s.Apply(&req, nil, true)
require.True(t, raw)
qwen := custom.(QwenChatCompletionRequest)
require.NotNil(t, qwen.EnableThinking)
assert.False(t, *qwen.EnableThinking)
})
}
// TestEnableThinking_ExtraConfigSemantics pins the extra_config "enable_thinking"
// override: only sent when explicitly requested.
func TestEnableThinking_ExtraConfigSemantics(t *testing.T) {
s := enableThinking{}
req := openai.ChatCompletionRequest{Model: "qwen3"}
custom, raw := s.Apply(&req, &ChatOptions{Thinking: ptrBool(true)}, true)
require.True(t, raw)
qwen := custom.(QwenChatCompletionRequest)
require.NotNil(t, qwen.EnableThinking)
assert.True(t, *qwen.EnableThinking)
custom, raw = s.Apply(&req, nil, true)
assert.Nil(t, custom)
assert.False(t, raw)
}
func TestThinkingTypeField(t *testing.T) {
s := thinkingTypeField{}
req := openai.ChatCompletionRequest{Model: "ds-v3"}
custom, raw := s.Apply(&req, &ChatOptions{Thinking: ptrBool(false)}, true)
require.True(t, raw)
typed, ok := custom.(ThinkingChatCompletionRequest)
require.True(t, ok)
require.NotNil(t, typed.Thinking)
assert.Equal(t, "disabled", typed.Thinking.Type)
custom, raw = s.Apply(&req, &ChatOptions{Thinking: ptrBool(true)}, true)
require.True(t, raw)
assert.Equal(t, "enabled", custom.(ThinkingChatCompletionRequest).Thinking.Type)
}
func TestChatTemplateKwargs(t *testing.T) {
s := chatTemplateKwargs{}
req := openai.ChatCompletionRequest{Model: "vllm"}
custom, raw := s.Apply(&req, &ChatOptions{Thinking: ptrBool(true)}, true)
require.True(t, raw)
out, ok := custom.(*openai.ChatCompletionRequest)
require.True(t, ok)
assert.Equal(t, true, out.ChatTemplateKwargs["enable_thinking"])
body, err := json.Marshal(custom)
require.NoError(t, err)
assert.Contains(t, string(body), "chat_template_kwargs")
}
func TestParseThinkingOverride(t *testing.T) {
cases := map[string]ThinkingStrategy{
"none": noThinking{},
"enable_thinking": enableThinking{},
"thinking_type": thinkingTypeField{},
"chat_template_kwargs": chatTemplateKwargs{},
"something-unknown": chatTemplateKwargs{}, // legacy default-mode fallback
}
for value, want := range cases {
got := parseThinkingOverride(map[string]string{ExtraConfigThinkingControl: value})
assert.IsType(t, want, got, "value=%q", value)
}
assert.Nil(t, parseThinkingOverride(nil))
assert.Nil(t, parseThinkingOverride(map[string]string{}))
assert.Nil(t, parseThinkingOverride(map[string]string{ExtraConfigThinkingControl: ""}))
}
func TestEffectiveThinkingControl(t *testing.T) {
assert.Equal(t, "enable_thinking", EffectiveThinkingControl(&ChatConfig{
Provider: "aliyun",
ModelName: "qwen3-32b",
}))
assert.Equal(t, "chat_template_kwargs", EffectiveThinkingControl(&ChatConfig{
Provider: "generic",
ModelName: "qwen3",
ExtraConfig: map[string]string{ExtraConfigThinkingControl: "chat_template_kwargs"},
}))
assert.Equal(t, "none", EffectiveThinkingControl(&ChatConfig{
Provider: "generic",
ModelName: "qwen3",
ExtraConfig: map[string]string{ExtraConfigThinkingControl: "none"},
}))
}
+86
View File
@@ -0,0 +1,86 @@
package chat
import (
"context"
"os"
"testing"
"time"
)
func TestWithLLMTimeout_NoParentDeadline_AppliesDefault(t *testing.T) {
ctx, cancel := withLLMTimeout(context.Background(), 50*time.Millisecond)
defer cancel()
dl, ok := ctx.Deadline()
if !ok {
t.Fatalf("expected deadline to be set when parent has none")
}
if remaining := time.Until(dl); remaining <= 0 || remaining > 50*time.Millisecond {
t.Fatalf("unexpected remaining duration: %v", remaining)
}
}
func TestWithLLMTimeout_ShorterParentDeadline_Respected(t *testing.T) {
parent, parentCancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
defer parentCancel()
ctx, cancel := withLLMTimeout(parent, 10*time.Second)
defer cancel()
dl, _ := ctx.Deadline()
if remaining := time.Until(dl); remaining > 30*time.Millisecond {
t.Fatalf("parent shorter deadline should be respected, got remaining=%v", remaining)
}
}
func TestWithLLMTimeout_LongerParentDeadline_NotTruncated(t *testing.T) {
parent, parentCancel := context.WithTimeout(context.Background(), 10*time.Second)
defer parentCancel()
ctx, cancel := withLLMTimeout(parent, 50*time.Millisecond)
defer cancel()
dl, ok := ctx.Deadline()
if !ok {
t.Fatalf("deadline must be set")
}
if remaining := time.Until(dl); remaining < 5*time.Second {
t.Fatalf("parent longer deadline should NOT be truncated by default, got remaining=%v", remaining)
}
}
func TestEnvDurationSeconds(t *testing.T) {
const key = "WEKNORA_TEST_TIMEOUT_SECONDS"
t.Run("unset returns fallback", func(t *testing.T) {
os.Unsetenv(key)
if got := envDurationSeconds(key, 7*time.Second); got != 7*time.Second {
t.Fatalf("got %v", got)
}
})
t.Run("valid value parsed", func(t *testing.T) {
t.Setenv(key, "42")
if got := envDurationSeconds(key, 1*time.Second); got != 42*time.Second {
t.Fatalf("got %v", got)
}
})
t.Run("invalid falls back", func(t *testing.T) {
t.Setenv(key, "not-a-number")
if got := envDurationSeconds(key, 9*time.Second); got != 9*time.Second {
t.Fatalf("got %v", got)
}
})
t.Run("non-positive falls back", func(t *testing.T) {
t.Setenv(key, "0")
if got := envDurationSeconds(key, 9*time.Second); got != 9*time.Second {
t.Fatalf("got %v", got)
}
t.Setenv(key, "-5")
if got := envDurationSeconds(key, 9*time.Second); got != 9*time.Second {
t.Fatalf("got %v", got)
}
})
}
+63
View File
@@ -0,0 +1,63 @@
package chat
import (
"context"
"net/http"
"os"
"strconv"
"strings"
"time"
"github.com/Tencent/WeKnora/internal/models/requestbudget"
secutils "github.com/Tencent/WeKnora/internal/utils"
)
// LLM 调用超时配置。仅作为"上层未设置 deadline 时"的兜底,避免 hung 请求
// 永久阻塞 worker。如果上层 ctx 已经设置了 deadline(无论比默认更短还是更长),
// 都会原样尊重,不再叠加默认超时。可通过环境变量覆盖:
// - WEKNORA_LLM_CHAT_TIMEOUT_SECONDS 非流式调用兜底超时(默认 600s)
// - WEKNORA_LLM_STREAM_TIMEOUT_SECONDS 流式调用兜底超时(默认 1800s)
var (
defaultChatTimeout = envDurationSeconds("WEKNORA_LLM_CHAT_TIMEOUT_SECONDS", 300*time.Second)
defaultStreamTimeout = envDurationSeconds("WEKNORA_LLM_STREAM_TIMEOUT_SECONDS", 600*time.Second)
)
// envDurationSeconds 读取以"秒"为单位的环境变量,解析失败或非正值时回退到 fallback。
func envDurationSeconds(key string, fallback time.Duration) time.Duration {
v := strings.TrimSpace(os.Getenv(key))
if v == "" {
return fallback
}
n, err := strconv.Atoi(v)
if err != nil || n <= 0 {
return fallback
}
return time.Duration(n) * time.Second
}
// withLLMTimeout 仅在上层 ctx 没有 deadline 时附加一个兜底超时;
// 如果上层已显式设置 deadline(无论更短或更长),则原样返回,
// 让调用方对自己的超时策略拥有最终决定权。
func withLLMTimeout(ctx context.Context, d time.Duration) (context.Context, context.CancelFunc) {
if _, ok := ctx.Deadline(); ok {
return ctx, func() {}
}
return context.WithTimeout(ctx, d)
}
// rawHTTPClient is a shared HTTP client for raw HTTP LLM calls with connection-level timeouts.
// Per-request timeout is enforced via context deadline (see defaultChatTimeout / defaultStreamTimeout)
// rather than http.Client.Timeout, so streaming calls are not prematurely terminated.
// Uses SSRFSafeDialContext to prevent DNS rebinding attacks at the connection layer.
var rawHTTPTransport = &http.Transport{
Proxy: http.ProxyFromEnvironment,
DialContext: secutils.SSRFSafeDialContext,
TLSHandshakeTimeout: 10 * time.Second,
IdleConnTimeout: 90 * time.Second,
MaxIdleConnsPerHost: 5,
}
var rawHTTPClient = secutils.NewSSRFSafeHTTPClientWithTransport(
secutils.SSRFSafeHTTPClientConfig{Timeout: 0, MaxRedirects: 10},
requestbudget.Transport{Base: rawHTTPTransport},
)
+46
View File
@@ -0,0 +1,46 @@
package chat
import (
"context"
"strings"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/types"
)
// logUsage emits the standard "[LLM Usage]" line shared by every Chat
// implementation. It is a no-op when usage is nil so callers can pass through
// optional usage blocks without guarding at each call site.
func logUsage(ctx context.Context, model string, u *types.TokenUsage) {
if u == nil {
return
}
purpose, prefixFingerprint := types.LLMCallMetadataFromContext(ctx)
logger.Infof(ctx,
"[LLM Usage] model=%s, purpose=%s, prompt_prefix=%s, prompt_tokens=%d, completion_tokens=%d, "+
"total_tokens=%d, cached_tokens=%d, cache_read_tokens=%d, cache_write_tokens=%d, "+
"cache_miss_tokens=%d, cache_hit_rate=%.1f%%, cache_reported=%t, cache_status=%s%s",
model, purpose, prefixFingerprint, u.PromptTokens, u.CompletionTokens, u.TotalTokens,
u.CachedTokens, u.CacheReadTokens, u.CacheWriteTokens, u.CacheMissTokens,
u.PromptCacheHitRate(), u.CacheReported, u.CacheStatus, usageAttribution(ctx))
}
// usageAttribution renders the ", session_id=…, principal=…" suffix that
// attributes a usage line to the session and terminal principal that
// triggered the call. Calls that run outside a session or without a resolved
// principal (document parsing, title generation, background jobs) render an
// empty suffix, keeping their lines byte-identical to before.
func usageAttribution(ctx context.Context) string {
var b strings.Builder
if sessionID, ok := types.SessionIDFromContext(ctx); ok && sessionID != "" {
b.WriteString(", session_id=")
b.WriteString(sessionID)
}
if principal, ok := types.PrincipalFromContext(ctx); ok {
b.WriteString(", principal=")
b.WriteString(principal.Type)
b.WriteString(":")
b.WriteString(principal.ID)
}
return b.String()
}
@@ -0,0 +1,34 @@
package chat
import (
"context"
"testing"
"github.com/Tencent/WeKnora/internal/types"
)
func TestUsageAttributionEmptyOutsideSessions(t *testing.T) {
// Calls without session or principal context (document parsing, title
// generation, background jobs) must keep the usage line byte-identical.
if got := usageAttribution(context.Background()); got != "" {
t.Fatalf("expected empty suffix, got %q", got)
}
}
func TestUsageAttributionCarriesSessionAndPrincipal(t *testing.T) {
ctx := types.WithSessionID(context.Background(), "sess-123")
ctx = types.WithPrincipal(ctx, types.Principal{Type: types.PrincipalAPIExternalUser, ID: "10000:42"})
got := usageAttribution(ctx)
want := ", session_id=sess-123, principal=api_external_user:10000:42"
if got != want {
t.Fatalf("got %q want %q", got, want)
}
}
func TestUsageAttributionSessionOnly(t *testing.T) {
ctx := types.WithSessionID(context.Background(), "sess-9")
if got := usageAttribution(ctx); got != ", session_id=sess-9" {
t.Fatalf("got %q", got)
}
}
+281
View File
@@ -0,0 +1,281 @@
package embedding
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
"time"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/models/requestbudget"
secutils "github.com/Tencent/WeKnora/internal/utils"
)
const (
// AliyunMultimodalEmbeddingEndpoint 阿里云 DashScope 多模态 Embedding API 端点
AliyunMultimodalEmbeddingEndpoint = "/api/v1/services/embeddings/multimodal-embedding/multimodal-embedding"
)
// AliyunEmbedder implements text vectorization using Aliyun DashScope multimodal embedding API
type AliyunEmbedder struct {
apiKey string
baseURL string
modelName string
truncatePromptTokens int
dimensions int
modelID string
httpClient *http.Client
timeout time.Duration
maxRetries int
customHeaders map[string]string
supportsDimensionOverride bool
EmbedderPooler
}
// SetCustomHeaders 设置用户自定义 HTTP 请求头(类似 OpenAI Python SDK 的 extra_headers)。
func (e *AliyunEmbedder) SetCustomHeaders(headers map[string]string) {
e.customHeaders = headers
}
func (e *AliyunEmbedder) SetSupportsDimensionOverride(supported bool) {
e.supportsDimensionOverride = supported
}
// AliyunEmbedParameters represents the parameters for Aliyun DashScope multimodal embedding
type AliyunEmbedParameters struct {
Dimension int `json:"dimension,omitempty"`
}
// AliyunEmbedRequest represents an Aliyun DashScope multimodal embedding request
type AliyunEmbedRequest struct {
Model string `json:"model"`
Input AliyunEmbedInput `json:"input"`
Parameters *AliyunEmbedParameters `json:"parameters,omitempty"`
}
// AliyunEmbedInput represents the input structure for Aliyun embedding
type AliyunEmbedInput struct {
Contents []AliyunContent `json:"contents"`
}
// AliyunContent represents a single content item in the input
type AliyunContent struct {
Text string `json:"text,omitempty"`
}
// AliyunEmbedResponse represents an Aliyun DashScope embedding response
type AliyunEmbedResponse struct {
Output struct {
Embeddings []struct {
Embedding []float32 `json:"embedding"`
TextIndex int `json:"text_index"`
} `json:"embeddings"`
} `json:"output"`
Usage struct {
TotalTokens int `json:"total_tokens"`
} `json:"usage"`
RequestID string `json:"request_id"`
}
// AliyunErrorResponse represents an error response from Aliyun DashScope
type AliyunErrorResponse struct {
Code string `json:"code"`
Message string `json:"message"`
RequestID string `json:"request_id"`
}
// NewAliyunEmbedder creates a new Aliyun DashScope embedder
func NewAliyunEmbedder(apiKey, baseURL, modelName string,
truncatePromptTokens int, dimensions int, modelID string, pooler EmbedderPooler,
) (*AliyunEmbedder, error) {
if baseURL == "" {
baseURL = "https://dashscope.aliyuncs.com"
}
// Remove trailing slash and any existing path suffix
baseURL = strings.TrimRight(baseURL, "/")
// If baseURL contains /compatible-mode/v1, strip it for multimodal API
if strings.Contains(baseURL, "/compatible-mode/v1") {
baseURL = strings.Replace(baseURL, "/compatible-mode/v1", "", 1)
}
if modelName == "" {
return nil, fmt.Errorf("model name is required")
}
if truncatePromptTokens == 0 {
truncatePromptTokens = 511
}
timeout := 60 * time.Second
if err := validateEmbeddingBaseURL(baseURL); err != nil {
return nil, err
}
return &AliyunEmbedder{
apiKey: apiKey,
baseURL: baseURL,
modelName: modelName,
httpClient: newEmbeddingHTTPClient(timeout),
truncatePromptTokens: truncatePromptTokens,
EmbedderPooler: pooler,
dimensions: dimensions,
modelID: modelID,
timeout: timeout,
maxRetries: 3,
}, nil
}
// Embed converts text to vector
func (e *AliyunEmbedder) Embed(ctx context.Context, text string) ([]float32, error) {
for range 3 {
embeddings, err := e.BatchEmbed(ctx, []string{text})
if err != nil {
return nil, err
}
if len(embeddings) > 0 {
return embeddings[0], nil
}
}
return nil, fmt.Errorf("no embedding returned")
}
func (e *AliyunEmbedder) doRequestWithRetry(ctx context.Context, jsonData []byte) (*http.Response, error) {
ctx = requestbudget.WithModel(ctx, e.modelID)
var resp *http.Response
var err error
url := e.baseURL + AliyunMultimodalEmbeddingEndpoint
for i := 0; i <= e.maxRetries; i++ {
if err := requestbudget.Check(ctx); err != nil {
return nil, err
}
if i > 0 {
backoffTime := time.Duration(1<<uint(i-1)) * time.Second
if backoffTime > 10*time.Second {
backoffTime = 10 * time.Second
}
logger.GetLogger(ctx).
Infof("AliyunEmbedder retrying request (%d/%d), waiting %v", i, e.maxRetries, backoffTime)
select {
case <-time.After(backoffTime):
case <-ctx.Done():
return nil, ctx.Err()
}
}
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(jsonData))
if err != nil {
logger.GetLogger(ctx).Errorf("AliyunEmbedder failed to create request: %v", err)
continue
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+e.apiKey)
secutils.ApplyCustomHeaders(req, e.customHeaders)
resp, err = e.httpClient.Do(req)
if stopProviderRetry(ctx, err) {
return nil, err
}
if err == nil {
return resp, nil
}
logger.GetLogger(ctx).Errorf("AliyunEmbedder request failed (attempt %d/%d): %v", i+1, e.maxRetries+1, err)
}
return nil, err
}
func (e *AliyunEmbedder) BatchEmbed(ctx context.Context, texts []string) ([][]float32, error) {
// Build contents array from texts
contents := make([]AliyunContent, 0, len(texts))
for _, text := range texts {
contents = append(contents, AliyunContent{Text: text})
}
// Create request body
reqBody := AliyunEmbedRequest{
Model: e.modelName,
Input: AliyunEmbedInput{
Contents: contents,
},
}
if e.supportsDimensionsParam() {
reqBody.Parameters = &AliyunEmbedParameters{Dimension: e.dimensions}
}
jsonData, err := json.Marshal(reqBody)
if err != nil {
logger.GetLogger(ctx).Errorf("AliyunEmbedder BatchEmbed marshal request error: %v", err)
return nil, fmt.Errorf("marshal request: %w", err)
}
resp, err := e.doRequestWithRetry(ctx, jsonData)
if err != nil {
logger.GetLogger(ctx).Errorf("AliyunEmbedder BatchEmbed send request error: %v", err)
return nil, fmt.Errorf("send request: %w", err)
}
if resp.Body != nil {
defer resp.Body.Close()
}
body, err := io.ReadAll(resp.Body)
if err != nil {
logger.GetLogger(ctx).Errorf("AliyunEmbedder BatchEmbed read response error: %v", err)
return nil, fmt.Errorf("read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
// Try to parse error response
var errResp AliyunErrorResponse
if json.Unmarshal(body, &errResp) == nil && errResp.Message != "" {
logger.GetLogger(ctx).Errorf("AliyunEmbedder BatchEmbed API error: %s - %s", errResp.Code, errResp.Message)
return nil, fmt.Errorf("API error: %s - %s", errResp.Code, errResp.Message)
}
logger.GetLogger(ctx).Errorf("AliyunEmbedder BatchEmbed API error: Http Status %s", resp.Status)
return nil, fmt.Errorf("BatchEmbed API error: Http Status %s", resp.Status)
}
// Parse response
var response AliyunEmbedResponse
if err := json.Unmarshal(body, &response); err != nil {
logger.GetLogger(ctx).Errorf("AliyunEmbedder BatchEmbed unmarshal response error: %v", err)
return nil, fmt.Errorf("unmarshal response: %w", err)
}
// Extract embedding vectors, preserving order by text_index
embeddings := make([][]float32, len(texts))
for _, emb := range response.Output.Embeddings {
if emb.TextIndex >= 0 && emb.TextIndex < len(embeddings) {
embeddings[emb.TextIndex] = emb.Embedding
}
}
return embeddings, nil
}
// GetModelName returns the model name
func (e *AliyunEmbedder) GetModelName() string {
return e.modelName
}
func (e *AliyunEmbedder) supportsDimensionsParam() bool {
return e.supportsDimensionOverride && e.dimensions > 0
}
// GetDimensions returns the vector dimensions
func (e *AliyunEmbedder) GetDimensions() int {
return e.dimensions
}
// GetModelID returns the model ID
func (e *AliyunEmbedder) GetModelID() string {
return e.modelID
}
+199
View File
@@ -0,0 +1,199 @@
package embedding
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"time"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/models/requestbudget"
secutils "github.com/Tencent/WeKnora/internal/utils"
)
// AzureOpenAIEmbedder implements text vectorization using Azure OpenAI API
type AzureOpenAIEmbedder struct {
apiKey string
baseURL string
modelName string
truncatePromptTokens int
dimensions int
modelID string
apiVersion string
httpClient *http.Client
maxRetries int
customHeaders map[string]string
supportsDimensionOverride bool
EmbedderPooler
}
// SetCustomHeaders 设置用户自定义 HTTP 请求头(类似 OpenAI Python SDK 的 extra_headers)。
func (e *AzureOpenAIEmbedder) SetCustomHeaders(headers map[string]string) {
e.customHeaders = headers
}
func (e *AzureOpenAIEmbedder) SetSupportsDimensionOverride(supported bool) {
e.supportsDimensionOverride = supported
}
type azureOpenAIEmbedRequest struct {
Model string `json:"model"`
Input []string `json:"input"`
EncodingFormat string `json:"encoding_format,omitempty"`
Dimensions int `json:"dimensions,omitempty"`
}
// NewAzureOpenAIEmbedder creates a new Azure OpenAI embedder
func NewAzureOpenAIEmbedder(apiKey, baseURL, modelName string,
truncatePromptTokens int, dimensions int, modelID string,
apiVersion string, pooler EmbedderPooler,
) (*AzureOpenAIEmbedder, error) {
if baseURL == "" {
return nil, fmt.Errorf("Azure resource endpoint (base URL) is required")
}
if modelName == "" {
return nil, fmt.Errorf("deployment name (model name) is required")
}
if apiVersion == "" {
apiVersion = "2024-10-21"
}
if truncatePromptTokens == 0 {
truncatePromptTokens = 511
}
if err := validateEmbeddingBaseURL(baseURL); err != nil {
return nil, err
}
return &AzureOpenAIEmbedder{
apiKey: apiKey,
baseURL: baseURL,
modelName: modelName,
truncatePromptTokens: truncatePromptTokens,
dimensions: dimensions,
modelID: modelID,
apiVersion: apiVersion,
httpClient: newEmbeddingHTTPClient(60 * time.Second),
maxRetries: 3,
EmbedderPooler: pooler,
}, nil
}
func (e *AzureOpenAIEmbedder) Embed(ctx context.Context, text string) ([]float32, error) {
for range 3 {
embeddings, err := e.BatchEmbed(ctx, []string{text})
if err != nil {
return nil, err
}
if len(embeddings) > 0 {
return embeddings[0], nil
}
}
return nil, fmt.Errorf("no embedding returned")
}
func (e *AzureOpenAIEmbedder) BatchEmbed(ctx context.Context, texts []string) ([][]float32, error) {
reqBody := azureOpenAIEmbedRequest{
Model: e.modelName,
Input: texts,
EncodingFormat: "float",
}
if e.supportsDimensionsParam() {
reqBody.Dimensions = e.dimensions
}
jsonData, err := json.Marshal(reqBody)
if err != nil {
return nil, fmt.Errorf("marshal request: %w", err)
}
logger.GetLogger(ctx).Debugf("AzureOpenAIEmbedder BatchEmbed: model=%s, input_count=%d",
e.modelName, len(texts))
resp, err := e.doRequestWithRetry(ctx, jsonData)
if err != nil {
return nil, fmt.Errorf("send request: %w", err)
}
if resp.Body != nil {
defer resp.Body.Close()
}
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
bodyStr := string(body)
if len(bodyStr) > 1000 {
bodyStr = bodyStr[:1000] + "... (truncated)"
}
return nil, fmt.Errorf("Azure Embedding API error: Http Status %s, Response: %s", resp.Status, bodyStr)
}
var response OpenAIEmbedResponse
if err := json.Unmarshal(body, &response); err != nil {
return nil, fmt.Errorf("unmarshal response: %w", err)
}
embeddings := make([][]float32, 0, len(response.Data))
for _, data := range response.Data {
embeddings = append(embeddings, data.Embedding)
}
return embeddings, nil
}
func (e *AzureOpenAIEmbedder) doRequestWithRetry(ctx context.Context, jsonData []byte) (*http.Response, error) {
ctx = requestbudget.WithModel(ctx, e.modelID)
url := fmt.Sprintf("%s/openai/deployments/%s/embeddings?api-version=%s",
e.baseURL, e.modelName, e.apiVersion)
var resp *http.Response
var err error
for i := 0; i <= e.maxRetries; i++ {
if err := requestbudget.Check(ctx); err != nil {
return nil, err
}
if i > 0 {
backoffTime := time.Duration(1<<uint(i-1)) * time.Second
if backoffTime > 10*time.Second {
backoffTime = 10 * time.Second
}
select {
case <-time.After(backoffTime):
case <-ctx.Done():
return nil, ctx.Err()
}
}
req, reqErr := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(jsonData))
if reqErr != nil {
err = reqErr
continue
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("api-key", e.apiKey)
secutils.ApplyCustomHeaders(req, e.customHeaders)
resp, err = e.httpClient.Do(req)
if stopProviderRetry(ctx, err) {
return nil, err
}
if err == nil {
return resp, nil
}
}
return nil, err
}
func (e *AzureOpenAIEmbedder) supportsDimensionsParam() bool {
return e.supportsDimensionOverride && e.dimensions > 0
}
func (e *AzureOpenAIEmbedder) GetModelName() string { return e.modelName }
func (e *AzureOpenAIEmbedder) GetDimensions() int { return e.dimensions }
func (e *AzureOpenAIEmbedder) GetModelID() string { return e.modelID }
@@ -0,0 +1,156 @@
package embedding
import (
"bytes"
"context"
"encoding/json"
"io"
"net/http"
"testing"
)
func TestAzureOpenAIEmbedderBatchEmbedSendsConfiguredDimensions(t *testing.T) {
t.Parallel()
var requestBody map[string]any
transport := roundTripFunc(func(r *http.Request) (*http.Response, error) {
if r.Method != http.MethodPost {
t.Fatalf("expected POST request, got %s", r.Method)
}
if got, want := r.URL.String(),
"https://example-resource.openai.azure.com/openai/deployments/text-embedding-3-large-deployment/embeddings?api-version=2024-10-21"; got != want {
t.Fatalf("unexpected request path: got %s want %s", got, want)
}
if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
t.Fatalf("decode request body: %v", err)
}
return &http.Response{
StatusCode: http.StatusOK,
Header: make(http.Header),
Body: io.NopCloser(bytes.NewBufferString(`{"data":[{"embedding":[0.1,0.2],"index":0}]}`)),
}, nil
})
embedder, err := NewAzureOpenAIEmbedder(
"test-key",
"https://example-resource.openai.azure.com",
"text-embedding-3-large-deployment",
511,
256,
"text-embedding-3-large",
"2024-10-21",
nil,
)
if err != nil {
t.Fatalf("create embedder: %v", err)
}
embedder.SetSupportsDimensionOverride(true)
embedder.httpClient = &http.Client{Transport: transport}
if _, err := embedder.BatchEmbed(context.Background(), []string{"hello"}); err != nil {
t.Fatalf("BatchEmbed returned error: %v", err)
}
got, ok := requestBody["dimensions"]
if !ok {
t.Fatalf("expected request body to include dimensions, got %v", requestBody)
}
if got != float64(256) {
t.Fatalf("unexpected dimensions value: got %v want 256", got)
}
}
func TestAzureOpenAIEmbedderBatchEmbedOmitsDimensionsByDefault(t *testing.T) {
t.Parallel()
var requestBody map[string]any
transport := roundTripFunc(func(r *http.Request) (*http.Response, error) {
if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
t.Fatalf("decode request body: %v", err)
}
return &http.Response{
StatusCode: http.StatusOK,
Header: make(http.Header),
Body: io.NopCloser(bytes.NewBufferString(`{"data":[{"embedding":[0.1,0.2],"index":0}]}`)),
}, nil
})
embedder, err := NewAzureOpenAIEmbedder(
"test-key",
"https://example-resource.openai.azure.com",
"ada-002-deployment",
511,
1536,
"text-embedding-ada-002",
"2024-10-21",
nil,
)
if err != nil {
t.Fatalf("create embedder: %v", err)
}
embedder.httpClient = &http.Client{Transport: transport}
if _, err := embedder.BatchEmbed(context.Background(), []string{"hello"}); err != nil {
t.Fatalf("BatchEmbed returned error: %v", err)
}
if _, ok := requestBody["dimensions"]; ok {
t.Fatalf("expected request body to omit dimensions for fixed-size model, got %v", requestBody)
}
}
func TestAzureOpenAIEmbedderBatchEmbedSendsDimensionsWhenOverrideEnabledRegardlessOfAPIVersion(t *testing.T) {
t.Parallel()
var requestBody map[string]any
transport := roundTripFunc(func(r *http.Request) (*http.Response, error) {
if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
t.Fatalf("decode request body: %v", err)
}
return &http.Response{
StatusCode: http.StatusOK,
Header: make(http.Header),
Body: io.NopCloser(bytes.NewBufferString(`{"data":[{"embedding":[0.1,0.2],"index":0}]}`)),
}, nil
})
embedder, err := NewAzureOpenAIEmbedder(
"test-key",
"https://example-resource.openai.azure.com",
"text-embedding-3-large-deployment",
511,
256,
"text-embedding-3-large",
"2024-02-15-preview",
nil,
)
if err != nil {
t.Fatalf("create embedder: %v", err)
}
embedder.SetSupportsDimensionOverride(true)
embedder.httpClient = &http.Client{Transport: transport}
if _, err := embedder.BatchEmbed(context.Background(), []string{"hello"}); err != nil {
t.Fatalf("BatchEmbed returned error: %v", err)
}
got, ok := requestBody["dimensions"]
if !ok {
t.Fatalf("expected request body to include dimensions when override is enabled, got %v", requestBody)
}
if got != float64(256) {
t.Fatalf("unexpected dimensions value: got %v want 256", got)
}
}
type roundTripFunc func(*http.Request) (*http.Response, error)
func (f roundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) {
return f(r)
}
+112
View File
@@ -0,0 +1,112 @@
package embedding
import (
"context"
"fmt"
"os"
"strconv"
"sync"
"github.com/Tencent/WeKnora/internal/models/utils"
"github.com/panjf2000/ants/v2"
)
type batchEmbedder struct {
pool *ants.Pool
}
func NewBatchEmbedder(pool *ants.Pool) EmbedderPooler {
return &batchEmbedder{pool: pool}
}
type textEmbedding struct {
text string
results []float32
}
func (e *batchEmbedder) BatchEmbedWithPool(ctx context.Context, model Embedder, texts []string) ([][]float32, error) {
// Create goroutine pool for concurrent processing of document chunks
var wg sync.WaitGroup
var mu sync.Mutex // For synchronizing access to error
var firstErr error // Record the first error that occurs
batchSizeStr := os.Getenv("BATCH_EMBED_SIZE")
if batchSizeStr == "" {
batchSizeStr = "5"
}
batchSize, err := strconv.Atoi(batchSizeStr)
if err != nil {
return nil, err
}
if batchSize <= 0 {
return nil, fmt.Errorf("embedding batch size must be positive")
}
textEmbeddings := utils.MapSlice(texts, func(text string) *textEmbedding {
return &textEmbedding{text: text}
})
// Function to process each document chunk
processChunk := func(texts []*textEmbedding) func() {
return func() {
defer wg.Done()
// If an error has already occurred, don't continue processing
mu.Lock()
failed := firstErr != nil
mu.Unlock()
if failed {
return
}
// Embed text
embedding, err := model.BatchEmbed(ctx, utils.MapSlice(texts, func(text *textEmbedding) string {
return text.text
}))
if err != nil {
mu.Lock()
if firstErr == nil {
firstErr = err
}
mu.Unlock()
return
}
if len(embedding) != len(texts) {
mu.Lock()
if firstErr == nil {
firstErr = fmt.Errorf("embedding model returned %d embeddings for %d inputs", len(embedding), len(texts))
}
mu.Unlock()
return
}
mu.Lock()
for i, text := range texts {
if text == nil {
continue
}
text.results = embedding[i]
}
mu.Unlock()
}
}
// Submit all tasks to the goroutine pool
for _, texts := range utils.ChunkSlice(textEmbeddings, batchSize) {
wg.Add(1)
err := e.pool.Submit(processChunk(texts))
if err != nil {
wg.Done()
wg.Wait()
return nil, err
}
}
// Wait for all tasks to complete
wg.Wait()
// Check if any errors occurred
if firstErr != nil {
return nil, firstErr
}
results := utils.MapSlice(textEmbeddings, func(text *textEmbedding) []float32 {
return text.results
})
return results, nil
}
@@ -0,0 +1,66 @@
package embedding
import (
"context"
"github.com/Tencent/WeKnora/internal/models/limiter"
)
// Embedding is the highest-volume background model call: document ingestion
// vectorises every chunk, so a single batch upload can burst the whole worker
// pool against one embedding provider. Like chat and vlm, embedding is governed
// at the client layer via the shared per-model concurrency governor. Only
// background (asynq worker) calls are throttled — see limiter.Gate /
// types.IsBackgroundTask; interactive query embedding is never gated.
//
// Placement note: unlike chat/vlm (outermost), this wrapper sits INNERMOST —
// directly around the real embedder, BELOW the debug/langfuse decorators.
// BatchEmbedWithPool fans a batch out into per-sub-batch BatchEmbed calls
// through the pooler, and the pooler invokes BatchEmbed on whichever Embedder
// was threaded down as `model`. Sitting innermost is what routes those
// per-sub-batch provider round-trips back through Gate, so the semaphore bounds
// real concurrent provider calls rather than one coarse per-document unit. The
// trade-off is that Gate wait time is included in debug/langfuse timing, which
// is acceptable for background ingestion.
type concurrencyEmbedder struct {
inner Embedder
// limit is this model's configured per-model background cap; 0 falls back
// to the process-wide default (see limiter.GateN).
limit int
}
func (w *concurrencyEmbedder) Embed(ctx context.Context, text string) ([]float32, error) {
release := limiter.GateNamedN(ctx, w.inner.GetModelID(), w.inner.GetModelName(), w.limit)
defer release()
return w.inner.Embed(ctx, text)
}
func (w *concurrencyEmbedder) BatchEmbed(ctx context.Context, texts []string) ([][]float32, error) {
release := limiter.GateNamedN(ctx, w.inner.GetModelID(), w.inner.GetModelName(), w.limit)
defer release()
return w.inner.BatchEmbed(ctx, texts)
}
// BatchEmbedWithPool threads THIS wrapper down as the model so the pooler's
// per-sub-batch callbacks land on our gated BatchEmbed above, rather than on
// the raw embedder. The wait for each slot is held only around the actual
// per-sub-batch provider round-trip.
func (w *concurrencyEmbedder) BatchEmbedWithPool(
ctx context.Context, model Embedder, texts []string,
) ([][]float32, error) {
return w.inner.BatchEmbedWithPool(ctx, w, texts)
}
func (w *concurrencyEmbedder) GetModelName() string { return w.inner.GetModelName() }
func (w *concurrencyEmbedder) GetDimensions() int { return w.inner.GetDimensions() }
func (w *concurrencyEmbedder) GetModelID() string { return w.inner.GetModelID() }
// wrapEmbeddingConcurrency installs the background concurrency governor directly
// around the real embedder. Always applied; a cheap passthrough when no limiter
// is installed or the call is interactive.
func wrapEmbeddingConcurrency(e Embedder, limit int) Embedder {
if e == nil {
return e
}
return &concurrencyEmbedder{inner: e, limit: limit}
}
@@ -0,0 +1,232 @@
package embedding
import (
"context"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/Tencent/WeKnora/internal/models/limiter"
"github.com/Tencent/WeKnora/internal/types"
"github.com/panjf2000/ants/v2"
)
// fakeEmbedder blocks every Embed/BatchEmbed call until release is closed and
// records the max number of calls it ever saw in flight simultaneously, so
// tests can assert the governor's per-model concurrency bound.
type fakeEmbedder struct {
id string
pooler EmbedderPooler // optional: exercises the BatchEmbedWithPool fan-out
inFlight int32
maxSeen int32
enter chan struct{} // one signal per call that reaches the provider
release chan struct{} // closed to unblock all (current + future) calls
}
func newFakeEmbedder(id string) *fakeEmbedder {
return &fakeEmbedder{
id: id,
enter: make(chan struct{}, 256),
release: make(chan struct{}),
}
}
func (f *fakeEmbedder) track() {
n := atomic.AddInt32(&f.inFlight, 1)
for {
old := atomic.LoadInt32(&f.maxSeen)
if n <= old || atomic.CompareAndSwapInt32(&f.maxSeen, old, n) {
break
}
}
f.enter <- struct{}{}
<-f.release
atomic.AddInt32(&f.inFlight, -1)
}
func (f *fakeEmbedder) Embed(ctx context.Context, _ string) ([]float32, error) {
f.track()
return []float32{1}, nil
}
func (f *fakeEmbedder) BatchEmbed(ctx context.Context, _ []string) ([][]float32, error) {
f.track()
return [][]float32{{1}}, nil
}
func (f *fakeEmbedder) BatchEmbedWithPool(
ctx context.Context, model Embedder, texts []string,
) ([][]float32, error) {
if f.pooler != nil {
return f.pooler.BatchEmbedWithPool(ctx, model, texts)
}
return model.BatchEmbed(ctx, texts)
}
func (f *fakeEmbedder) GetModelName() string { return f.id }
func (f *fakeEmbedder) GetDimensions() int { return 1 }
func (f *fakeEmbedder) GetModelID() string { return f.id }
// TestConcurrencyEmbedderBackgroundGated verifies background BatchEmbed calls
// are capped at the per-model limit.
func TestConcurrencyEmbedderBackgroundGated(t *testing.T) {
t.Cleanup(func() { limiter.SetGovernor(nil, 0) })
limiter.SetGovernor(limiter.NewLocalLimiter(), 2)
f := newFakeEmbedder("emb-bg")
w := wrapEmbeddingConcurrency(f, 0)
ctx := types.WithBackgroundTask(context.Background())
const n = 5
var wg sync.WaitGroup
for range n {
wg.Add(1)
go func() {
defer wg.Done()
_, _ = w.BatchEmbed(ctx, []string{"x"})
}()
}
// With limit=2 exactly two calls may reach the provider concurrently.
for i := range 2 {
select {
case <-f.enter:
case <-time.After(2 * time.Second):
t.Fatalf("expected call %d to enter, inFlight=%d", i, atomic.LoadInt32(&f.inFlight))
}
}
// A third must not sneak in while both slots are held.
select {
case <-f.enter:
t.Fatal("a third call entered while limit=2 slots were held")
case <-time.After(150 * time.Millisecond):
}
close(f.release)
wg.Wait()
if got := atomic.LoadInt32(&f.maxSeen); got > 2 {
t.Fatalf("max in-flight %d exceeded limit 2", got)
}
}
// TestConcurrencyEmbedderPerModelLimitOverridesDefault verifies a model's own
// configured limit takes precedence over the process-wide default.
func TestConcurrencyEmbedderPerModelLimitOverridesDefault(t *testing.T) {
t.Cleanup(func() { limiter.SetGovernor(nil, 0) })
// Global default is generous (10), but this model is pinned to 1.
limiter.SetGovernor(limiter.NewLocalLimiter(), 10)
f := newFakeEmbedder("emb-permodel")
w := wrapEmbeddingConcurrency(f, 1)
ctx := types.WithBackgroundTask(context.Background())
const n = 3
var wg sync.WaitGroup
for range n {
wg.Add(1)
go func() {
defer wg.Done()
_, _ = w.BatchEmbed(ctx, []string{"x"})
}()
}
// Only one may be in flight because the per-model limit is 1.
select {
case <-f.enter:
case <-time.After(2 * time.Second):
t.Fatalf("expected one call to enter, inFlight=%d", atomic.LoadInt32(&f.inFlight))
}
select {
case <-f.enter:
t.Fatal("a second call entered while per-model limit=1 slot was held")
case <-time.After(150 * time.Millisecond):
}
close(f.release)
wg.Wait()
if got := atomic.LoadInt32(&f.maxSeen); got > 1 {
t.Fatalf("max in-flight %d exceeded per-model limit 1", got)
}
}
// TestConcurrencyEmbedderInteractiveNotGated verifies interactive calls bypass
// the governor entirely, even at limit 1.
func TestConcurrencyEmbedderInteractiveNotGated(t *testing.T) {
t.Cleanup(func() { limiter.SetGovernor(nil, 0) })
limiter.SetGovernor(limiter.NewLocalLimiter(), 1)
f := newFakeEmbedder("emb-interactive")
w := wrapEmbeddingConcurrency(f, 0)
ctx := context.Background() // no background marker
const n = 3
var wg sync.WaitGroup
for range n {
wg.Add(1)
go func() {
defer wg.Done()
_, _ = w.Embed(ctx, "x")
}()
}
// All three must be able to run at once despite limit=1.
for i := range n {
select {
case <-f.enter:
case <-time.After(2 * time.Second):
t.Fatalf("interactive call %d did not enter (should be ungated), inFlight=%d",
i, atomic.LoadInt32(&f.inFlight))
}
}
close(f.release)
wg.Wait()
}
// TestConcurrencyEmbedderPoolFanOutGated verifies that BatchEmbedWithPool's
// per-sub-batch provider round-trips are individually gated — the reason the
// wrapper sits innermost.
func TestConcurrencyEmbedderPoolFanOutGated(t *testing.T) {
t.Setenv("BATCH_EMBED_SIZE", "1") // one provider round-trip per text
t.Cleanup(func() { limiter.SetGovernor(nil, 0) })
limiter.SetGovernor(limiter.NewLocalLimiter(), 2)
pool, err := ants.NewPool(16)
if err != nil {
t.Fatalf("new pool: %v", err)
}
defer pool.Release()
f := newFakeEmbedder("emb-pool")
f.pooler = NewBatchEmbedder(pool)
w := wrapEmbeddingConcurrency(f, 0)
ctx := types.WithBackgroundTask(context.Background())
done := make(chan struct{})
go func() {
defer close(done)
_, _ = w.BatchEmbedWithPool(ctx, w, []string{"a", "b", "c", "d", "e"})
}()
for i := range 2 {
select {
case <-f.enter:
case <-time.After(2 * time.Second):
t.Fatalf("expected sub-batch %d to enter, inFlight=%d", i, atomic.LoadInt32(&f.inFlight))
}
}
select {
case <-f.enter:
t.Fatal("a third sub-batch entered while limit=2 slots were held")
case <-time.After(150 * time.Millisecond):
}
close(f.release)
select {
case <-done:
case <-time.After(2 * time.Second):
t.Fatal("BatchEmbedWithPool did not complete after release")
}
if got := atomic.LoadInt32(&f.maxSeen); got > 2 {
t.Fatalf("max in-flight sub-batches %d exceeded limit 2", got)
}
}
@@ -0,0 +1,47 @@
package embedding
import (
"testing"
"github.com/Tencent/WeKnora/internal/types"
)
func TestConfigFromModel(t *testing.T) {
m := &types.Model{
ID: "emb-1",
Name: "text-embedding-3-small",
Source: types.ModelSourceRemote,
Parameters: types.ModelParameters{
BaseURL: "https://api.example.com/v1",
APIKey: "sk-xxx",
Provider: "openai",
EmbeddingParameters: types.EmbeddingParameters{
Dimension: 1536,
TruncatePromptTokens: 512,
SupportsDimensionOverride: true,
},
ExtraConfig: map[string]string{"region": "us-east"},
CustomHeaders: map[string]string{"X-Gateway": "g1"},
},
}
cfg := ConfigFromModel(m, "app", "secret")
if cfg.ModelID != "emb-1" || cfg.ModelName != "text-embedding-3-small" {
t.Errorf("identity mismatch: %+v", cfg)
}
if cfg.Dimensions != 1536 || cfg.TruncatePromptTokens != 512 {
t.Errorf("embedding params mismatch: %+v", cfg)
}
if !cfg.SupportsDimensionOverride {
t.Errorf("SupportsDimensionOverride not propagated: %+v", cfg)
}
if cfg.CustomHeaders["X-Gateway"] != "g1" {
t.Errorf("CustomHeaders not propagated: %+v", cfg.CustomHeaders)
}
if cfg.ExtraConfig["region"] != "us-east" {
t.Errorf("ExtraConfig not propagated: %+v", cfg.ExtraConfig)
}
if cfg.AppID != "app" || cfg.AppSecret != "secret" {
t.Errorf("cloud creds mismatch: %+v", cfg)
}
}
@@ -0,0 +1,70 @@
package embedding
import "testing"
func TestEmbeddingDimensionsParamRequiresExplicitOverride(t *testing.T) {
tests := []struct {
name string
supports func(bool) bool
}{
{
name: "aliyun",
supports: func(enabled bool) bool {
e := &AliyunEmbedder{dimensions: 1024}
e.SetSupportsDimensionOverride(enabled)
return e.supportsDimensionsParam()
},
},
{
name: "nvidia",
supports: func(enabled bool) bool {
e := &NvidiaEmbedder{dimensions: 1024}
e.SetSupportsDimensionOverride(enabled)
return e.supportsDimensionsParam()
},
},
{
name: "volcengine",
supports: func(enabled bool) bool {
e := &VolcengineEmbedder{dimensions: 1024}
e.SetSupportsDimensionOverride(enabled)
return e.supportsDimensionsParam()
},
},
{
name: "zhipu",
supports: func(enabled bool) bool {
e := &ZhipuEmbedder{dimensions: 1024}
e.SetSupportsDimensionOverride(enabled)
return e.supportsDimensionsParam()
},
},
{
name: "azure_openai",
supports: func(enabled bool) bool {
e := &AzureOpenAIEmbedder{dimensions: 1024}
e.SetSupportsDimensionOverride(enabled)
return e.supportsDimensionsParam()
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := tt.supports(false); got {
t.Fatalf("supportsDimensionsParam() = true with override disabled")
}
if got := tt.supports(true); !got {
t.Fatalf("supportsDimensionsParam() = false with override enabled")
}
})
}
}
func TestEmbeddingDimensionsParamRequiresPositiveDimension(t *testing.T) {
e := &OpenAIEmbedder{dimensions: 0}
e.SetSupportsDimensionOverride(true)
if e.supportsDimensionsParam() {
t.Fatal("expected dimensions param to be omitted when configured dimension is zero")
}
}
+281
View File
@@ -0,0 +1,281 @@
package embedding
import (
"context"
"fmt"
"strings"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/models/provider"
"github.com/Tencent/WeKnora/internal/models/utils/ollama"
"github.com/Tencent/WeKnora/internal/tracing/langfuse"
"github.com/Tencent/WeKnora/internal/types"
)
// Embedder defines the interface for text vectorization
type Embedder interface {
// Embed converts text to vector
Embed(ctx context.Context, text string) ([]float32, error)
// BatchEmbed converts multiple texts to vectors in batch
BatchEmbed(ctx context.Context, texts []string) ([][]float32, error)
// GetModelName returns the model name
GetModelName() string
// GetDimensions returns the vector dimensions
GetDimensions() int
// GetModelID returns the model ID
GetModelID() string
EmbedderPooler
}
type EmbedderPooler interface {
BatchEmbedWithPool(ctx context.Context, model Embedder, texts []string) ([][]float32, error)
}
// EmbedderType represents the embedder type
type EmbedderType string
// Config represents the embedder configuration
type Config struct {
Source types.ModelSource `json:"source"`
BaseURL string `json:"base_url"`
ModelName string `json:"model_name"`
APIKey string `json:"api_key"`
TruncatePromptTokens int `json:"truncate_prompt_tokens"`
Dimensions int `json:"dimensions"`
SupportsDimensionOverride bool `json:"supports_dimension_override"`
ModelID string `json:"model_id"`
Provider string `json:"provider"`
// MaxConcurrency caps concurrent background calls to this model; 0 falls
// back to the process-wide default (see limiter.GateN).
MaxConcurrency int `json:"max_concurrency"`
ExtraConfig map[string]string `json:"extra_config"`
// CustomHeaders 允许在调用远程 API 时附加自定义 HTTP 请求头(类似 OpenAI Python SDK 的 extra_headers)。
CustomHeaders map[string]string `json:"custom_headers"`
AppID string
AppSecret string // 加密值,工厂函数调用方传入,使用前已解密
}
// ConfigFromModel 根据 types.Model 构造 embedding.Config。
// 生产路径(从 DB 拉起)和测试连接路径(临时表单)共享这份映射。
// appID / appSecret 是已解密的 WeKnoraCloud 凭证,调用方负责传入。
func ConfigFromModel(m *types.Model, appID, appSecret string) Config {
if m == nil {
return Config{}
}
return Config{
Source: m.Source,
BaseURL: m.Parameters.BaseURL,
APIKey: m.Parameters.APIKey,
ModelID: m.ID,
ModelName: m.Name,
Dimensions: m.Parameters.EmbeddingParameters.Dimension,
SupportsDimensionOverride: m.Parameters.EmbeddingParameters.SupportsDimensionOverride,
TruncatePromptTokens: m.Parameters.EmbeddingParameters.TruncatePromptTokens,
Provider: m.Parameters.Provider,
MaxConcurrency: m.Parameters.MaxConcurrency,
ExtraConfig: m.Parameters.ExtraConfig,
CustomHeaders: m.Parameters.CustomHeaders,
AppID: appID,
AppSecret: appSecret,
}
}
// NewEmbedder creates an embedder based on the configuration
func NewEmbedder(config Config, pooler EmbedderPooler, ollamaService *ollama.OllamaService) (Embedder, error) {
e, err := newEmbedder(config, pooler, ollamaService)
if err != nil {
return e, err
}
if setter, ok := e.(interface{ SetSupportsDimensionOverride(bool) }); ok {
setter.SetSupportsDimensionOverride(config.SupportsDimensionOverride)
}
// Innermost: gate the real provider round-trips (including the per-sub-batch
// pool callbacks) before debug/langfuse wrap for logging/tracing. See
// concurrencyEmbedder for why this sits below the observability decorators.
e = wrapEmbeddingConcurrency(wrapEmbeddingRequestBudget(e), config.MaxConcurrency)
if logger.LLMDebugEnabled() {
e = &debugEmbedder{inner: e}
}
if langfuse.GetManager().Enabled() {
e = &langfuseEmbedder{inner: e}
}
return e, nil
}
func newEmbedder(config Config, pooler EmbedderPooler, ollamaService *ollama.OllamaService) (Embedder, error) {
var embedder Embedder
var err error
switch strings.ToLower(string(config.Source)) {
case string(types.ModelSourceLocal):
embedder, err = NewOllamaEmbedder(config.BaseURL,
config.ModelName, config.TruncatePromptTokens, config.Dimensions, config.ModelID, pooler, ollamaService)
return embedder, err
case string(types.ModelSourceRemote):
// Detect or use configured provider for routing
providerName := provider.ProviderName(config.Provider)
if providerName == "" {
providerName = provider.DetectProvider(config.BaseURL)
}
// Route to provider-specific embedders
switch providerName {
case provider.ProviderAliyun:
// 检查是否是多模态嵌入模型
// 多模态模型: tongyi-embedding-vision-*, multimodal-embedding-*
// tex-only模型: text-embedding-v1/v2/v3/v4 应该使用 OpenAI 兼容接口,否则响应格式不匹配、embedding 返回空数组
isMultimodalModel := strings.Contains(strings.ToLower(config.ModelName), "vision") ||
strings.Contains(strings.ToLower(config.ModelName), "multimodal")
if isMultimodalModel {
// 多模态模型需要使用DashScope专用 API 端点
// 如果用户填写了 OpenAI 兼容模式的 URL,自动修正为多模态 API 的baseURL
baseURL := config.BaseURL
if baseURL == "" {
baseURL = "https://dashscope.aliyuncs.com"
} else if strings.Contains(baseURL, "/compatible-mode/") {
// 移除 compatible-mode 路径,AliyunEmbedder 会自动添加多模态端点
baseURL = strings.Replace(baseURL, "/compatible-mode/v1", "", 1)
baseURL = strings.Replace(baseURL, "/compatible-mode", "", 1)
}
aliyunEmb, aErr := NewAliyunEmbedder(config.APIKey,
baseURL,
config.ModelName,
config.TruncatePromptTokens,
config.Dimensions,
config.ModelID,
pooler)
if aliyunEmb != nil {
aliyunEmb.SetCustomHeaders(config.CustomHeaders)
}
embedder, err = aliyunEmb, aErr
} else {
baseURL := config.BaseURL
if baseURL == "" || !strings.Contains(baseURL, "/compatible-mode/") {
baseURL = "https://dashscope.aliyuncs.com/compatible-mode/v1"
}
openaiEmb, oErr := NewOpenAIEmbedder(config.APIKey,
baseURL,
config.ModelName,
config.TruncatePromptTokens,
config.Dimensions,
config.ModelID,
pooler)
if openaiEmb != nil {
openaiEmb.SetCustomHeaders(config.CustomHeaders)
}
embedder, err = openaiEmb, oErr
}
return embedder, err
case provider.ProviderVolcengine:
// Volcengine Ark uses multimodal embedding API
volcEmb, vErr := NewVolcengineEmbedder(config.APIKey,
config.BaseURL,
config.ModelName,
config.TruncatePromptTokens,
config.Dimensions,
config.ModelID,
pooler)
if volcEmb != nil {
volcEmb.SetCustomHeaders(config.CustomHeaders)
}
embedder, err = volcEmb, vErr
return embedder, err
case provider.ProviderJina:
// Jina AI uses different API format (truncate instead of truncate_prompt_tokens)
jinaEmb, jErr := NewJinaEmbedder(config.APIKey,
config.BaseURL,
config.ModelName,
config.TruncatePromptTokens,
config.Dimensions,
config.ModelID,
pooler)
if jinaEmb != nil {
jinaEmb.SetCustomHeaders(config.CustomHeaders)
}
embedder, err = jinaEmb, jErr
return embedder, err
case provider.ProviderAzureOpenAI:
apiVersion := "2024-10-21"
if config.ExtraConfig != nil {
if v, ok := config.ExtraConfig["api_version"]; ok {
apiVersion = v
}
}
azureEmb, azErr := NewAzureOpenAIEmbedder(config.APIKey,
config.BaseURL,
config.ModelName,
config.TruncatePromptTokens,
config.Dimensions,
config.ModelID,
apiVersion,
pooler)
if azureEmb != nil {
azureEmb.SetCustomHeaders(config.CustomHeaders)
}
embedder, err = azureEmb, azErr
return embedder, err
case provider.ProviderNvidia:
nvEmb, nErr := NewNvidiaEmbedder(config.APIKey,
config.BaseURL,
config.ModelName,
config.Dimensions,
config.ModelID,
pooler)
if nvEmb != nil {
nvEmb.SetCustomHeaders(config.CustomHeaders)
}
embedder, err = nvEmb, nErr
return embedder, err
case provider.ProviderGemini:
geminiEmb, gErr := NewGeminiEmbedder(config.APIKey,
config.BaseURL,
config.ModelName,
config.TruncatePromptTokens,
config.Dimensions,
config.ModelID,
pooler)
if geminiEmb != nil {
geminiEmb.SetCustomHeaders(config.CustomHeaders)
}
embedder, err = geminiEmb, gErr
return embedder, err
case provider.ProviderZhipu:
zhipuEmb, zErr := NewZhipuEmbedder(config.APIKey,
config.BaseURL,
config.ModelName,
config.TruncatePromptTokens,
config.Dimensions,
config.ModelID,
pooler)
if zhipuEmb != nil {
zhipuEmb.SetCustomHeaders(config.CustomHeaders)
}
embedder, err = zhipuEmb, zErr
return embedder, err
case provider.ProviderWeKnoraCloud:
embedder, err = NewWeKnoraCloudEmbedder(config)
return embedder, err
default:
// Use OpenAI-compatible embedder for other providers
openaiEmb, oErr := NewOpenAIEmbedder(config.APIKey,
config.BaseURL,
config.ModelName,
config.TruncatePromptTokens,
config.Dimensions,
config.ModelID,
pooler)
if openaiEmb != nil {
openaiEmb.SetCustomHeaders(config.CustomHeaders)
}
embedder, err = openaiEmb, oErr
return embedder, err
}
default:
return nil, fmt.Errorf("unsupported embedder source: %s", config.Source)
}
}
+248
View File
@@ -0,0 +1,248 @@
package embedding
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
"time"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/models/requestbudget"
secutils "github.com/Tencent/WeKnora/internal/utils"
)
const geminiEmbeddingBaseURL = "https://generativelanguage.googleapis.com/v1beta"
// GeminiEmbedder implements text vectorization using the native Gemini
// embedContent / batchEmbedContents REST API.
type GeminiEmbedder struct {
apiKey string
baseURL string
modelName string
truncatePromptTokens int
dimensions int
modelID string
httpClient *http.Client
timeout time.Duration
maxRetries int
customHeaders map[string]string
supportsDimensionOverride bool
EmbedderPooler
}
type geminiBatchEmbedRequest struct {
Requests []geminiEmbedRequest `json:"requests"`
}
type geminiEmbedRequest struct {
Model string `json:"model"`
Content geminiContent `json:"content"`
TaskType string `json:"taskType,omitempty"`
OutputDimensionality int `json:"output_dimensionality,omitempty"`
}
type geminiContent struct {
Parts []geminiPart `json:"parts"`
}
type geminiPart struct {
Text string `json:"text"`
}
type geminiBatchEmbedResponse struct {
Embeddings []geminiEmbedding `json:"embeddings"`
}
type geminiEmbedding struct {
Values []float32 `json:"values"`
}
func NewGeminiEmbedder(apiKey, baseURL, modelName string,
truncatePromptTokens int, dimensions int, modelID string, pooler EmbedderPooler,
) (*GeminiEmbedder, error) {
if modelName == "" {
return nil, fmt.Errorf("model name is required")
}
if truncatePromptTokens == 0 {
truncatePromptTokens = 511
}
if baseURL == "" {
baseURL = geminiEmbeddingBaseURL
}
baseURL = strings.TrimRight(baseURL, "/")
if strings.HasSuffix(baseURL, "/openai") {
baseURL = strings.TrimSuffix(baseURL, "/openai")
}
timeout := 60 * time.Second
if err := validateEmbeddingBaseURL(baseURL); err != nil {
return nil, err
}
return &GeminiEmbedder{
apiKey: apiKey,
baseURL: baseURL,
modelName: strings.TrimPrefix(modelName, "models/"),
truncatePromptTokens: truncatePromptTokens,
dimensions: dimensions,
modelID: modelID,
httpClient: newEmbeddingHTTPClient(timeout),
timeout: timeout,
maxRetries: 3,
EmbedderPooler: pooler,
}, nil
}
func (e *GeminiEmbedder) SetCustomHeaders(headers map[string]string) {
e.customHeaders = headers
}
func (e *GeminiEmbedder) SetSupportsDimensionOverride(supported bool) {
e.supportsDimensionOverride = supported
}
func (e *GeminiEmbedder) Embed(ctx context.Context, text string) ([]float32, error) {
embeddings, err := e.BatchEmbed(ctx, []string{text})
if err != nil {
return nil, err
}
if len(embeddings) == 0 {
return nil, fmt.Errorf("no embedding returned")
}
return embeddings[0], nil
}
func (e *GeminiEmbedder) BatchEmbed(ctx context.Context, texts []string) ([][]float32, error) {
if len(texts) == 0 {
return nil, nil
}
requests := make([]geminiEmbedRequest, 0, len(texts))
for _, text := range texts {
req := geminiEmbedRequest{
Model: "models/" + e.modelName,
Content: geminiContent{Parts: []geminiPart{
{Text: text},
}},
}
if e.supportsDimensionOverride && e.dimensions > 0 {
req.OutputDimensionality = e.dimensions
}
requests = append(requests, req)
}
jsonData, err := json.Marshal(geminiBatchEmbedRequest{Requests: requests})
if err != nil {
logger.GetLogger(ctx).Errorf("GeminiEmbedder BatchEmbed marshal request error: %v", err)
return nil, fmt.Errorf("marshal request: %w", err)
}
logger.GetLogger(ctx).Debugf("GeminiEmbedder BatchEmbed: model=%s, input_count=%d",
e.modelName, len(texts))
resp, err := e.doRequestWithRetry(ctx, jsonData)
if err != nil {
logger.GetLogger(ctx).Errorf("GeminiEmbedder BatchEmbed send request error: %v", err)
return nil, fmt.Errorf("send request: %w", err)
}
if resp.Body != nil {
defer resp.Body.Close()
}
body, err := io.ReadAll(resp.Body)
if err != nil {
logger.GetLogger(ctx).Errorf("GeminiEmbedder BatchEmbed read response error: %v", err)
return nil, fmt.Errorf("read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
bodyStr := string(body)
if len(bodyStr) > 1000 {
bodyStr = bodyStr[:1000] + "... (truncated)"
}
logger.GetLogger(ctx).Errorf("GeminiEmbedder BatchEmbed API error: Http Status %s, Response Body: %s", resp.Status, bodyStr)
return nil, fmt.Errorf("Gemini BatchEmbed API error: Http Status %s, Response: %s", resp.Status, bodyStr)
}
var response geminiBatchEmbedResponse
if err := json.Unmarshal(body, &response); err != nil {
logger.GetLogger(ctx).Errorf("GeminiEmbedder BatchEmbed unmarshal response error: %v", err)
return nil, fmt.Errorf("unmarshal response: %w", err)
}
if len(response.Embeddings) != len(texts) {
return nil, fmt.Errorf("Gemini BatchEmbed returned %d embeddings for %d inputs", len(response.Embeddings), len(texts))
}
embeddings := make([][]float32, 0, len(response.Embeddings))
for _, embedding := range response.Embeddings {
embeddings = append(embeddings, embedding.Values)
}
return embeddings, nil
}
func (e *GeminiEmbedder) doRequestWithRetry(ctx context.Context, jsonData []byte) (*http.Response, error) {
ctx = requestbudget.WithModel(ctx, e.modelID)
var resp *http.Response
var err error
url := fmt.Sprintf("%s/models/%s:batchEmbedContents", e.baseURL, e.modelName)
for i := 0; i <= e.maxRetries; i++ {
if err := requestbudget.Check(ctx); err != nil {
return nil, err
}
if i > 0 {
backoffTime := time.Duration(1<<uint(i-1)) * time.Second
if backoffTime > 10*time.Second {
backoffTime = 10 * time.Second
}
logger.GetLogger(ctx).
Infof("GeminiEmbedder retrying request (%d/%d), waiting %v", i, e.maxRetries, backoffTime)
select {
case <-time.After(backoffTime):
case <-ctx.Done():
return nil, ctx.Err()
}
}
var req *http.Request
req, err = http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(jsonData))
if err != nil {
logger.GetLogger(ctx).Errorf("GeminiEmbedder failed to create request: %v", err)
continue
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("x-goog-api-key", e.apiKey)
secutils.ApplyCustomHeaders(req, e.customHeaders)
resp, err = e.httpClient.Do(req)
if stopProviderRetry(ctx, err) {
return nil, err
}
if err == nil {
return resp, nil
}
logger.GetLogger(ctx).Errorf("GeminiEmbedder request failed (attempt %d/%d): %v", i+1, e.maxRetries+1, err)
}
return nil, err
}
func (e *GeminiEmbedder) GetModelName() string {
return e.modelName
}
func (e *GeminiEmbedder) GetDimensions() int {
return e.dimensions
}
func (e *GeminiEmbedder) GetModelID() string {
return e.modelID
}
+116
View File
@@ -0,0 +1,116 @@
package embedding
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"strings"
"testing"
)
func TestGeminiEmbedderBatchEmbedUsesNativeAPI(t *testing.T) {
t.Setenv("SSRF_WHITELIST", "127.0.0.1")
var gotPath string
var gotAPIKey string
var gotReq geminiBatchEmbedRequest
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
gotPath = r.URL.Path
gotAPIKey = r.Header.Get("x-goog-api-key")
if err := json.NewDecoder(r.Body).Decode(&gotReq); err != nil {
t.Fatalf("decode request: %v", err)
}
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{
"embeddings": [
{"values": [0.1, 0.2]},
{"values": [0.3, 0.4]}
]
}`))
}))
defer server.Close()
embedder, err := NewGeminiEmbedder("test-key", server.URL+"/openai", "gemini-embedding-2",
0, 768, "model-id", nil)
if err != nil {
t.Fatalf("NewGeminiEmbedder: %v", err)
}
embeddings, err := embedder.BatchEmbed(context.Background(), []string{"hello", "world"})
if err != nil {
t.Fatalf("BatchEmbed: %v", err)
}
if gotPath != "/models/gemini-embedding-2:batchEmbedContents" {
t.Fatalf("path = %q, want native batchEmbedContents path", gotPath)
}
if gotAPIKey != "test-key" {
t.Fatalf("x-goog-api-key = %q", gotAPIKey)
}
if len(gotReq.Requests) != 2 {
t.Fatalf("requests len = %d", len(gotReq.Requests))
}
if gotReq.Requests[0].Model != "models/gemini-embedding-2" {
t.Fatalf("request model = %q", gotReq.Requests[0].Model)
}
if gotReq.Requests[0].OutputDimensionality != 0 {
t.Fatalf("output_dimensionality = %d, want omitted by default", gotReq.Requests[0].OutputDimensionality)
}
if gotReq.Requests[0].Content.Parts[0].Text != "hello" {
t.Fatalf("first text = %q", gotReq.Requests[0].Content.Parts[0].Text)
}
if len(embeddings) != 2 || len(embeddings[0]) != 2 || embeddings[1][1] != 0.4 {
t.Fatalf("unexpected embeddings: %#v", embeddings)
}
}
func TestGeminiEmbedderBatchEmbedSendsOutputDimensionalityWhenOverrideEnabled(t *testing.T) {
t.Setenv("SSRF_WHITELIST", "127.0.0.1")
var gotReq geminiBatchEmbedRequest
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if err := json.NewDecoder(r.Body).Decode(&gotReq); err != nil {
t.Fatalf("decode request: %v", err)
}
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"embeddings":[{"values":[0.1,0.2]}]}`))
}))
defer server.Close()
embedder, err := NewGeminiEmbedder("test-key", server.URL, "gemini-embedding-2",
0, 768, "model-id", nil)
if err != nil {
t.Fatalf("NewGeminiEmbedder: %v", err)
}
embedder.SetSupportsDimensionOverride(true)
if _, err := embedder.BatchEmbed(context.Background(), []string{"hello"}); err != nil {
t.Fatalf("BatchEmbed: %v", err)
}
if gotReq.Requests[0].OutputDimensionality != 768 {
t.Fatalf("output_dimensionality = %d, want 768", gotReq.Requests[0].OutputDimensionality)
}
}
func TestGeminiEmbedderReturnsAPIErrorBody(t *testing.T) {
t.Setenv("SSRF_WHITELIST", "127.0.0.1")
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
http.Error(w, `{"error":"not found"}`, http.StatusNotFound)
}))
defer server.Close()
embedder, err := NewGeminiEmbedder("test-key", server.URL, "gemini-embedding-2",
0, 0, "model-id", nil)
if err != nil {
t.Fatalf("NewGeminiEmbedder: %v", err)
}
_, err = embedder.BatchEmbed(context.Background(), []string{"hello"})
if err == nil || !strings.Contains(err.Error(), "404") {
t.Fatalf("expected 404 error, got %v", err)
}
}
+222
View File
@@ -0,0 +1,222 @@
package embedding
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"time"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/models/requestbudget"
secutils "github.com/Tencent/WeKnora/internal/utils"
)
// JinaEmbedder implements text vectorization functionality using Jina AI API
// Jina API is mostly OpenAI-compatible but does NOT support truncate_prompt_tokens
type JinaEmbedder struct {
apiKey string
baseURL string
modelName string
dimensions int
modelID string
httpClient *http.Client
timeout time.Duration
maxRetries int
customHeaders map[string]string
supportsDimensionOverride bool
EmbedderPooler
}
// SetCustomHeaders 设置用户自定义 HTTP 请求头(类似 OpenAI Python SDK 的 extra_headers)。
func (e *JinaEmbedder) SetCustomHeaders(headers map[string]string) {
e.customHeaders = headers
}
func (e *JinaEmbedder) SetSupportsDimensionOverride(supported bool) {
e.supportsDimensionOverride = supported
}
// JinaEmbedRequest represents a Jina embedding request
// Note: Jina uses 'truncate' (boolean) instead of 'truncate_prompt_tokens' (integer)
type JinaEmbedRequest struct {
Model string `json:"model"`
Input []string `json:"input"`
Truncate bool `json:"truncate,omitempty"` // Whether to truncate text exceeding max token length
Dimensions int `json:"dimensions,omitempty"` // Output embedding dimensions (for models that support it)
}
// JinaEmbedResponse represents a Jina embedding response
type JinaEmbedResponse struct {
Data []struct {
Embedding []float32 `json:"embedding"`
Index int `json:"index"`
} `json:"data"`
}
// NewJinaEmbedder creates a new Jina embedder
func NewJinaEmbedder(apiKey, baseURL, modelName string,
truncatePromptTokens int, dimensions int, modelID string, pooler EmbedderPooler,
) (*JinaEmbedder, error) {
if baseURL == "" {
baseURL = "https://api.jina.ai/v1"
}
if modelName == "" {
return nil, fmt.Errorf("model name is required")
}
timeout := 60 * time.Second
if err := validateEmbeddingBaseURL(baseURL); err != nil {
return nil, err
}
return &JinaEmbedder{
apiKey: apiKey,
baseURL: baseURL,
modelName: modelName,
httpClient: newEmbeddingHTTPClient(timeout),
EmbedderPooler: pooler,
dimensions: dimensions,
modelID: modelID,
timeout: timeout,
maxRetries: 3,
}, nil
}
// Embed converts text to vector
func (e *JinaEmbedder) Embed(ctx context.Context, text string) ([]float32, error) {
for range 3 {
embeddings, err := e.BatchEmbed(ctx, []string{text})
if err != nil {
return nil, err
}
if len(embeddings) > 0 {
return embeddings[0], nil
}
}
return nil, fmt.Errorf("no embedding returned")
}
func (e *JinaEmbedder) doRequestWithRetry(ctx context.Context, jsonData []byte) (*http.Response, error) {
ctx = requestbudget.WithModel(ctx, e.modelID)
var resp *http.Response
var err error
url := e.baseURL + "/embeddings"
for i := 0; i <= e.maxRetries; i++ {
if err := requestbudget.Check(ctx); err != nil {
return nil, err
}
if i > 0 {
backoffTime := time.Duration(1<<uint(i-1)) * time.Second
if backoffTime > 10*time.Second {
backoffTime = 10 * time.Second
}
logger.GetLogger(ctx).
Infof("JinaEmbedder retrying request (%d/%d), waiting %v", i, e.maxRetries, backoffTime)
select {
case <-time.After(backoffTime):
case <-ctx.Done():
return nil, ctx.Err()
}
}
// Rebuild request each time to ensure Body is valid
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(jsonData))
if err != nil {
logger.GetLogger(ctx).Errorf("JinaEmbedder failed to create request: %v", err)
continue
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+e.apiKey)
secutils.ApplyCustomHeaders(req, e.customHeaders)
resp, err = e.httpClient.Do(req)
if stopProviderRetry(ctx, err) {
return nil, err
}
if err == nil {
return resp, nil
}
logger.GetLogger(ctx).Errorf("JinaEmbedder request failed (attempt %d/%d): %v", i+1, e.maxRetries+1, err)
}
return nil, err
}
func (e *JinaEmbedder) BatchEmbed(ctx context.Context, texts []string) ([][]float32, error) {
// Create request body - Jina uses 'truncate' boolean instead of 'truncate_prompt_tokens'
reqBody := JinaEmbedRequest{
Model: e.modelName,
Input: texts,
Truncate: true, // Enable truncation for long texts
}
if e.supportsDimensionOverride && e.dimensions > 0 {
reqBody.Dimensions = e.dimensions
}
jsonData, err := json.Marshal(reqBody)
if err != nil {
logger.GetLogger(ctx).Errorf("JinaEmbedder EmbedBatch marshal request error: %v", err)
return nil, fmt.Errorf("marshal request: %w", err)
}
// Send request
resp, err := e.doRequestWithRetry(ctx, jsonData)
if err != nil {
logger.GetLogger(ctx).Errorf("JinaEmbedder EmbedBatch send request error: %v", err)
return nil, fmt.Errorf("send request: %w", err)
}
if resp.Body != nil {
defer resp.Body.Close()
}
// Read response
body, err := io.ReadAll(resp.Body)
if err != nil {
logger.GetLogger(ctx).Errorf("JinaEmbedder EmbedBatch read response error: %v", err)
return nil, fmt.Errorf("read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
logger.GetLogger(ctx).Errorf("JinaEmbedder EmbedBatch API error: Http Status %s, Body: %s", resp.Status, string(body))
return nil, fmt.Errorf("EmbedBatch API error: Http Status %s", resp.Status)
}
// Parse response
var response JinaEmbedResponse
if err := json.Unmarshal(body, &response); err != nil {
logger.GetLogger(ctx).Errorf("JinaEmbedder EmbedBatch unmarshal response error: %v", err)
return nil, fmt.Errorf("unmarshal response: %w", err)
}
// Extract embedding vectors
embeddings := make([][]float32, 0, len(response.Data))
for _, data := range response.Data {
embeddings = append(embeddings, data.Embedding)
}
return embeddings, nil
}
// GetModelName returns the model name
func (e *JinaEmbedder) GetModelName() string {
return e.modelName
}
// GetDimensions returns the vector dimensions
func (e *JinaEmbedder) GetDimensions() int {
return e.dimensions
}
// GetModelID returns the model ID
func (e *JinaEmbedder) GetModelID() string {
return e.modelID
}
@@ -0,0 +1,136 @@
package embedding
import (
"context"
"github.com/Tencent/WeKnora/internal/tracing/langfuse"
)
// langfuseEmbedder wraps an Embedder and reports each call as a Langfuse
// generation observation. Input token counts are approximated from the text
// lengths when the underlying provider doesn't return usage data, because
// Langfuse's cost reports require non-zero input tokens.
type langfuseEmbedder struct {
inner Embedder
}
func (l *langfuseEmbedder) Embed(ctx context.Context, text string) ([]float32, error) {
mgr := langfuse.GetManager()
if !mgr.Enabled() {
return l.inner.Embed(ctx, text)
}
genCtx, gen := mgr.StartGeneration(ctx, langfuse.GenerationOptions{
Name: "embedding.embed",
Model: l.inner.GetModelName(),
Input: text,
Metadata: map[string]interface{}{
"model_id": l.inner.GetModelID(),
"dimensions": l.inner.GetDimensions(),
},
})
result, err := l.inner.Embed(genCtx, text)
usage := approxEmbeddingUsage([]string{text})
var out interface{}
if len(result) > 0 {
out = map[string]interface{}{
"dimensions": len(result),
"vector_preview": result[:min(3, len(result))],
}
}
gen.Finish(out, usage, err)
return result, err
}
func (l *langfuseEmbedder) BatchEmbed(ctx context.Context, texts []string) ([][]float32, error) {
mgr := langfuse.GetManager()
if !mgr.Enabled() {
return l.inner.BatchEmbed(ctx, texts)
}
genCtx, gen := mgr.StartGeneration(ctx, langfuse.GenerationOptions{
Name: "embedding.batch_embed",
Model: l.inner.GetModelName(),
Input: map[string]interface{}{
"count": len(texts),
// Avoid sending megabytes of full text — Langfuse truncates but
// the network cost is still real. Keep a short preview instead.
"preview": previewTexts(texts, 5),
},
Metadata: map[string]interface{}{
"model_id": l.inner.GetModelID(),
"dimensions": l.inner.GetDimensions(),
"batch_size": len(texts),
},
})
result, err := l.inner.BatchEmbed(genCtx, texts)
usage := approxEmbeddingUsage(texts)
var out interface{}
if len(result) > 0 {
out = map[string]interface{}{
"count": len(result),
"dimensions": len(result[0]),
}
}
gen.Finish(out, usage, err)
return result, err
}
func (l *langfuseEmbedder) BatchEmbedWithPool(ctx context.Context, model Embedder, texts []string) ([][]float32, error) {
return l.inner.BatchEmbedWithPool(ctx, l, texts)
}
func (l *langfuseEmbedder) GetModelName() string { return l.inner.GetModelName() }
func (l *langfuseEmbedder) GetDimensions() int { return l.inner.GetDimensions() }
func (l *langfuseEmbedder) GetModelID() string { return l.inner.GetModelID() }
// approxEmbeddingUsage estimates input tokens as ~rune_count / 4, matching the
// rule of thumb OpenAI uses in their tokenizer docs. This is purely for cost
// reporting — Langfuse lets users define per-model cost multipliers, so the
// approximation need only be proportional to length.
func approxEmbeddingUsage(texts []string) *langfuse.TokenUsage {
total := 0
for _, t := range texts {
runes := len([]rune(t))
if runes == 0 {
continue
}
total += runes/4 + 1
}
if total == 0 {
return nil
}
return &langfuse.TokenUsage{
Input: total,
Total: total,
Unit: "TOKENS",
}
}
func previewTexts(texts []string, n int) []string {
if len(texts) <= n {
out := make([]string, len(texts))
for i, t := range texts {
out[i] = truncateRunes(t, 120)
}
return out
}
out := make([]string, n)
for i := 0; i < n; i++ {
out[i] = truncateRunes(texts[i], 120)
}
return out
}
func truncateRunes(s string, maxRunes int) string {
r := []rune(s)
if len(r) <= maxRunes {
return s
}
return string(r[:maxRunes]) + "..."
}
func min(a, b int) int {
if a < b {
return a
}
return b
}
+93
View File
@@ -0,0 +1,93 @@
package embedding
import (
"context"
"fmt"
"strings"
"time"
"github.com/Tencent/WeKnora/internal/logger"
)
// debugEmbedder wraps an Embedder with LLM debug logging.
type debugEmbedder struct {
inner Embedder
}
func (d *debugEmbedder) Embed(ctx context.Context, text string) ([]float32, error) {
start := time.Now()
result, err := d.inner.Embed(ctx, text)
logEmbeddingDebug(ctx, d.inner.GetModelName(), []string{text}, singleToDouble(result), err, time.Since(start))
return result, err
}
func (d *debugEmbedder) BatchEmbed(ctx context.Context, texts []string) ([][]float32, error) {
start := time.Now()
result, err := d.inner.BatchEmbed(ctx, texts)
logEmbeddingDebug(ctx, d.inner.GetModelName(), texts, result, err, time.Since(start))
return result, err
}
func (d *debugEmbedder) BatchEmbedWithPool(ctx context.Context, model Embedder, texts []string) ([][]float32, error) {
return d.inner.BatchEmbedWithPool(ctx, d, texts)
}
func (d *debugEmbedder) GetModelName() string { return d.inner.GetModelName() }
func (d *debugEmbedder) GetDimensions() int { return d.inner.GetDimensions() }
func (d *debugEmbedder) GetModelID() string { return d.inner.GetModelID() }
func singleToDouble(v []float32) [][]float32 {
if v == nil {
return nil
}
return [][]float32{v}
}
func logEmbeddingDebug(ctx context.Context, model string, inputs []string, outputs [][]float32, callErr error, dur time.Duration) {
if !logger.LLMDebugEnabled() {
return
}
record := &logger.LLMCallRecord{
CallType: "Embedding",
Model: model,
Duration: dur,
}
// Input section: show each text with a preview
var inputBuf strings.Builder
inputBuf.WriteString(fmt.Sprintf("count=%d\n", len(inputs)))
for i, t := range inputs {
preview := strings.ReplaceAll(t, "\n", "\\n")
preview = logger.TruncateRunes(preview, 200)
inputBuf.WriteString(fmt.Sprintf("[%d] (len=%d) %s\n", i, len([]rune(t)), preview))
}
record.Sections = append(record.Sections, logger.RecordSection{Title: "Input", Content: inputBuf.String()})
// Output section
if outputs != nil {
var outBuf strings.Builder
outBuf.WriteString(fmt.Sprintf("count=%d\n", len(outputs)))
for i, vec := range outputs {
if len(vec) > 0 {
outBuf.WriteString(fmt.Sprintf("[%d] dims=%d, first_3=[%.6f, %.6f, %.6f]\n", i, len(vec),
safeIdx(vec, 0), safeIdx(vec, 1), safeIdx(vec, 2)))
} else {
outBuf.WriteString(fmt.Sprintf("[%d] empty\n", i))
}
}
record.Sections = append(record.Sections, logger.RecordSection{Title: "Output", Content: outBuf.String()})
}
if callErr != nil {
record.Error = callErr.Error()
}
logger.LLMDebugLog(ctx, record)
}
func safeIdx(v []float32, i int) float32 {
if i < len(v) {
return v[i]
}
return 0
}
+231
View File
@@ -0,0 +1,231 @@
package embedding
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"time"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/models/requestbudget"
"github.com/Tencent/WeKnora/internal/types"
secutils "github.com/Tencent/WeKnora/internal/utils"
)
// NvidiaEmbedder implements text vectorization functionality using NVIDIA API
type NvidiaEmbedder struct {
apiKey string
baseURL string
modelName string
dimensions int
modelID string
httpClient *http.Client
timeout time.Duration
maxRetries int
customHeaders map[string]string
supportsDimensionOverride bool
EmbedderPooler
}
// SetCustomHeaders 设置用户自定义 HTTP 请求头(类似 OpenAI Python SDK 的 extra_headers)。
func (e *NvidiaEmbedder) SetCustomHeaders(headers map[string]string) {
e.customHeaders = headers
}
func (e *NvidiaEmbedder) SetSupportsDimensionOverride(supported bool) {
e.supportsDimensionOverride = supported
}
// NvidiaEmbedRequest represents an NVIDIA embedding request
type NvidiaEmbedRequest struct {
Model string `json:"model"`
Input []string `json:"input"`
EncodingFormat string `json:"encoding_format,omitempty"`
Dimensions int `json:"dimensions,omitempty"`
TruncatePromptTokens int `json:"truncate_prompt_tokens,omitempty"`
InputType string `json:"input_type,omitempty"`
}
// NvidiaEmbedResponse represents an NVIDIA embedding response
type NvidiaEmbedResponse struct {
Data []struct {
Embedding []float32 `json:"embedding"`
Index int `json:"index"`
} `json:"data"`
}
// NewNvidiaEmbedder creates a new NVIDIA embedder
func NewNvidiaEmbedder(apiKey, baseURL, modelName string,
dimensions int, modelID string, pooler EmbedderPooler,
) (*NvidiaEmbedder, error) {
if baseURL == "" {
baseURL = "https://integrate.api.nvidia.com/v1"
}
if modelName == "" {
return nil, fmt.Errorf("model name is required")
}
timeout := 60 * time.Second
if err := validateEmbeddingBaseURL(baseURL); err != nil {
return nil, err
}
return &NvidiaEmbedder{
apiKey: apiKey,
baseURL: baseURL,
modelName: modelName,
httpClient: newEmbeddingHTTPClient(timeout),
EmbedderPooler: pooler,
dimensions: dimensions,
modelID: modelID,
timeout: timeout,
maxRetries: 3, // Maximum retry count
}, nil
}
// Embed converts text to vector
func (e *NvidiaEmbedder) Embed(ctx context.Context, text string) ([]float32, error) {
for range 3 {
embeddings, err := e.BatchEmbed(ctx, []string{text})
if err != nil {
return nil, err
}
if len(embeddings) > 0 {
return embeddings[0], nil
}
}
return nil, fmt.Errorf("no embedding returned")
}
func (e *NvidiaEmbedder) doRequestWithRetry(ctx context.Context, jsonData []byte) (*http.Response, error) {
ctx = requestbudget.WithModel(ctx, e.modelID)
var resp *http.Response
var err error
url := e.baseURL + "/embeddings"
for i := 0; i <= e.maxRetries; i++ {
if err := requestbudget.Check(ctx); err != nil {
return nil, err
}
if i > 0 {
backoffTime := time.Duration(1<<uint(i-1)) * time.Second
if backoffTime > 10*time.Second {
backoffTime = 10 * time.Second
}
logger.GetLogger(ctx).
Infof("NvidiaEmbedder retrying request (%d/%d), waiting %v", i, e.maxRetries, backoffTime)
select {
case <-time.After(backoffTime):
case <-ctx.Done():
return nil, ctx.Err()
}
}
// Rebuild request each time to ensure Body is valid
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(jsonData))
if err != nil {
logger.GetLogger(ctx).Errorf("NvidiaEmbedder failed to create request: %v", err)
continue
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+e.apiKey)
secutils.ApplyCustomHeaders(req, e.customHeaders)
resp, err = e.httpClient.Do(req)
if stopProviderRetry(ctx, err) {
return nil, err
}
if err == nil {
return resp, nil
}
logger.GetLogger(ctx).Errorf("NvidiaEmbedder request failed (attempt %d/%d): %v", i+1, e.maxRetries+1, err)
}
return nil, err
}
func (e *NvidiaEmbedder) BatchEmbed(ctx context.Context, texts []string) ([][]float32, error) {
// Create request body
reqBody := NvidiaEmbedRequest{
Model: e.modelName,
Input: texts,
EncodingFormat: "float",
InputType: "passage",
}
if e.supportsDimensionsParam() {
reqBody.Dimensions = e.dimensions
}
isQuery, _ := ctx.Value(types.EmbedQueryContextKey).(bool)
if isQuery {
reqBody.InputType = "query"
}
jsonData, err := json.Marshal(reqBody)
if err != nil {
logger.GetLogger(ctx).Errorf("NvidiaEmbedder EmbedBatch marshal request error: %v", err)
return nil, fmt.Errorf("marshal request: %w", err)
}
// Send request (passing jsonData instead of constructing http.Request)
resp, err := e.doRequestWithRetry(ctx, jsonData)
if err != nil {
logger.GetLogger(ctx).Errorf("NvidiaEmbedder EmbedBatch send request error: %v", err)
return nil, fmt.Errorf("send request: %w", err)
}
if resp.Body != nil {
defer resp.Body.Close()
}
// Read response
body, err := io.ReadAll(resp.Body)
if err != nil {
logger.GetLogger(ctx).Errorf("NvidiaEmbedder EmbedBatch read response error: %v", err)
return nil, fmt.Errorf("read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
logger.GetLogger(ctx).Errorf("NvidiaEmbedder EmbedBatch API error: Http Status %s", resp.Status)
return nil, fmt.Errorf("EmbedBatch API error: Http Status %s", resp.Status)
}
// Parse response
var response NvidiaEmbedResponse
if err := json.Unmarshal(body, &response); err != nil {
logger.GetLogger(ctx).Errorf("NvidiaEmbedder EmbedBatch unmarshal response error: %v", err)
return nil, fmt.Errorf("unmarshal response: %w", err)
}
// Extract embedding vectors
embeddings := make([][]float32, 0, len(response.Data))
for _, data := range response.Data {
embeddings = append(embeddings, data.Embedding)
}
return embeddings, nil
}
// GetModelName returns the model name
func (e *NvidiaEmbedder) GetModelName() string {
return e.modelName
}
func (e *NvidiaEmbedder) supportsDimensionsParam() bool {
return e.supportsDimensionOverride && e.dimensions > 0
}
// GetDimensions returns the vector dimensions
func (e *NvidiaEmbedder) GetDimensions() int {
return e.dimensions
}
// GetModelID returns the model ID
func (e *NvidiaEmbedder) GetModelID() string {
return e.modelID
}
+139
View File
@@ -0,0 +1,139 @@
package embedding
import (
"context"
"fmt"
"time"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/models/requestbudget"
"github.com/Tencent/WeKnora/internal/models/utils/ollama"
ollamaapi "github.com/ollama/ollama/api"
)
// OllamaEmbedder implements text vectorization functionality using Ollama
type OllamaEmbedder struct {
modelName string
truncatePromptTokens int
ollamaService *ollama.OllamaService
dimensions int
modelID string
supportsDimensionOverride bool
EmbedderPooler
}
// OllamaEmbedRequest represents an Ollama embedding request
type OllamaEmbedRequest struct {
Model string `json:"model"`
Prompt string `json:"prompt"`
TruncatePromptTokens int `json:"truncate_prompt_tokens"`
}
// OllamaEmbedResponse represents an Ollama embedding response
type OllamaEmbedResponse struct {
Embedding []float32 `json:"embedding"`
}
// NewOllamaEmbedder creates a new Ollama embedder
func NewOllamaEmbedder(baseURL,
modelName string,
truncatePromptTokens int,
dimensions int,
modelID string,
pooler EmbedderPooler,
ollamaService *ollama.OllamaService,
) (*OllamaEmbedder, error) {
if modelName == "" {
modelName = "nomic-embed-text"
}
if truncatePromptTokens == 0 {
truncatePromptTokens = 511
}
return &OllamaEmbedder{
modelName: modelName,
truncatePromptTokens: truncatePromptTokens,
ollamaService: ollamaService,
EmbedderPooler: pooler,
dimensions: dimensions,
modelID: modelID,
}, nil
}
// ensureModelAvailable ensures that the model is available
func (e *OllamaEmbedder) ensureModelAvailable(ctx context.Context) error {
logger.GetLogger(ctx).Infof("Ensuring model %s is available", e.modelName)
return e.ollamaService.EnsureModelAvailable(ctx, e.modelName)
}
func (e *OllamaEmbedder) SetSupportsDimensionOverride(supported bool) {
e.supportsDimensionOverride = supported
}
// Embed converts text to vector
func (e *OllamaEmbedder) Embed(ctx context.Context, text string) ([]float32, error) {
embedding, err := e.BatchEmbed(ctx, []string{text})
if err != nil {
return nil, fmt.Errorf("failed to embed text: %w", err)
}
if len(embedding) == 0 {
return nil, fmt.Errorf("failed to embed text: %w", err)
}
return embedding[0], nil
}
// BatchEmbed converts multiple texts to vectors in batch
func (e *OllamaEmbedder) BatchEmbed(ctx context.Context, texts []string) ([][]float32, error) {
if err := requestbudget.RequireSupported(ctx, false); err != nil {
return nil, err
}
// Ensure model is available
if err := e.ensureModelAvailable(ctx); err != nil {
return nil, err
}
// Create request
req := &ollamaapi.EmbedRequest{
Model: e.modelName,
Input: texts,
Options: make(map[string]interface{}),
}
if e.supportsDimensionOverride && e.dimensions > 0 {
req.Dimensions = e.dimensions
}
// Set truncation parameters
if e.truncatePromptTokens > 0 {
req.Options["num_ctx"] = e.truncatePromptTokens
truncate := true
req.Truncate = &truncate
}
// Send request
startTime := time.Now()
resp, err := e.ollamaService.Embeddings(ctx, req)
if err != nil {
return nil, fmt.Errorf("failed to get embedding vectors: %w", err)
}
logger.GetLogger(ctx).Debugf("Embedding vector retrieval took: %v", time.Since(startTime))
return resp.Embeddings, nil
}
// GetModelName returns the model name
func (e *OllamaEmbedder) GetModelName() string {
return e.modelName
}
// GetDimensions returns the vector dimensions
func (e *OllamaEmbedder) GetDimensions() int {
return e.dimensions
}
// GetModelID returns the model ID
func (e *OllamaEmbedder) GetModelID() string {
return e.modelID
}
+277
View File
@@ -0,0 +1,277 @@
package embedding
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"time"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/models/requestbudget"
secutils "github.com/Tencent/WeKnora/internal/utils"
)
// OpenAIEmbedder implements text vectorization functionality using OpenAI API
type OpenAIEmbedder struct {
apiKey string
baseURL string
modelName string
truncatePromptTokens int
dimensions int
modelID string
httpClient *http.Client
timeout time.Duration
maxRetries int
customHeaders map[string]string
supportsDimensionOverride bool
EmbedderPooler
}
// OpenAIEmbedRequest represents an OpenAI embedding request
type OpenAIEmbedRequest struct {
Model string `json:"model"`
Input []string `json:"input"`
EncodingFormat string `json:"encoding_format,omitempty"`
Dimensions int `json:"dimensions,omitempty"`
TruncatePromptTokens int `json:"truncate_prompt_tokens,omitempty"`
}
// OpenAIEmbedResponse represents an OpenAI embedding response
type OpenAIEmbedResponse struct {
Data []struct {
Embedding []float32 `json:"embedding"`
Index int `json:"index"`
} `json:"data"`
}
// NewOpenAIEmbedder creates a new OpenAI embedder
func NewOpenAIEmbedder(apiKey, baseURL, modelName string,
truncatePromptTokens int, dimensions int, modelID string, pooler EmbedderPooler,
) (*OpenAIEmbedder, error) {
if baseURL == "" {
baseURL = "https://api.openai.com/v1"
}
if modelName == "" {
return nil, fmt.Errorf("model name is required")
}
if truncatePromptTokens == 0 {
truncatePromptTokens = 511
}
timeout := 60 * time.Second
if err := validateEmbeddingBaseURL(baseURL); err != nil {
return nil, err
}
return &OpenAIEmbedder{
apiKey: apiKey,
baseURL: baseURL,
modelName: modelName,
httpClient: newEmbeddingHTTPClient(timeout),
truncatePromptTokens: truncatePromptTokens,
EmbedderPooler: pooler,
dimensions: dimensions,
modelID: modelID,
timeout: timeout,
maxRetries: 3, // Maximum retry count
}, nil
}
// SetCustomHeaders 设置用户自定义 HTTP 请求头(类似 OpenAI Python SDK 的 extra_headers)。
// 保留头(Authorization、Content-Type 等)会在发送时被自动跳过。
func (e *OpenAIEmbedder) SetCustomHeaders(headers map[string]string) {
e.customHeaders = headers
}
func (e *OpenAIEmbedder) SetSupportsDimensionOverride(supported bool) {
e.supportsDimensionOverride = supported
}
// Embed converts text to vector
func (e *OpenAIEmbedder) Embed(ctx context.Context, text string) ([]float32, error) {
for range 3 {
embeddings, err := e.BatchEmbed(ctx, []string{text})
if err != nil {
return nil, err
}
if len(embeddings) > 0 {
return embeddings[0], nil
}
}
return nil, fmt.Errorf("no embedding returned")
}
func (e *OpenAIEmbedder) doRequestWithRetry(ctx context.Context, jsonData []byte) (*http.Response, error) {
ctx = requestbudget.WithModel(ctx, e.modelID)
var resp *http.Response
var err error
url := e.baseURL + "/embeddings"
for i := 0; i <= e.maxRetries; i++ {
if err := requestbudget.Check(ctx); err != nil {
return nil, err
}
if i > 0 {
backoffTime := time.Duration(1<<uint(i-1)) * time.Second
if backoffTime > 10*time.Second {
backoffTime = 10 * time.Second
}
logger.GetLogger(ctx).
Infof("OpenAIEmbedder retrying request (%d/%d), waiting %v", i, e.maxRetries, backoffTime)
select {
case <-time.After(backoffTime):
case <-ctx.Done():
return nil, ctx.Err()
}
}
// Rebuild request each time to ensure Body is valid.
// IMPORTANT: declare `req` separately (var) so the assignment to `err`
// below uses the outer-scope variable, not a fresh loop-local one.
// Previously this read `req, err := http.NewRequestWithContext(...)`,
// where `:=` introduced a new `err` shadowing the outer one. The
// `resp, err = httpClient.Do(req)` line then wrote to the shadowed
// `err` only, so when all retries failed with connection errors the
// outer `err` stayed nil. The function returned `(nil, nil)`, and
// callers (BatchEmbed line 195) blindly dereferenced `resp.Body` →
// SIGSEGV nil-pointer panic that took down the whole process.
// Reproduce: stop the embedding upstream (e.g. localhost:3130), make
// any RAG query → backend SIGSEGV instead of returning HTTP 500.
var req *http.Request
req, err = http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(jsonData))
if err != nil {
logger.GetLogger(ctx).Errorf("OpenAIEmbedder failed to create request: %v", err)
continue
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+e.apiKey)
secutils.ApplyCustomHeaders(req, e.customHeaders)
resp, err = e.httpClient.Do(req)
if stopProviderRetry(ctx, err) {
return nil, err
}
if err == nil {
return resp, nil
}
logger.GetLogger(ctx).Errorf("OpenAIEmbedder request failed (attempt %d/%d): %v", i+1, e.maxRetries+1, err)
}
return nil, err
}
func (e *OpenAIEmbedder) BatchEmbed(ctx context.Context, texts []string) ([][]float32, error) {
// Create request body
reqBody := OpenAIEmbedRequest{
Model: e.modelName,
Input: texts,
EncodingFormat: "float",
TruncatePromptTokens: e.truncatePromptTokens,
}
if e.supportsDimensionsParam() {
reqBody.Dimensions = e.dimensions
}
jsonData, err := json.Marshal(reqBody)
if err != nil {
logger.GetLogger(ctx).Errorf("OpenAIEmbedder EmbedBatch marshal request error: %v", err)
return nil, fmt.Errorf("marshal request: %w", err)
}
// Log request details for debugging
logger.GetLogger(ctx).Debugf("OpenAIEmbedder BatchEmbed: model=%s, input_count=%d, truncate_tokens=%d",
e.modelName, len(texts), e.truncatePromptTokens)
// Check for invalid input lengths and log details
hasInvalidLength := false
for i, text := range texts {
textLen := len(text)
textPreview := text
if len(textPreview) > 200 {
textPreview = textPreview[:200] + "..."
}
// Log warning if length is outside valid range [1, 8192]
if textLen == 0 || textLen > 8192 {
hasInvalidLength = true
logger.GetLogger(ctx).Errorf("OpenAIEmbedder BatchEmbed input[%d]: INVALID length=%d (must be [1, 8192]), preview=%s",
i, textLen, textPreview)
} else {
logger.GetLogger(ctx).Debugf("OpenAIEmbedder BatchEmbed input[%d]: length=%d, preview=%s",
i, textLen, textPreview)
}
}
if hasInvalidLength {
logger.GetLogger(ctx).Errorf("OpenAIEmbedder BatchEmbed: Found invalid input lengths, this will likely cause API error")
}
// Send request (passing jsonData instead of constructing http.Request)
resp, err := e.doRequestWithRetry(ctx, jsonData)
if err != nil {
logger.GetLogger(ctx).Errorf("OpenAIEmbedder EmbedBatch send request error: %v", err)
return nil, fmt.Errorf("send request: %w", err)
}
if resp.Body != nil {
defer resp.Body.Close()
}
// Read response
body, err := io.ReadAll(resp.Body)
if err != nil {
logger.GetLogger(ctx).Errorf("OpenAIEmbedder EmbedBatch read response error: %v", err)
return nil, fmt.Errorf("read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
// Log detailed error response from OpenAI API
bodyStr := string(body)
if len(bodyStr) > 1000 {
bodyStr = bodyStr[:1000] + "... (truncated)"
}
logger.GetLogger(ctx).Errorf("OpenAIEmbedder EmbedBatch API error: Http Status %s, Response Body: %s", resp.Status, bodyStr)
return nil, fmt.Errorf("EmbedBatch API error: Http Status %s, Response: %s", resp.Status, bodyStr)
}
// Parse response
var response OpenAIEmbedResponse
if err := json.Unmarshal(body, &response); err != nil {
logger.GetLogger(ctx).Errorf("OpenAIEmbedder EmbedBatch unmarshal response error: %v", err)
return nil, fmt.Errorf("unmarshal response: %w", err)
}
// Extract embedding vectors
embeddings := make([][]float32, 0, len(response.Data))
for _, data := range response.Data {
embeddings = append(embeddings, data.Embedding)
}
return embeddings, nil
}
// GetModelName returns the model name
func (e *OpenAIEmbedder) GetModelName() string {
return e.modelName
}
func (e *OpenAIEmbedder) supportsDimensionsParam() bool {
return e.supportsDimensionOverride && e.dimensions > 0
}
// GetDimensions returns the vector dimensions
func (e *OpenAIEmbedder) GetDimensions() int {
return e.dimensions
}
// GetModelID returns the model ID
func (e *OpenAIEmbedder) GetModelID() string {
return e.modelID
}
+83
View File
@@ -0,0 +1,83 @@
package embedding
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
)
func TestOpenAIEmbedderBatchEmbedOmitsDimensionsByDefault(t *testing.T) {
requestBody := captureOpenAIEmbeddingRequest(t, "text-embedding-3-small", 256, false)
if _, ok := requestBody["dimensions"]; ok {
t.Fatalf("expected request body to omit dimensions by default, got %v", requestBody)
}
}
func TestOpenAIEmbedderBatchEmbedSendsDimensionsWhenOverrideEnabled(t *testing.T) {
requestBody := captureOpenAIEmbeddingRequest(t, "text-embedding-3-small", 256, true)
got, ok := requestBody["dimensions"]
if !ok {
t.Fatalf("expected request body to include dimensions, got %v", requestBody)
}
if got != float64(256) {
t.Fatalf("unexpected dimensions value: got %v want 256", got)
}
}
func TestOpenAIEmbedderBatchEmbedOmitsDimensionsForOpenAICompatibleModels(t *testing.T) {
requestBody := captureOpenAIEmbeddingRequest(t, "text-embedding-v3", 1024, false)
if _, ok := requestBody["dimensions"]; ok {
t.Fatalf("expected request body to omit dimensions for OpenAI-compatible model, got %v", requestBody)
}
}
func TestOpenAIEmbedderBatchEmbedOmitsDimensionsForFixedSizeModels(t *testing.T) {
requestBody := captureOpenAIEmbeddingRequest(t, "text-embedding-ada-002", 1536, false)
if _, ok := requestBody["dimensions"]; ok {
t.Fatalf("expected request body to omit dimensions for fixed-size model, got %v", requestBody)
}
}
func captureOpenAIEmbeddingRequest(t *testing.T, modelName string, dimensions int, supportsDimensionOverride bool) map[string]any {
t.Helper()
t.Setenv("SSRF_WHITELIST", "127.0.0.1")
requestBody := map[string]any{}
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/embeddings" {
t.Fatalf("unexpected request path: %s", r.URL.Path)
}
if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
t.Fatalf("decode request body: %v", err)
}
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"data":[{"embedding":[0.1,0.2],"index":0}]}`))
}))
defer server.Close()
embedder, err := NewOpenAIEmbedder(
"test-key",
server.URL,
modelName,
511,
dimensions,
"8f7d6082-5a15-4f84-ae55-88b2bdac4ba0",
nil,
)
if err != nil {
t.Fatalf("NewOpenAIEmbedder: %v", err)
}
embedder.SetSupportsDimensionOverride(supportsDimensionOverride)
if _, err := embedder.BatchEmbed(context.Background(), []string{"hello"}); err != nil {
t.Fatalf("BatchEmbed: %v", err)
}
return requestBody
}
@@ -0,0 +1,49 @@
package embedding
import (
"context"
"errors"
"github.com/Tencent/WeKnora/internal/models/requestbudget"
)
func stopProviderRetry(ctx context.Context, err error) bool {
return err != nil && (ctx.Err() != nil || errors.Is(err, requestbudget.ErrStopped))
}
type requestBudgetEmbedder struct {
Embedder
supported bool
}
func wrapEmbeddingRequestBudget(inner Embedder) Embedder {
supported := false
switch inner.(type) {
case *OpenAIEmbedder, *AzureOpenAIEmbedder, *AliyunEmbedder, *VolcengineEmbedder,
*JinaEmbedder, *NvidiaEmbedder, *GeminiEmbedder, *ZhipuEmbedder, *WeKnoraCloudEmbedder:
supported = true
}
return &requestBudgetEmbedder{Embedder: inner, supported: supported}
}
func (w *requestBudgetEmbedder) Embed(ctx context.Context, text string) ([]float32, error) {
if err := requestbudget.RequireSupported(ctx, w.supported); err != nil {
return nil, err
}
return w.Embedder.Embed(ctx, text)
}
func (w *requestBudgetEmbedder) BatchEmbed(ctx context.Context, texts []string) ([][]float32, error) {
if err := requestbudget.RequireSupported(ctx, w.supported); err != nil {
return nil, err
}
return w.Embedder.BatchEmbed(ctx, texts)
}
func (w *requestBudgetEmbedder) BatchEmbedWithPool(ctx context.Context, model Embedder, texts []string) ([][]float32, error) {
if err := requestbudget.RequireSupported(ctx, w.supported); err != nil {
return nil, err
}
// Preserve the concurrency wrapper supplied by the factory for each batch.
return w.Embedder.BatchEmbedWithPool(ctx, model, texts)
}
@@ -0,0 +1,257 @@
package embedding
import (
"context"
"encoding/json"
"errors"
"fmt"
"net/http"
"net/http/httptest"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/Tencent/WeKnora/internal/models/requestbudget"
"github.com/Tencent/WeKnora/internal/types"
"github.com/panjf2000/ants/v2"
)
type embeddingBudgetGuard struct {
mu sync.Mutex
limit int
models, phases []string
}
func (g *embeddingBudgetGuard) Key() string { return "local-embedding-test" }
func (g *embeddingBudgetGuard) Check(ctx context.Context) error { return ctx.Err() }
func (g *embeddingBudgetGuard) Reserve(ctx context.Context, model, phase string) error {
g.mu.Lock()
defer g.mu.Unlock()
if err := ctx.Err(); err != nil {
return err
}
if len(g.models) >= g.limit {
return requestbudget.ErrStopped
}
g.models, g.phases = append(g.models, model), append(g.phases, phase)
return nil
}
func embeddingBudgetContext(g *embeddingBudgetGuard) context.Context {
return requestbudget.WithModel(requestbudget.WithPhase(requestbudget.WithGuard(context.Background(), g), "wiki_embedding"), "caller-forged-model")
}
func assertEmbeddingBudget(t *testing.T, g *embeddingBudgetGuard, sends int) {
t.Helper()
g.mu.Lock()
defer g.mu.Unlock()
if len(g.models) != sends {
t.Fatalf("sends=%d reservations=%d", sends, len(g.models))
}
for i := range g.models {
if g.models[i] != "concrete-embedding" || g.phases[i] != "wiki_embedding" {
t.Fatalf("lost adapter identity/context: %v %v", g.models, g.phases)
}
}
}
func TestRequestBudgetEmbeddingProviderTransports(t *testing.T) {
t.Setenv("SSRF_WHITELIST", "127.0.0.1")
var sends atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
sends.Add(1)
w.Header().Set("Content-Type", "application/json")
fmt.Fprint(w, `{"data":[{"index":0,"embedding":[0.1]}]}`)
}))
defer server.Close()
client := newEmbeddingHTTPClient(3 * time.Second)
calls := map[string]func(context.Context, []byte) (*http.Response, error){
"openai": (&OpenAIEmbedder{baseURL: server.URL, modelID: "concrete-embedding", httpClient: client}).doRequestWithRetry,
"azure": (&AzureOpenAIEmbedder{baseURL: server.URL, modelID: "concrete-embedding", httpClient: client}).doRequestWithRetry,
"aliyun": (&AliyunEmbedder{baseURL: server.URL, modelID: "concrete-embedding", httpClient: client}).doRequestWithRetry,
"volcengine": (&VolcengineEmbedder{baseURL: server.URL, modelID: "concrete-embedding", httpClient: client}).doRequestWithRetry,
"jina": (&JinaEmbedder{baseURL: server.URL, modelID: "concrete-embedding", httpClient: client}).doRequestWithRetry,
"nvidia": (&NvidiaEmbedder{baseURL: server.URL, modelID: "concrete-embedding", httpClient: client}).doRequestWithRetry,
"gemini": (&GeminiEmbedder{baseURL: server.URL, modelID: "concrete-embedding", httpClient: client}).doRequestWithRetry,
"zhipu": (&ZhipuEmbedder{baseURL: server.URL, modelID: "concrete-embedding", httpClient: client}).doRequestWithRetry,
}
for name, call := range calls {
for _, limit := range []int{0, 1} {
t.Run(fmt.Sprintf("%s/limit=%d", name, limit), func(t *testing.T) {
sends.Store(0)
g := &embeddingBudgetGuard{limit: limit}
resp, err := call(embeddingBudgetContext(g), []byte(`{}`))
if resp != nil {
resp.Body.Close()
}
if limit == 0 && !errors.Is(err, requestbudget.ErrStopped) {
t.Fatalf("guard not applied: %v", err)
}
if limit == 1 && err != nil {
t.Fatal(err)
}
if int(sends.Load()) != limit {
t.Fatalf("sends=%d want=%d", sends.Load(), limit)
}
assertEmbeddingBudget(t, g, limit)
})
}
}
for _, limit := range []int{0, 1} {
t.Run(fmt.Sprintf("weknoracloud/limit=%d", limit), func(t *testing.T) {
sends.Store(0)
e := &WeKnoraCloudEmbedder{baseURL: server.URL, modelName: "local-test", modelID: "concrete-embedding", client: client}
g := &embeddingBudgetGuard{limit: limit}
_, err := e.BatchEmbed(embeddingBudgetContext(g), []string{"local test"})
if limit == 0 && !errors.Is(err, requestbudget.ErrStopped) {
t.Fatalf("guard not applied: %v", err)
}
if limit == 1 && err != nil {
t.Fatal(err)
}
if int(sends.Load()) != limit {
t.Fatalf("sends=%d want=%d", sends.Load(), limit)
}
assertEmbeddingBudget(t, g, limit)
})
}
}
func TestRequestBudgetEmbeddingPoolCountsEachBatch(t *testing.T) {
t.Setenv("SSRF_WHITELIST", "127.0.0.1")
t.Setenv("BATCH_EMBED_SIZE", "2")
for _, limit := range []int{0, 2, 4} {
t.Run(fmt.Sprint(limit), func(t *testing.T) {
var sends atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
sends.Add(1)
var body struct {
Input []string `json:"input"`
}
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
http.Error(w, "bad request", http.StatusBadRequest)
return
}
data := make([]map[string]any, len(body.Input))
for i := range body.Input {
data[i] = map[string]any{"index": i, "embedding": []float32{0.1}}
}
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(map[string]any{"data": data})
}))
defer server.Close()
pool, err := ants.NewPool(3)
if err != nil {
t.Fatal(err)
}
defer pool.Release()
e, err := NewEmbedder(Config{Source: types.ModelSourceRemote, Provider: "openai", BaseURL: server.URL, APIKey: "local-only", ModelID: "concrete-embedding", ModelName: "local-test"}, NewBatchEmbedder(pool), nil)
if err != nil {
t.Fatal(err)
}
g := &embeddingBudgetGuard{limit: limit}
ctx, cancel := context.WithTimeout(embeddingBudgetContext(g), 3*time.Second)
defer cancel()
vectors, err := e.BatchEmbedWithPool(ctx, e, []string{"a", "b", "c", "d", "e", "f", "g"})
if limit < 4 && !errors.Is(err, requestbudget.ErrStopped) {
t.Fatalf("batch escaped budget: %v", err)
}
if limit == 4 && (err != nil || len(vectors) != 7) {
t.Fatalf("batch failed: %d %v", len(vectors), err)
}
if int(sends.Load()) != limit {
t.Fatalf("batch sends=%d want=%d", sends.Load(), limit)
}
assertEmbeddingBudget(t, g, limit)
})
}
}
func TestRequestBudgetEmbeddingRetriesConsumeReservations(t *testing.T) {
t.Setenv("SSRF_WHITELIST", "127.0.0.1")
for _, limit := range []int{1, 2} {
t.Run(fmt.Sprint(limit), func(t *testing.T) {
var sends atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if sends.Add(1) == 1 {
connection, _, err := w.(http.Hijacker).Hijack()
if err == nil {
connection.Close()
}
return
}
w.Header().Set("Content-Type", "application/json")
fmt.Fprint(w, `{"data":[{"index":0,"embedding":[0.1]}]}`)
}))
defer server.Close()
e, err := NewOpenAIEmbedder("local-only", server.URL, "local-test", 0, 0, "concrete-embedding", nil)
if err != nil {
t.Fatal(err)
}
e.maxRetries = 1
g := &embeddingBudgetGuard{limit: limit}
_, err = e.BatchEmbed(embeddingBudgetContext(g), []string{"test"})
if limit == 1 && !errors.Is(err, requestbudget.ErrStopped) {
t.Fatalf("retry escaped budget: %v", err)
}
if limit == 2 && err != nil {
t.Fatal(err)
}
if int(sends.Load()) != limit {
t.Fatalf("retry sends=%d want=%d", sends.Load(), limit)
}
assertEmbeddingBudget(t, g, limit)
})
}
}
func TestRequestBudgetVolcengineCountsEachText(t *testing.T) {
t.Setenv("SSRF_WHITELIST", "127.0.0.1")
var sends atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
sends.Add(1)
w.Header().Set("Content-Type", "application/json")
fmt.Fprint(w, `{"data":{"embedding":[0.1]}}`)
}))
defer server.Close()
e := &VolcengineEmbedder{baseURL: server.URL, modelID: "concrete-embedding", modelName: "local-test", httpClient: newEmbeddingHTTPClient(time.Second)}
g := &embeddingBudgetGuard{limit: 2}
_, err := e.BatchEmbed(embeddingBudgetContext(g), []string{"a", "b", "c"})
if !errors.Is(err, requestbudget.ErrStopped) || sends.Load() != 2 {
t.Fatalf("per-text budget not enforced: sends=%d err=%v", sends.Load(), err)
}
assertEmbeddingBudget(t, g, 2)
}
func TestRequestBudgetEmbeddingFactoryFailsClosed(t *testing.T) {
g := &embeddingBudgetGuard{limit: 1}
ctx := embeddingBudgetContext(g)
e, err := NewEmbedder(Config{Source: types.ModelSourceLocal, ModelID: "local-model"}, nil, nil)
if err != nil {
t.Fatal(err)
}
if _, err := e.Embed(ctx, "test"); !errors.Is(err, requestbudget.ErrStopped) {
t.Fatalf("local embed invoked: %v", err)
}
if _, err := e.BatchEmbed(ctx, []string{"test"}); !errors.Is(err, requestbudget.ErrStopped) {
t.Fatalf("local batch invoked: %v", err)
}
if _, err := e.BatchEmbedWithPool(ctx, e, []string{"test"}); !errors.Is(err, requestbudget.ErrStopped) {
t.Fatalf("local pool invoked: %v", err)
}
assertEmbeddingBudget(t, g, 0)
fake := newFakeEmbedder("unknown")
close(fake.release)
unknown := wrapEmbeddingRequestBudget(fake)
if _, err := unknown.Embed(ctx, "test"); !errors.Is(err, requestbudget.ErrStopped) {
t.Fatalf("unknown embed adapter allowed: %v", err)
}
if _, err := unknown.BatchEmbed(ctx, []string{"test"}); !errors.Is(err, requestbudget.ErrStopped) {
t.Fatalf("unknown batch adapter allowed: %v", err)
}
if _, err := unknown.BatchEmbedWithPool(ctx, unknown, []string{"test"}); !errors.Is(err, requestbudget.ErrStopped) {
t.Fatalf("unknown pool adapter allowed: %v", err)
}
if len(fake.enter) != 0 {
t.Fatal("unsupported adapter was invoked")
}
}
+40
View File
@@ -0,0 +1,40 @@
package embedding
import (
"fmt"
"net/http"
"time"
"github.com/Tencent/WeKnora/internal/models/requestbudget"
secutils "github.com/Tencent/WeKnora/internal/utils"
)
// sharedEmbeddingHTTPTransport keeps a single SSRF-safe connection pool for
// all embedding clients. Embedders are recreated as model configuration changes,
// but their outbound connections can be safely reused across client instances,
// so the transport (and its keep-alive pool) is built once at package load.
var sharedEmbeddingHTTPTransport = requestbudget.Transport{Base: secutils.NewSSRFSafeTransport(
secutils.DefaultSSRFSafeHTTPClientConfig(),
)}
// validateEmbeddingBaseURL checks that a resolved embedding API base URL is safe
// for outbound requests. Empty URLs are allowed (callers apply provider defaults).
func validateEmbeddingBaseURL(baseURL string) error {
if baseURL == "" {
return nil
}
if err := secutils.ValidateURLForSSRF(baseURL); err != nil {
return fmt.Errorf("base URL SSRF check failed: %w", err)
}
return nil
}
// newEmbeddingHTTPClient returns an HTTP client with connection-level SSRF
// protection and redirect validation, aligned with internal/models/chat/transport.go.
// All clients share sharedEmbeddingHTTPTransport so keep-alive connections are
// pooled globally, while each keeps its own timeout.
func newEmbeddingHTTPClient(timeout time.Duration) *http.Client {
cfg := secutils.DefaultSSRFSafeHTTPClientConfig()
cfg.Timeout = timeout
return secutils.NewSSRFSafeHTTPClientWithTransport(cfg, sharedEmbeddingHTTPTransport)
}
@@ -0,0 +1,72 @@
package embedding
import (
"net/http"
"strings"
"testing"
"time"
secutils "github.com/Tencent/WeKnora/internal/utils"
)
func TestNewEmbeddingHTTPClient_ReusesTransport(t *testing.T) {
firstTimeout := 15 * time.Second
secondTimeout := 45 * time.Second
first := newEmbeddingHTTPClient(firstTimeout)
second := newEmbeddingHTTPClient(secondTimeout)
if first == second {
t.Fatal("expected distinct HTTP clients")
}
firstGuard, ok := first.Transport.(*secutils.SSRFValidatingRoundTripper)
if !ok {
t.Fatalf("expected SSRF-validating transport, got %T", first.Transport)
}
secondGuard, ok := second.Transport.(*secutils.SSRFValidatingRoundTripper)
if !ok {
t.Fatalf("expected SSRF-validating transport, got %T", second.Transport)
}
if firstGuard.Base != secondGuard.Base {
t.Fatal("expected embedding HTTP clients to share a base transport")
}
if firstGuard.Base != http.RoundTripper(sharedEmbeddingHTTPTransport) {
t.Fatal("expected embedding HTTP client to use the shared transport")
}
if first.Timeout != firstTimeout {
t.Fatalf("unexpected first client timeout: got %v, want %v", first.Timeout, firstTimeout)
}
if second.Timeout != secondTimeout {
t.Fatalf("unexpected second client timeout: got %v, want %v", second.Timeout, secondTimeout)
}
}
func TestValidateEmbeddingBaseURL_RejectsLoopback(t *testing.T) {
err := validateEmbeddingBaseURL("http://169.254.169.254/latest/meta-data")
if err == nil {
t.Fatal("expected SSRF error for link-local metadata URL")
}
if !strings.Contains(err.Error(), "SSRF") {
t.Fatalf("unexpected error: %v", err)
}
}
func TestValidateEmbeddingBaseURL_AllowsEmpty(t *testing.T) {
if err := validateEmbeddingBaseURL(""); err != nil {
t.Fatalf("empty base URL should be allowed: %v", err)
}
}
func TestNewOpenAIEmbedder_RejectsPrivateBaseURL(t *testing.T) {
_, err := NewOpenAIEmbedder(
"test-key",
"http://169.254.169.254/latest/meta-data",
"text-embedding-3-small",
511,
256,
"model-id",
nil,
)
if err == nil {
t.Fatal("expected SSRF rejection for link-local metadata URL")
}
}
+282
View File
@@ -0,0 +1,282 @@
package embedding
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
"time"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/models/requestbudget"
secutils "github.com/Tencent/WeKnora/internal/utils"
)
const (
// VolcengineMultimodalEmbeddingPath 火山引擎 Ark 多模态 Embedding API 路径
VolcengineMultimodalEmbeddingPath = "/api/v3/embeddings/multimodal"
)
// VolcengineEmbedder implements text vectorization using Volcengine Ark multimodal embedding API
type VolcengineEmbedder struct {
apiKey string
baseURL string
modelName string
truncatePromptTokens int
dimensions int
modelID string
httpClient *http.Client
timeout time.Duration
maxRetries int
customHeaders map[string]string
supportsDimensionOverride bool
EmbedderPooler
}
// SetCustomHeaders 设置用户自定义 HTTP 请求头(类似 OpenAI Python SDK 的 extra_headers)。
func (e *VolcengineEmbedder) SetCustomHeaders(headers map[string]string) {
e.customHeaders = headers
}
func (e *VolcengineEmbedder) SetSupportsDimensionOverride(supported bool) {
e.supportsDimensionOverride = supported
}
// VolcengineEmbedRequest represents a Volcengine Ark multimodal embedding request
type VolcengineEmbedRequest struct {
Model string `json:"model"`
Input []VolcengineInputContent `json:"input"`
Dimensions int `json:"dimensions,omitempty"`
}
// VolcengineInputContent represents a single input item for Volcengine
type VolcengineInputContent struct {
Type string `json:"type"`
Text string `json:"text,omitempty"`
ImageURL *VolcengineImageURL `json:"image_url,omitempty"`
}
// VolcengineImageURL represents the image URL structure for Volcengine
type VolcengineImageURL struct {
URL string `json:"url"`
}
// VolcengineEmbedResponse represents a Volcengine Ark multimodal embedding response
// Multimodal API returns data as an object with embedding array directly
type VolcengineEmbedResponse struct {
Object string `json:"object"`
Data struct {
Embedding []float32 `json:"embedding"`
} `json:"data"`
Model string `json:"model"`
Usage struct {
PromptTokens int `json:"prompt_tokens"`
TotalTokens int `json:"total_tokens"`
} `json:"usage"`
}
// VolcengineErrorResponse represents an error response from Volcengine
type VolcengineErrorResponse struct {
Error struct {
Code string `json:"code"`
Message string `json:"message"`
Type string `json:"type"`
} `json:"error"`
}
// NewVolcengineEmbedder creates a new Volcengine Ark embedder
func NewVolcengineEmbedder(apiKey, baseURL, modelName string,
truncatePromptTokens int, dimensions int, modelID string, pooler EmbedderPooler,
) (*VolcengineEmbedder, error) {
if baseURL == "" {
baseURL = "https://ark.cn-beijing.volces.com"
}
// Remove trailing slash
baseURL = strings.TrimRight(baseURL, "/")
// Extract base host if URL contains the full multimodal path
if strings.Contains(baseURL, "/embeddings/multimodal") {
// Strip the path to get base URL
if idx := strings.Index(baseURL, "/api/"); idx != -1 {
baseURL = baseURL[:idx]
}
} else if strings.HasSuffix(baseURL, "/api/v3") {
// If it ends with /api/v3, keep just the host
baseURL = strings.TrimSuffix(baseURL, "/api/v3")
}
if modelName == "" {
return nil, fmt.Errorf("model name is required")
}
if truncatePromptTokens == 0 {
truncatePromptTokens = 511
}
timeout := 60 * time.Second
if err := validateEmbeddingBaseURL(baseURL); err != nil {
return nil, err
}
return &VolcengineEmbedder{
apiKey: apiKey,
baseURL: baseURL,
modelName: modelName,
httpClient: newEmbeddingHTTPClient(timeout),
truncatePromptTokens: truncatePromptTokens,
EmbedderPooler: pooler,
dimensions: dimensions,
modelID: modelID,
timeout: timeout,
maxRetries: 3,
}, nil
}
// Embed converts text to vector
func (e *VolcengineEmbedder) Embed(ctx context.Context, text string) ([]float32, error) {
for range 3 {
embeddings, err := e.BatchEmbed(ctx, []string{text})
if err != nil {
return nil, err
}
if len(embeddings) > 0 {
return embeddings[0], nil
}
}
return nil, fmt.Errorf("no embedding returned")
}
func (e *VolcengineEmbedder) doRequestWithRetry(ctx context.Context, jsonData []byte) (*http.Response, error) {
ctx = requestbudget.WithModel(ctx, e.modelID)
var resp *http.Response
var err error
url := e.baseURL + VolcengineMultimodalEmbeddingPath
for i := 0; i <= e.maxRetries; i++ {
if err := requestbudget.Check(ctx); err != nil {
return nil, err
}
if i > 0 {
backoffTime := time.Duration(1<<uint(i-1)) * time.Second
if backoffTime > 10*time.Second {
backoffTime = 10 * time.Second
}
logger.GetLogger(ctx).
Infof("VolcengineEmbedder retrying request (%d/%d), waiting %v", i, e.maxRetries, backoffTime)
select {
case <-time.After(backoffTime):
case <-ctx.Done():
return nil, ctx.Err()
}
}
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(jsonData))
if err != nil {
logger.GetLogger(ctx).Errorf("VolcengineEmbedder failed to create request: %v", err)
continue
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+e.apiKey)
secutils.ApplyCustomHeaders(req, e.customHeaders)
resp, err = e.httpClient.Do(req)
if stopProviderRetry(ctx, err) {
return nil, err
}
if err == nil {
return resp, nil
}
logger.GetLogger(ctx).Errorf("VolcengineEmbedder request failed (attempt %d/%d): %v", i+1, e.maxRetries+1, err)
}
return nil, err
}
func (e *VolcengineEmbedder) BatchEmbed(ctx context.Context, texts []string) ([][]float32, error) {
embeddings := make([][]float32, len(texts))
// Volcengine multimodal API returns a single combined embedding for all inputs,
// so we need to call the API once per text for proper batch embedding
for i, text := range texts {
input := []VolcengineInputContent{
{
Type: "text",
Text: text,
},
}
reqBody := VolcengineEmbedRequest{
Model: e.modelName,
Input: input,
}
if e.supportsDimensionsParam() {
reqBody.Dimensions = e.dimensions
}
jsonData, err := json.Marshal(reqBody)
if err != nil {
logger.GetLogger(ctx).Errorf("VolcengineEmbedder BatchEmbed marshal request error: %v", err)
return nil, fmt.Errorf("marshal request: %w", err)
}
resp, err := e.doRequestWithRetry(ctx, jsonData)
if err != nil {
logger.GetLogger(ctx).Errorf("VolcengineEmbedder BatchEmbed send request error: %v", err)
return nil, fmt.Errorf("send request: %w", err)
}
body, err := io.ReadAll(resp.Body)
resp.Body.Close()
if err != nil {
logger.GetLogger(ctx).Errorf("VolcengineEmbedder BatchEmbed read response error: %v", err)
return nil, fmt.Errorf("read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
var errResp VolcengineErrorResponse
if json.Unmarshal(body, &errResp) == nil && errResp.Error.Message != "" {
logger.GetLogger(ctx).Errorf("VolcengineEmbedder BatchEmbed API error: %s - %s", errResp.Error.Code, errResp.Error.Message)
return nil, fmt.Errorf("API error: %s - %s", errResp.Error.Code, errResp.Error.Message)
}
logger.GetLogger(ctx).Errorf("VolcengineEmbedder BatchEmbed API error: Http Status %s", resp.Status)
return nil, fmt.Errorf("BatchEmbed API error: Http Status %s", resp.Status)
}
var response VolcengineEmbedResponse
if err := json.Unmarshal(body, &response); err != nil {
logger.GetLogger(ctx).Errorf("VolcengineEmbedder BatchEmbed unmarshal response error: %v", err)
return nil, fmt.Errorf("unmarshal response: %w", err)
}
embeddings[i] = response.Data.Embedding
}
return embeddings, nil
}
// GetModelName returns the model name
func (e *VolcengineEmbedder) GetModelName() string {
return e.modelName
}
func (e *VolcengineEmbedder) supportsDimensionsParam() bool {
return e.supportsDimensionOverride && e.dimensions > 0
}
// GetDimensions returns the vector dimensions
func (e *VolcengineEmbedder) GetDimensions() int {
return e.dimensions
}
// GetModelID returns the model ID
func (e *VolcengineEmbedder) GetModelID() string {
return e.modelID
}
+178
View File
@@ -0,0 +1,178 @@
package embedding
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
"time"
"github.com/Tencent/WeKnora/internal/models/provider"
"github.com/Tencent/WeKnora/internal/models/requestbudget"
"github.com/Tencent/WeKnora/internal/models/utils"
"github.com/google/uuid"
)
const weKnoraCloudEmbedPath = "/api/v1/embeddings"
// WeKnoraCloudEmbedder 实现 embedding.Embedder 接口,对接 WeKnoraCloud /api/v1/embeddings
type WeKnoraCloudEmbedder struct {
modelName string
remoteModelName string
modelID string
appID string
apiKey string
baseURL string
dimensions int
supportsDimensionOverride bool
client *http.Client
EmbedderPooler
}
// NewWeKnoraCloudEmbedder 构造 WeKnoraCloudEmbedder
func NewWeKnoraCloudEmbedder(config Config) (*WeKnoraCloudEmbedder, error) {
if config.AppID == "" {
return nil, fmt.Errorf("WeKnoraCloud embedder: AppID is required")
}
if config.AppSecret == "" {
return nil, fmt.Errorf("WeKnoraCloud embedder: AppSecret is required")
}
remoteModelName := ""
if config.ExtraConfig != nil {
remoteModelName = strings.TrimSpace(config.ExtraConfig["remote_model_name"])
}
baseURL := strings.TrimRight(config.BaseURL, "/")
if baseURL == "" {
baseURL = provider.WeKnoraCloudBaseURL
}
if err := validateEmbeddingBaseURL(baseURL); err != nil {
return nil, err
}
return &WeKnoraCloudEmbedder{
modelName: config.ModelName,
remoteModelName: remoteModelName,
modelID: config.ModelID,
appID: config.AppID,
apiKey: config.AppSecret,
baseURL: baseURL,
dimensions: config.Dimensions,
supportsDimensionOverride: config.SupportsDimensionOverride,
client: newEmbeddingHTTPClient(60 * time.Second),
}, nil
}
type weKnoraCloudEmbedRequest struct {
Model string `json:"model"`
Input []string `json:"input"`
Dimensions int `json:"dimensions,omitempty"`
TruncatePromptTokens int `json:"truncate_prompt_tokens,omitempty"`
}
type weKnoraCloudEmbedResponse struct {
Data []struct {
Index int `json:"index"`
Embedding []float32 `json:"embedding"`
} `json:"data"`
}
func (e *WeKnoraCloudEmbedder) Embed(ctx context.Context, text string) ([]float32, error) {
results, err := e.BatchEmbed(ctx, []string{text})
if err != nil {
return nil, err
}
if len(results) == 0 {
return nil, fmt.Errorf("weknoracloud embedder: empty response")
}
return results[0], nil
}
func (e *WeKnoraCloudEmbedder) BatchEmbed(ctx context.Context, texts []string) ([][]float32, error) {
ctx = requestbudget.WithModel(ctx, e.modelID)
if err := requestbudget.Check(ctx); err != nil {
return nil, err
}
reqBody := weKnoraCloudEmbedRequest{Model: e.effectiveModelName(), Input: texts}
if e.supportsDimensionOverride && e.dimensions > 0 {
reqBody.Dimensions = e.dimensions
}
bodyBytes, err := json.Marshal(reqBody)
if err != nil {
return nil, fmt.Errorf("weknoracloud embedder: marshal: %w", err)
}
requestID := uuid.New().String()
headers := utils.Sign(e.appID, e.apiKey, requestID, string(bodyBytes))
req, err := http.NewRequestWithContext(ctx, http.MethodPost, e.baseURL+weKnoraCloudEmbedPath, bytes.NewReader(bodyBytes))
if err != nil {
return nil, fmt.Errorf("weknoracloud embedder: create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
for k, v := range headers {
req.Header.Set(k, v)
}
resp, err := e.client.Do(req)
if err != nil {
return nil, fmt.Errorf("weknoracloud embedder: do request: %w", err)
}
defer resp.Body.Close()
respBytes, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("weknoracloud embedder: read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("weknoracloud embedder: status %d: %s", resp.StatusCode, string(respBytes))
}
var embedResp weKnoraCloudEmbedResponse
if err := json.Unmarshal(respBytes, &embedResp); err != nil {
return nil, fmt.Errorf("weknoracloud embedder: unmarshal: %w", err)
}
result := make([][]float32, len(texts))
seen := make([]bool, len(texts))
for _, item := range embedResp.Data {
if item.Index < 0 || item.Index >= len(result) {
return nil, fmt.Errorf(
"weknoracloud embedder: response index %d out of range for %d inputs",
item.Index,
len(texts),
)
}
if seen[item.Index] {
return nil, fmt.Errorf("weknoracloud embedder: duplicate response index %d", item.Index)
}
result[item.Index] = item.Embedding
seen[item.Index] = true
}
for index, found := range seen {
if !found {
return nil, fmt.Errorf("weknoracloud embedder: missing embedding for input index %d", index)
}
}
return result, nil
}
func (e *WeKnoraCloudEmbedder) BatchEmbedWithPool(ctx context.Context, model Embedder, texts []string) ([][]float32, error) {
return e.BatchEmbed(ctx, texts)
}
func (e *WeKnoraCloudEmbedder) SetSupportsDimensionOverride(supported bool) {
e.supportsDimensionOverride = supported
}
func (e *WeKnoraCloudEmbedder) effectiveModelName() string {
if e.remoteModelName != "" {
return e.remoteModelName
}
return e.modelName
}
func (e *WeKnoraCloudEmbedder) GetModelName() string { return e.modelName }
func (e *WeKnoraCloudEmbedder) GetModelID() string { return e.modelID }
func (e *WeKnoraCloudEmbedder) GetDimensions() int { return e.dimensions }
@@ -0,0 +1,100 @@
package embedding
import (
"context"
"fmt"
"net/http"
"net/http/httptest"
"reflect"
"strings"
"testing"
)
func newWeKnoraCloudEmbedderTestServer(t *testing.T, response string) *WeKnoraCloudEmbedder {
t.Helper()
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != weKnoraCloudEmbedPath {
http.Error(w, fmt.Sprintf("request path = %q, want %q", r.URL.Path, weKnoraCloudEmbedPath), http.StatusNotFound)
return
}
w.Header().Set("Content-Type", "application/json")
_, _ = fmt.Fprint(w, response)
}))
t.Cleanup(server.Close)
return &WeKnoraCloudEmbedder{
modelName: "test-embedding",
modelID: "test-model-id",
appID: "test-app-id",
apiKey: "test-api-key",
baseURL: server.URL,
client: server.Client(),
}
}
func TestWeKnoraCloudBatchEmbedPreservesInputOrder(t *testing.T) {
embedder := newWeKnoraCloudEmbedderTestServer(t, `{
"data": [
{"index": 1, "embedding": [0.3, 0.4]},
{"index": 0, "embedding": [0.1, 0.2]}
]
}`)
got, err := embedder.BatchEmbed(context.Background(), []string{"first", "second"})
if err != nil {
t.Fatalf("BatchEmbed returned error: %v", err)
}
want := [][]float32{{0.1, 0.2}, {0.3, 0.4}}
if !reflect.DeepEqual(got, want) {
t.Fatalf("BatchEmbed result = %v, want %v", got, want)
}
}
func TestWeKnoraCloudBatchEmbedRejectsMalformedResponse(t *testing.T) {
tests := []struct {
name string
response string
texts []string
wantErrMsg string
}{
{
name: "negative index",
response: `{"data": [{"index": -1, "embedding": [0.1, 0.2]}]}`,
texts: []string{"first"},
wantErrMsg: "response index -1 out of range",
},
{
name: "index above input range",
response: `{"data": [{"index": 1, "embedding": [0.1, 0.2]}]}`,
texts: []string{"first"},
wantErrMsg: "response index 1 out of range",
},
{
name: "duplicate index",
response: `{"data": [{"index": 0, "embedding": [0.1]}, {"index": 0, "embedding": [0.2]}]}`,
texts: []string{"first"},
wantErrMsg: "duplicate response index 0",
},
{
name: "missing result",
response: `{"data": [{"index": 0, "embedding": [0.1, 0.2]}]}`,
texts: []string{"first", "second"},
wantErrMsg: "missing embedding for input index 1",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
embedder := newWeKnoraCloudEmbedderTestServer(t, tt.response)
_, err := embedder.BatchEmbed(context.Background(), tt.texts)
if err == nil {
t.Fatalf("BatchEmbed returned nil error, want error containing %q", tt.wantErrMsg)
}
if !strings.Contains(err.Error(), tt.wantErrMsg) {
t.Fatalf("BatchEmbed error = %q, want it to contain %q", err, tt.wantErrMsg)
}
})
}
}
+264
View File
@@ -0,0 +1,264 @@
package embedding
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"time"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/models/provider"
"github.com/Tencent/WeKnora/internal/models/requestbudget"
secutils "github.com/Tencent/WeKnora/internal/utils"
)
// ZhipuEmbedder implements text vectorization functionality using Zhipu AI API
type ZhipuEmbedder struct {
apiKey string
baseURL string
modelName string
truncatePromptTokens int
dimensions int
modelID string
httpClient *http.Client
timeout time.Duration
maxRetries int
customHeaders map[string]string
supportsDimensionOverride bool
EmbedderPooler
}
// ZhipuEmbedRequest represents a Zhipu embedding request
type ZhipuEmbedRequest struct {
Model string `json:"model"`
Input []string `json:"input"`
Dimensions int `json:"dimensions,omitempty"`
TruncatePromptTokens int `json:"truncate_prompt_tokens,omitempty"`
}
// ZhipuEmbedResponse represents a Zhipu embedding response
type ZhipuEmbedResponse struct {
Data []struct {
Embedding []float32 `json:"embedding"`
Index int `json:"index"`
} `json:"data"`
Model string `json:"model"`
}
// NewZhipuEmbedder creates a new Zhipu embedder
func NewZhipuEmbedder(apiKey, baseURL, modelName string,
truncatePromptTokens int, dimensions int, modelID string, pooler EmbedderPooler,
) (*ZhipuEmbedder, error) {
if baseURL == "" {
baseURL = provider.ZhipuEmbeddingBaseURL
}
if modelName == "" {
return nil, fmt.Errorf("model name is required")
}
if truncatePromptTokens == 0 {
truncatePromptTokens = 511
}
timeout := 60 * time.Second
if err := validateEmbeddingBaseURL(baseURL); err != nil {
return nil, err
}
return &ZhipuEmbedder{
apiKey: apiKey,
baseURL: baseURL,
modelName: modelName,
httpClient: newEmbeddingHTTPClient(timeout),
truncatePromptTokens: truncatePromptTokens,
EmbedderPooler: pooler,
dimensions: dimensions,
modelID: modelID,
timeout: timeout,
maxRetries: 3, // Maximum retry count
}, nil
}
// SetCustomHeaders sets custom HTTP headers for the embedder
func (e *ZhipuEmbedder) SetCustomHeaders(headers map[string]string) {
e.customHeaders = headers
}
func (e *ZhipuEmbedder) SetSupportsDimensionOverride(supported bool) {
e.supportsDimensionOverride = supported
}
// Embed converts text to vector
func (e *ZhipuEmbedder) Embed(ctx context.Context, text string) ([]float32, error) {
for range 3 {
embeddings, err := e.BatchEmbed(ctx, []string{text})
if err != nil {
return nil, err
}
if len(embeddings) > 0 {
return embeddings[0], nil
}
}
return nil, fmt.Errorf("no embedding returned")
}
func (e *ZhipuEmbedder) doRequestWithRetry(ctx context.Context, jsonData []byte) (*http.Response, error) {
ctx = requestbudget.WithModel(ctx, e.modelID)
var resp *http.Response
var err error
url := e.baseURL + "/embeddings"
for i := 0; i <= e.maxRetries; i++ {
if err := requestbudget.Check(ctx); err != nil {
return nil, err
}
if i > 0 {
backoffTime := time.Duration(1<<uint(i-1)) * time.Second
if backoffTime > 10*time.Second {
backoffTime = 10 * time.Second
}
logger.GetLogger(ctx).
Infof("ZhipuEmbedder retrying request (%d/%d), waiting %v", i, e.maxRetries, backoffTime)
select {
case <-time.After(backoffTime):
case <-ctx.Done():
return nil, ctx.Err()
}
}
var req *http.Request
req, err = http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(jsonData))
if err != nil {
logger.GetLogger(ctx).Errorf("ZhipuEmbedder failed to create request: %v", err)
continue
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+e.apiKey)
secutils.ApplyCustomHeaders(req, e.customHeaders)
resp, err = e.httpClient.Do(req)
if stopProviderRetry(ctx, err) {
return nil, err
}
if err == nil {
return resp, nil
}
logger.GetLogger(ctx).Errorf("ZhipuEmbedder request failed (attempt %d/%d): %v", i+1, e.maxRetries+1, err)
}
return nil, err
}
func (e *ZhipuEmbedder) BatchEmbed(ctx context.Context, texts []string) ([][]float32, error) {
// Create request body
reqBody := ZhipuEmbedRequest{
Model: e.modelName,
Input: texts,
TruncatePromptTokens: e.truncatePromptTokens,
}
if e.supportsDimensionsParam() {
reqBody.Dimensions = e.dimensions
}
jsonData, err := json.Marshal(reqBody)
if err != nil {
logger.GetLogger(ctx).Errorf("ZhipuEmbedder BatchEmbed marshal request error: %v", err)
return nil, fmt.Errorf("marshal request: %w", err)
}
// Log request details for debugging
logger.GetLogger(ctx).Debugf("ZhipuEmbedder BatchEmbed: model=%s, input_count=%d, truncate_tokens=%d",
e.modelName, len(texts), e.truncatePromptTokens)
// Check for invalid input lengths and log details
hasInvalidLength := false
for i, text := range texts {
textLen := len(text)
textPreview := text
if len(textPreview) > 200 {
textPreview = textPreview[:200] + "..."
}
// Log warning if length is outside valid range [1, 8192]
if textLen == 0 || textLen > 8192 {
hasInvalidLength = true
logger.GetLogger(ctx).Errorf("ZhipuEmbedder BatchEmbed input[%d]: INVALID length=%d (must be [1, 8192]), preview=%s",
i, textLen, textPreview)
} else {
logger.GetLogger(ctx).Debugf("ZhipuEmbedder BatchEmbed input[%d]: length=%d, preview=%s",
i, textLen, textPreview)
}
}
if hasInvalidLength {
logger.GetLogger(ctx).Errorf("ZhipuEmbedder BatchEmbed: Found invalid input lengths, this will likely cause API error")
}
// Send request (passing jsonData instead of constructing http.Request)
resp, err := e.doRequestWithRetry(ctx, jsonData)
if err != nil {
logger.GetLogger(ctx).Errorf("ZhipuEmbedder BatchEmbed send request error: %v", err)
return nil, fmt.Errorf("send request: %w", err)
}
if resp.Body != nil {
defer resp.Body.Close()
}
// Read response
body, err := io.ReadAll(resp.Body)
if err != nil {
logger.GetLogger(ctx).Errorf("ZhipuEmbedder BatchEmbed read response error: %v", err)
return nil, fmt.Errorf("read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
// Log detailed error response from OpenAI API
bodyStr := string(body)
if len(bodyStr) > 1000 {
bodyStr = bodyStr[:1000] + "... (truncated)"
}
logger.GetLogger(ctx).Errorf("ZhipuEmbedder BatchEmbed API error: Http Status %s, Response Body: %s", resp.Status, bodyStr)
return nil, fmt.Errorf("BatchEmbed API error: Http Status %s, Response: %s", resp.Status, bodyStr)
}
// Parse response
var response ZhipuEmbedResponse
if err := json.Unmarshal(body, &response); err != nil {
logger.GetLogger(ctx).Errorf("ZhipuEmbedder BatchEmbed unmarshal response error: %v", err)
return nil, fmt.Errorf("unmarshal response: %w", err)
}
// Extract embedding vectors
embeddings := make([][]float32, 0, len(response.Data))
for _, data := range response.Data {
embeddings = append(embeddings, data.Embedding)
}
return embeddings, nil
}
// GetModelName returns the model name
func (e *ZhipuEmbedder) GetModelName() string {
return e.modelName
}
func (e *ZhipuEmbedder) supportsDimensionsParam() bool {
return e.supportsDimensionOverride && e.dimensions > 0
}
// GetDimensions returns the vector dimensions
func (e *ZhipuEmbedder) GetDimensions() int {
return e.dimensions
}
// GetModelID returns the model ID
func (e *ZhipuEmbedder) GetModelID() string {
return e.modelID
}
+174
View File
@@ -0,0 +1,174 @@
package limiter
import (
"context"
"sort"
"sync"
"github.com/Tencent/WeKnora/internal/types"
)
// The concurrency governor is process-wide, shared by every model-client layer
// that fronts a provider (chat, vlm). Keeping the singleton here — rather than
// inside one client package — lets all of them gate against the same limiter
// and per-model limit without importing each other. Wired once at startup (see
// container.registerModelConcurrencyLimiter) via SetGovernor.
var (
governorMu sync.RWMutex
governor ModelConcurrencyLimiter
governorN int
)
// SetGovernor installs the process-wide background concurrency governor and the
// default per-model limit. Passing a nil limiter or a non-positive limit
// disables governance (all calls pass through). Safe to call at startup.
func SetGovernor(l ModelConcurrencyLimiter, limit int) {
governorMu.Lock()
defer governorMu.Unlock()
governor = l
governorN = limit
}
// SetGlobalLimit updates ONLY the process-wide default per-model limit,
// leaving the installed limiter backend intact. Used by the system-settings
// runtime bridge so an operator can retune model.max_concurrency without a
// restart. A non-positive value disables the default (models that carry their
// own MaxConcurrency still honour it).
func SetGlobalLimit(limit int) {
governorMu.Lock()
defer governorMu.Unlock()
governorN = limit
}
// Gate acquires a per-model concurrency slot using the process-wide default
// limit. Equivalent to GateN(ctx, modelID, 0).
func Gate(ctx context.Context, modelID string) func() {
return GateN(ctx, modelID, 0)
}
// GateN acquires a per-model concurrency slot when the call is a background task
// (see types.IsBackgroundTask) and a governor is installed. modelLimit is the
// model's own configured cap; a value <= 0 means "fall back to the process-wide
// default" (governorN). It returns a release func that is ALWAYS safe to call:
// on the passthrough / fail-open paths it is a cheap no-op. The gate never
// blocks a call permanently — a limiter/Redis outage or a cancelled context
// fails open.
func GateN(ctx context.Context, modelID string, modelLimit int) func() {
return GateNamedN(ctx, modelID, "", modelLimit)
}
func GateNamedN(ctx context.Context, modelID, modelName string, modelLimit int) func() {
governorMu.RLock()
l, defaultLimit := governor, governorN
governorMu.RUnlock()
limit := modelLimit
if limit <= 0 {
limit = defaultLimit
}
if l == nil || limit <= 0 || !types.IsBackgroundTask(ctx) {
return noop
}
if named, ok := l.(interface{ SetModelName(string, string) }); ok {
named.SetModelName(modelID, modelName)
}
release, err := l.Acquire(ctx, modelID, limit)
if err != nil || release == nil {
return noop
}
return release
}
// RuntimeStats returns the semaphores observed by this process. Redis-backed
// limiters report cluster-wide active holders; waiting callers remain local to
// this instance by design.
func RuntimeStats(ctx context.Context) ([]RuntimeStat, bool, error) {
governorMu.RLock()
l := governor
governorMu.RUnlock()
inspector, ok := l.(runtimeInspectable)
if !ok || inspector == nil {
return []RuntimeStat{}, false, nil
}
stats, err := inspector.RuntimeStats(ctx)
if stats == nil {
stats = []RuntimeStat{}
}
return stats, true, err
}
// localLimiter is an in-process (single-node) counting semaphore keyed by
// model ID. It is the Lite-mode counterpart to the Redis limiter: Lite runs a
// single process with no Redis, so a shared distributed semaphore is neither
// available nor needed — but background ingestion can still burst the whole
// worker pool against one provider, so we still cap concurrency locally.
type localLimiter struct {
mu sync.Mutex
sems map[string]chan struct{}
tracked map[string]*trackedSemaphore
}
// NewLocalLimiter builds an in-process per-key concurrency limiter.
func NewLocalLimiter() ModelConcurrencyLimiter {
return &localLimiter{sems: make(map[string]chan struct{}), tracked: make(map[string]*trackedSemaphore)}
}
func (l *localLimiter) Acquire(ctx context.Context, key string, limit int) (func(), error) {
if l == nil || limit <= 0 || key == "" {
return noop, nil
}
l.mu.Lock()
sem, ok := l.sems[key]
tracked := l.tracked[key]
if tracked == nil {
tracked = &trackedSemaphore{}
l.tracked[key] = tracked
}
tracked.limit.Store(int64(limit))
if !ok {
// Capacity is fixed at first use for a key; the limit is a
// process-wide constant, so it never changes across acquires.
sem = make(chan struct{}, limit)
l.sems[key] = sem
}
l.mu.Unlock()
tracked.waiting.Add(1)
defer tracked.waiting.Add(-1)
select {
case sem <- struct{}{}:
var once sync.Once
return func() { once.Do(func() { <-sem }) }, nil
case <-ctx.Done():
// Fail open on cancellation, mirroring the Redis limiter.
return noop, nil
}
}
func (l *localLimiter) RuntimeStats(_ context.Context) ([]RuntimeStat, error) {
l.mu.Lock()
defer l.mu.Unlock()
stats := make([]RuntimeStat, 0, len(l.sems))
for modelID, sem := range l.sems {
tracked := l.tracked[modelID]
name, _ := tracked.name.Load().(string)
stats = append(stats, RuntimeStat{ModelID: modelID, Name: name, Active: int64(len(sem)), Waiting: tracked.waiting.Load(), Limit: int(tracked.limit.Load())})
}
sort.Slice(stats, func(i, j int) bool { return stats[i].ModelID < stats[j].ModelID })
return stats, nil
}
func (l *localLimiter) SetModelName(modelID, name string) {
if modelID == "" || name == "" {
return
}
l.mu.Lock()
tracked := l.tracked[modelID]
if tracked == nil {
tracked = &trackedSemaphore{}
l.tracked[modelID] = tracked
}
l.mu.Unlock()
tracked.name.Store(name)
}
+133
View File
@@ -0,0 +1,133 @@
package limiter
import (
"context"
"sync"
"testing"
"time"
"github.com/Tencent/WeKnora/internal/types"
)
func withBackground() context.Context {
return types.WithBackgroundTask(context.Background())
}
// TestGateOnlyGovernsBackground verifies interactive calls always pass through
// (no-op release) and only background calls consult the installed limiter.
func TestGateOnlyGovernsBackground(t *testing.T) {
t.Cleanup(func() { SetGovernor(nil, 0) })
SetGovernor(NewLocalLimiter(), 1)
// Interactive (non-background) ctx: never gated, even at limit 1.
rel1 := Gate(context.Background(), "m")
rel2 := Gate(context.Background(), "m")
rel1()
rel2()
// Background ctx: gated. First slot taken; second would block, so a
// cancelled second acquire must fail open (return a usable release).
held := Gate(withBackground(), "m")
defer held()
ctx, cancel := context.WithTimeout(withBackground(), 30*time.Millisecond)
defer cancel()
rel := Gate(ctx, "m")
if rel == nil {
t.Fatal("cancelled background gate must fail open with a usable release")
}
rel()
}
// TestGateDisabledWhenNoGovernor verifies Gate is a passthrough when no
// governor is installed.
func TestGateDisabledWhenNoGovernor(t *testing.T) {
SetGovernor(nil, 0)
if rel := Gate(withBackground(), "m"); rel == nil {
t.Fatal("gate with no governor must return a usable release")
} else {
rel()
}
}
// TestLocalLimiterCaps verifies the in-process limiter enforces the cap and
// releases slots.
func TestLocalLimiterCaps(t *testing.T) {
l := NewLocalLimiter()
ctx := context.Background()
r1, _ := l.Acquire(ctx, "k", 2)
r2, _ := l.Acquire(ctx, "k", 2)
// Third acquire must block while full.
done := make(chan func(), 1)
go func() {
r3, _ := l.Acquire(ctx, "k", 2)
done <- r3
}()
select {
case <-done:
t.Fatal("third acquire should block while at limit")
case <-time.After(50 * time.Millisecond):
}
r1() // free a slot; the waiter proceeds
select {
case r3 := <-done:
r3()
case <-time.After(time.Second):
t.Fatal("waiter should proceed after a slot frees")
}
r2()
}
// TestLocalLimiterIndependentKeys verifies per-key budgets are independent.
func TestLocalLimiterIndependentKeys(t *testing.T) {
l := NewLocalLimiter()
ctx := context.Background()
ra, _ := l.Acquire(ctx, "a", 1)
rb, _ := l.Acquire(ctx, "b", 1) // different key: must not block
ra()
rb()
}
// TestLocalLimiterFailOpen verifies degraded inputs pass through.
func TestLocalLimiterFailOpen(t *testing.T) {
l := NewLocalLimiter()
ctx := context.Background()
if r, err := l.Acquire(ctx, "", 4); err != nil || r == nil {
t.Fatal("empty key should fail open")
} else {
r()
}
if r, err := l.Acquire(ctx, "k", 0); err != nil || r == nil {
t.Fatal("limit<=0 should fail open")
} else {
r()
}
}
// TestLocalLimiterReleaseIdempotent verifies double release doesn't
// over-free a slot.
func TestLocalLimiterReleaseIdempotent(t *testing.T) {
l := NewLocalLimiter()
ctx := context.Background()
r1, _ := l.Acquire(ctx, "k", 1)
var wg sync.WaitGroup
wg.Add(2)
go func() { defer wg.Done(); r1() }()
go func() { defer wg.Done(); r1() }()
wg.Wait()
// After idempotent release the single slot is free again.
r2, _ := l.Acquire(ctx, "k", 1)
done := make(chan struct{})
go func() { r3, _ := l.Acquire(ctx, "k", 1); r3(); close(done) }()
select {
case <-done:
t.Fatal("a double-release must not leave the slot permanently free")
case <-time.After(50 * time.Millisecond):
}
r2()
<-done
}
+239
View File
@@ -0,0 +1,239 @@
// Package limiter provides a distributed, per-key concurrency governor for
// outbound model-provider calls. The shared finite resource is the model
// provider (its request/concurrency budget), so concurrency is capped at the
// model-client layer — keyed by model ID — rather than at the asynq queue layer
// (queue weights are scheduling priority, not throttling).
//
// The Redis implementation is a self-healing distributed semaphore built on a
// sorted set: each held slot is a ZSET member (unique token) scored by its
// lease expiry. Acquire atomically prunes expired leases, counts live holders,
// and admits a new one only while under the limit. A background heartbeat
// refreshes the lease so long calls keep their slot; a crashed holder's lease
// simply expires and is reclaimed. Every backend error fails OPEN (the call is
// allowed) so a limiter/Redis outage can never halt model traffic.
package limiter
import (
"context"
"sort"
"strconv"
"sync"
"sync/atomic"
"time"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/google/uuid"
"github.com/redis/go-redis/v9"
)
// ModelConcurrencyLimiter caps the number of concurrent in-flight calls per
// key (typically a model ID) across all processes sharing the same backend.
type ModelConcurrencyLimiter interface {
// Acquire blocks until a slot for key is available or ctx is done. It
// returns a release func that MUST be invoked to free the slot. On any
// backend error (or ctx cancellation) it fails open: release is a no-op and
// err is nil, so callers proceed without a slot rather than dropping the
// call.
Acquire(ctx context.Context, key string, limit int) (release func(), err error)
}
// RuntimeStat is a point-in-time view of a model semaphore. Active is
// cluster-wide for the Redis backend and process-local in Lite mode. Waiting is
// deliberately process-local: waiters block in application processes and are
// not represented in Redis.
type RuntimeStat struct {
ModelID string `json:"model_id"`
Name string `json:"name"`
Active int64 `json:"active"`
Waiting int64 `json:"waiting"`
Limit int `json:"limit"`
}
type runtimeInspectable interface {
RuntimeStats(context.Context) ([]RuntimeStat, error)
}
type trackedSemaphore struct {
limit atomic.Int64
waiting atomic.Int64
name atomic.Value // string
}
// noop is the release returned on the fail-open / passthrough paths.
func noop() {}
const (
// defaultLeaseTTL is the crash-recovery window, not a request timeout.
// Live calls refresh their lease every ttl/3, so even a very long provider
// request keeps its slot. Keeping this short prevents an app/container
// restart from leaving an entire model budget apparently occupied for
// minutes while the replacement workers are blocked behind dead holders.
defaultLeaseTTL = 30 * time.Second
// defaultPollInterval is how often a waiting acquirer re-checks for a free
// slot. Small enough to stay responsive, large enough to avoid hammering
// Redis under contention.
defaultPollInterval = 200 * time.Millisecond
// keyPrefix namespaces the semaphore ZSETs in Redis.
keyPrefix = "weknora:modelsem:"
)
// acquireScript atomically prunes expired leases, counts live holders, and
// admits the caller (adding its token scored by lease expiry) only while the
// live count is below the limit. Returns 1 on admission, 0 when full.
//
// KEYS[1] = semaphore ZSET key
// ARGV[1] = now (unix ms)
// ARGV[2] = limit
// ARGV[3] = caller token
// ARGV[4] = lease TTL (ms)
var acquireScript = redis.NewScript(`
redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1])
local count = redis.call('ZCARD', KEYS[1])
if count < tonumber(ARGV[2]) then
redis.call('ZADD', KEYS[1], ARGV[1] + ARGV[4], ARGV[3])
redis.call('PEXPIRE', KEYS[1], ARGV[4] * 2)
return 1
end
return 0
`)
type redisLimiter struct {
rdb *redis.Client
ttl time.Duration
pollInterval time.Duration
tracked sync.Map // model ID -> *trackedSemaphore
}
// NewRedisLimiter builds a distributed limiter backed by rdb. A nil client
// yields a limiter that always fails open.
func NewRedisLimiter(rdb *redis.Client) ModelConcurrencyLimiter {
return &redisLimiter{
rdb: rdb,
ttl: defaultLeaseTTL,
pollInterval: defaultPollInterval,
}
}
func (l *redisLimiter) Acquire(ctx context.Context, key string, limit int) (func(), error) {
if l == nil || l.rdb == nil || limit <= 0 || key == "" {
return noop, nil
}
zkey := keyPrefix + key
entry, _ := l.tracked.LoadOrStore(key, &trackedSemaphore{})
tracked := entry.(*trackedSemaphore)
tracked.limit.Store(int64(limit))
tracked.waiting.Add(1)
defer tracked.waiting.Add(-1)
token := uuid.NewString()
ttlMs := l.ttl.Milliseconds()
// Reuse a single timer across poll iterations rather than allocating a new
// one via time.After each loop: under sustained contention a waiter can
// spin thousands of times, and every time.After timer lives until it fires.
// Start it stopped so the first Reset below arms it cleanly.
timer := time.NewTimer(0)
if !timer.Stop() {
<-timer.C
}
defer timer.Stop()
for {
now := time.Now().UnixMilli()
res, err := acquireScript.Run(ctx, l.rdb, []string{zkey},
now, limit, token, ttlMs).Int()
if err != nil {
// Fail open: a limiter outage must never block model traffic.
logger.Warnf(ctx, "[ModelLimiter] acquire failed for key=%s, failing open: %v", key, err)
return noop, nil
}
if res == 1 {
return l.hold(zkey, token), nil
}
timer.Reset(l.pollInterval)
select {
case <-ctx.Done():
// Fail open on cancellation too: let the inner call observe the
// cancelled context and return its own error, rather than us
// synthesising one here.
return noop, nil
case <-timer.C:
}
}
}
func (l *redisLimiter) RuntimeStats(ctx context.Context) ([]RuntimeStat, error) {
stats := make([]RuntimeStat, 0)
if l == nil || l.rdb == nil {
return stats, nil
}
var firstErr error
now := time.Now().UnixMilli()
l.tracked.Range(func(rawKey, rawValue any) bool {
modelID := rawKey.(string)
tracked := rawValue.(*trackedSemaphore)
active, err := l.rdb.ZCount(ctx, keyPrefix+modelID, strconv.FormatInt(now+1, 10), "+inf").Result()
if err != nil {
if firstErr == nil {
firstErr = err
}
return true
}
name, _ := tracked.name.Load().(string)
stats = append(stats, RuntimeStat{ModelID: modelID, Name: name, Active: active, Waiting: tracked.waiting.Load(), Limit: int(tracked.limit.Load())})
return true
})
sort.Slice(stats, func(i, j int) bool { return stats[i].ModelID < stats[j].ModelID })
return stats, firstErr
}
func (l *redisLimiter) SetModelName(modelID, name string) {
if modelID == "" || name == "" {
return
}
entry, _ := l.tracked.LoadOrStore(modelID, &trackedSemaphore{})
entry.(*trackedSemaphore).name.Store(name)
}
// hold starts a heartbeat that refreshes the lease and returns an idempotent
// release that stops the heartbeat and drops the slot.
func (l *redisLimiter) hold(zkey, token string) func() {
stop := make(chan struct{})
go func() {
t := time.NewTicker(l.ttl / 3)
defer t.Stop()
for {
select {
case <-stop:
return
case <-t.C:
now := time.Now().UnixMilli()
// Detached ctx: the heartbeat must outlive request ctx up to
// release. Best-effort; a failed refresh just risks early
// reclamation, which the limit already tolerates.
//
// Refresh BOTH the member lease score AND the ZSET key's own
// TTL. The acquire script only PEXPIREs the key on admission,
// so a semaphore that stays saturated with no slot turnover
// would otherwise let the whole key expire after ttl*2 —
// dropping every live lease and admitting over the limit. The
// heartbeat pushes the key TTL out in lockstep with the lease.
bg := context.Background()
_ = l.rdb.ZAdd(bg, zkey, redis.Z{
Score: float64(now + l.ttl.Milliseconds()),
Member: token,
}).Err()
_ = l.rdb.PExpire(bg, zkey, l.ttl*2).Err()
}
}
}()
var once sync.Once
return func() {
once.Do(func() {
close(stop)
_ = l.rdb.ZRem(context.Background(), zkey, token).Err()
})
}
}
+175
View File
@@ -0,0 +1,175 @@
package limiter
import (
"context"
"testing"
"time"
"github.com/alicebob/miniredis/v2"
"github.com/redis/go-redis/v9"
)
func newTestLimiter(t *testing.T, ttl, poll time.Duration) (*redisLimiter, *miniredis.Miniredis, *redis.Client) {
t.Helper()
s, err := miniredis.Run()
if err != nil {
t.Fatalf("start miniredis: %v", err)
}
t.Cleanup(s.Close)
rdb := redis.NewClient(&redis.Options{Addr: s.Addr()})
t.Cleanup(func() { _ = rdb.Close() })
return &redisLimiter{rdb: rdb, ttl: ttl, pollInterval: poll}, s, rdb
}
func zcard(t *testing.T, rdb *redis.Client, key string) int64 {
t.Helper()
n, err := rdb.ZCard(context.Background(), keyPrefix+key).Result()
if err != nil {
t.Fatalf("zcard: %v", err)
}
return n
}
// TestRedisLimiterCapsConcurrency verifies the live-holder count never exceeds
// the limit and that a freed slot lets a waiting acquirer through.
func TestRedisLimiterCapsConcurrency(t *testing.T) {
lim, _, rdb := newTestLimiter(t, time.Minute, 10*time.Millisecond)
ctx := context.Background()
const key = "m1"
r1, err := lim.Acquire(ctx, key, 2)
if err != nil {
t.Fatalf("acquire 1: %v", err)
}
r2, err := lim.Acquire(ctx, key, 2)
if err != nil {
t.Fatalf("acquire 2: %v", err)
}
if got := zcard(t, rdb, key); got != 2 {
t.Fatalf("expected 2 live holders, got %d", got)
}
// The third acquire must block while the semaphore is full.
done := make(chan func(), 1)
go func() {
r3, _ := lim.Acquire(ctx, key, 2)
done <- r3
}()
select {
case <-done:
t.Fatal("third acquire should block while at limit")
case <-time.After(100 * time.Millisecond):
}
if got := zcard(t, rdb, key); got != 2 {
t.Fatalf("count must not exceed limit while blocked, got %d", got)
}
// Freeing a slot lets the waiter in.
r1()
var r3 func()
select {
case r3 = <-done:
case <-time.After(2 * time.Second):
t.Fatal("third acquire should proceed after a slot frees")
}
if got := zcard(t, rdb, key); got != 2 {
t.Fatalf("expected 2 live holders after handoff, got %d", got)
}
r2()
r3()
if got := zcard(t, rdb, key); got != 0 {
t.Fatalf("expected 0 live holders after release, got %d", got)
}
}
// TestRedisLimiterReclaimsExpiredLease verifies a crashed holder's stale lease
// (no heartbeat) is pruned so its slot is reclaimed by a new acquirer.
func TestRedisLimiterReclaimsExpiredLease(t *testing.T) {
lim, _, rdb := newTestLimiter(t, time.Minute, 10*time.Millisecond)
ctx := context.Background()
const key = "m2"
// Simulate a dead holder: a member whose lease already expired.
expired := float64(time.Now().Add(-time.Second).UnixMilli())
if err := rdb.ZAdd(ctx, keyPrefix+key, redis.Z{Score: expired, Member: "dead"}).Err(); err != nil {
t.Fatalf("seed dead holder: %v", err)
}
// limit=1 would be full if the dead lease still counted; pruning reclaims it.
release, err := lim.Acquire(ctx, key, 1)
if err != nil {
t.Fatalf("acquire after expiry: %v", err)
}
if got := zcard(t, rdb, key); got != 1 {
t.Fatalf("expected only the live holder, got %d", got)
}
if exists, _ := rdb.ZScore(ctx, keyPrefix+key, "dead").Result(); exists != 0 {
t.Fatal("expired lease should have been pruned")
}
release()
}
// TestRedisLimiterIndependentKeys verifies budgets are per key.
func TestRedisLimiterIndependentKeys(t *testing.T) {
lim, _, _ := newTestLimiter(t, time.Minute, 10*time.Millisecond)
ctx := context.Background()
ra, err := lim.Acquire(ctx, "a", 1)
if err != nil {
t.Fatalf("acquire a: %v", err)
}
// A different key has its own budget, so this must not block.
rb, err := lim.Acquire(ctx, "b", 1)
if err != nil {
t.Fatalf("acquire b: %v", err)
}
ra()
rb()
}
// TestRedisLimiterFailsOpen verifies every degraded path allows the call
// instead of blocking model traffic.
func TestRedisLimiterFailsOpen(t *testing.T) {
ctx := context.Background()
// Nil client.
lim := &redisLimiter{rdb: nil, ttl: time.Minute, pollInterval: 10 * time.Millisecond}
if release, err := lim.Acquire(ctx, "k", 4); err != nil || release == nil {
t.Fatalf("nil client should fail open, got release==nil=%v err=%v", release == nil, err)
}
// Non-positive limit is a no-op governor.
lim2, _, _ := newTestLimiter(t, time.Minute, 10*time.Millisecond)
if release, err := lim2.Acquire(ctx, "k", 0); err != nil || release == nil {
t.Fatalf("limit<=0 should fail open, got release==nil=%v err=%v", release == nil, err)
}
// Backend error (server down) must fail open too.
lim3, s, _ := newTestLimiter(t, time.Minute, 10*time.Millisecond)
s.Close()
if release, err := lim3.Acquire(ctx, "k", 4); err != nil || release == nil {
t.Fatalf("backend error should fail open, got release==nil=%v err=%v", release == nil, err)
}
}
// TestRedisLimiterCancelledContextFailsOpen verifies a waiter whose context is
// cancelled while blocked returns a usable (no-op) release rather than erroring.
func TestRedisLimiterCancelledContextFailsOpen(t *testing.T) {
lim, _, _ := newTestLimiter(t, time.Minute, 10*time.Millisecond)
r1, err := lim.Acquire(context.Background(), "m3", 1)
if err != nil {
t.Fatalf("acquire: %v", err)
}
defer r1()
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
defer cancel()
release, err := lim.Acquire(ctx, "m3", 1)
if err != nil || release == nil {
t.Fatalf("cancelled wait should fail open, got release==nil=%v err=%v", release == nil, err)
}
release()
}
+76
View File
@@ -0,0 +1,76 @@
package provider
import (
"fmt"
"strings"
"github.com/Tencent/WeKnora/internal/types"
)
const (
// AliyunChatBaseURL 阿里云 DashScope Chat/Embedding 的默认 BaseURL
AliyunChatBaseURL = "https://dashscope.aliyuncs.com/compatible-mode/v1"
// AliyunRerankBaseURL 阿里云 DashScope Rerank 的默认 BaseURL
AliyunRerankBaseURL = "https://dashscope.aliyuncs.com/api/v1/services/rerank/text-rerank/text-rerank"
)
// AliyunProvider 实现阿里云 DashScope 的 Provider 接口
type AliyunProvider struct{}
func init() {
Register(&AliyunProvider{})
}
// Info 返回阿里云 provider 的元数据
func (p *AliyunProvider) Info() ProviderInfo {
return ProviderInfo{
Name: ProviderAliyun,
DisplayName: "阿里云 DashScope",
Description: "qwen-plus, tongyi-embedding-vision-plus, qwen3-rerank, etc.",
DefaultURLs: map[types.ModelType]string{
types.ModelTypeKnowledgeQA: AliyunChatBaseURL,
types.ModelTypeEmbedding: AliyunChatBaseURL,
types.ModelTypeRerank: AliyunRerankBaseURL,
types.ModelTypeVLLM: AliyunChatBaseURL,
},
ModelTypes: []types.ModelType{
types.ModelTypeKnowledgeQA,
types.ModelTypeEmbedding,
types.ModelTypeRerank,
types.ModelTypeVLLM,
},
RequiresAuth: true,
}
}
// ValidateConfig 验证阿里云 provider 配置
func (p *AliyunProvider) ValidateConfig(config *Config) error {
if config.APIKey == "" {
return fmt.Errorf("API key is required for Aliyun DashScope")
}
if config.ModelName == "" {
return fmt.Errorf("model name is required")
}
return nil
}
// IsQwenThinkingModel 检查模型名是否为支持思维链的 Qwen 模型
// 支持思维链的模型需要特殊处理 enable_thinking 参数
func IsQwenThinkingModel(modelName string) bool {
lowerName := strings.ToLower(modelName)
return strings.HasPrefix(lowerName, "qwen3") ||
strings.HasPrefix(lowerName, "qwen-plus") ||
strings.HasPrefix(lowerName, "qwen-max") ||
strings.HasPrefix(lowerName, "qwen-turbo")
}
// IsQwen3Model checks whether the model belongs to the Qwen3 family only.
func IsQwen3Model(modelName string) bool {
return strings.HasPrefix(strings.ToLower(modelName), "qwen3")
}
// IsDeepSeekModel 检查模型名是否为 DeepSeek 模型
// DeepSeek 模型不支持 tool_choice 参数
func IsDeepSeekModel(modelName string) bool {
return strings.Contains(strings.ToLower(modelName), "deepseek")
}
+41
View File
@@ -0,0 +1,41 @@
package provider
import (
"fmt"
"github.com/Tencent/WeKnora/internal/types"
)
const AnthropicBaseURL = "https://api.anthropic.com/v1"
// AnthropicProvider implements native Anthropic Messages API metadata.
type AnthropicProvider struct{}
func init() {
Register(&AnthropicProvider{})
}
func (p *AnthropicProvider) Info() ProviderInfo {
return ProviderInfo{
Name: ProviderAnthropic,
DisplayName: "Anthropic",
Description: "Claude models via native Anthropic Messages API",
DefaultURLs: map[types.ModelType]string{
types.ModelTypeKnowledgeQA: AnthropicBaseURL,
},
ModelTypes: []types.ModelType{
types.ModelTypeKnowledgeQA,
},
RequiresAuth: true,
}
}
func (p *AnthropicProvider) ValidateConfig(config *Config) error {
if config.APIKey == "" {
return fmt.Errorf("API key is required for Anthropic provider")
}
if config.ModelName == "" {
return fmt.Errorf("model name is required")
}
return nil
}
+61
View File
@@ -0,0 +1,61 @@
package provider
import (
"fmt"
"github.com/Tencent/WeKnora/internal/types"
)
// AzureOpenAIProvider 实现 Azure OpenAI 的 Provider 接口
type AzureOpenAIProvider struct{}
func init() {
Register(&AzureOpenAIProvider{})
}
// Info 返回 Azure OpenAI provider 的元数据
func (p *AzureOpenAIProvider) Info() ProviderInfo {
return ProviderInfo{
Name: ProviderAzureOpenAI,
DisplayName: "Azure OpenAI",
Description: "gpt-4o, gpt-4, text-embedding-ada-002, etc.",
DefaultURLs: map[types.ModelType]string{
types.ModelTypeKnowledgeQA: "https://{resource}.openai.azure.com",
types.ModelTypeEmbedding: "https://{resource}.openai.azure.com",
types.ModelTypeRerank: "https://{resource}.openai.azure.com",
types.ModelTypeVLLM: "https://{resource}.openai.azure.com",
types.ModelTypeASR: "https://{resource}.openai.azure.com",
},
ModelTypes: []types.ModelType{
types.ModelTypeKnowledgeQA,
types.ModelTypeEmbedding,
types.ModelTypeVLLM,
types.ModelTypeASR,
},
RequiresAuth: true,
ExtraFields: []ExtraFieldConfig{
{
Key: "api_version",
Label: "API Version",
Type: "string",
Required: false,
Default: "2024-10-21",
Placeholder: "e.g. 2024-10-21",
},
},
}
}
// ValidateConfig 验证 Azure OpenAI provider 配置
func (p *AzureOpenAIProvider) ValidateConfig(config *Config) error {
if config.APIKey == "" {
return fmt.Errorf("API key is required for Azure OpenAI provider")
}
if config.ModelName == "" {
return fmt.Errorf("deployment name (model name) is required")
}
if config.BaseURL == "" {
return fmt.Errorf("Azure resource endpoint (base URL) is required")
}
return nil
}
+46
View File
@@ -0,0 +1,46 @@
package provider
import (
"fmt"
"github.com/Tencent/WeKnora/internal/types"
)
const (
// DeepSeekBaseURL DeepSeek 官方 API BaseURL
DeepSeekBaseURL = "https://api.deepseek.com/v1"
)
// DeepSeekProvider 实现 DeepSeek 的 Provider 接口
type DeepSeekProvider struct{}
func init() {
Register(&DeepSeekProvider{})
}
// Info 返回 DeepSeek provider 的元数据
func (p *DeepSeekProvider) Info() ProviderInfo {
return ProviderInfo{
Name: ProviderDeepSeek,
DisplayName: "DeepSeek",
Description: "deepseek-chat, deepseek-reasoner, etc.",
DefaultURLs: map[types.ModelType]string{
types.ModelTypeKnowledgeQA: DeepSeekBaseURL,
},
ModelTypes: []types.ModelType{
types.ModelTypeKnowledgeQA,
},
RequiresAuth: true,
}
}
// ValidateConfig 验证 DeepSeek provider 配置
func (p *DeepSeekProvider) ValidateConfig(config *Config) error {
if config.APIKey == "" {
return fmt.Errorf("API key is required for DeepSeek provider")
}
if config.ModelName == "" {
return fmt.Errorf("model name is required")
}
return nil
}
+50
View File
@@ -0,0 +1,50 @@
package provider
import (
"fmt"
"github.com/Tencent/WeKnora/internal/types"
)
const (
// GeminiBaseURL Google Gemini API BaseURL
GeminiBaseURL = "https://generativelanguage.googleapis.com/v1beta"
// GeminiOpenAICompatBaseURL Gemini OpenAI 兼容模式 BaseURL
GeminiOpenAICompatBaseURL = "https://generativelanguage.googleapis.com/v1beta/openai"
)
// GeminiProvider 实现 Google Gemini 的 Provider 接口
type GeminiProvider struct{}
func init() {
Register(&GeminiProvider{})
}
// Info 返回 Gemini provider 的元数据
func (p *GeminiProvider) Info() ProviderInfo {
return ProviderInfo{
Name: ProviderGemini,
DisplayName: "Google Gemini",
Description: "gemini-3-flash-preview, gemini-2.5-pro, gemini-embedding-2, etc.",
DefaultURLs: map[types.ModelType]string{
types.ModelTypeKnowledgeQA: GeminiOpenAICompatBaseURL,
types.ModelTypeEmbedding: GeminiBaseURL,
},
ModelTypes: []types.ModelType{
types.ModelTypeKnowledgeQA,
types.ModelTypeEmbedding,
},
RequiresAuth: true,
}
}
// ValidateConfig 验证 Gemini provider 配置
func (p *GeminiProvider) ValidateConfig(config *Config) error {
if config.APIKey == "" {
return fmt.Errorf("API key is required for Google Gemini provider")
}
if config.ModelName == "" {
return fmt.Errorf("model name is required")
}
return nil
}
+43
View File
@@ -0,0 +1,43 @@
package provider
import (
"fmt"
"github.com/Tencent/WeKnora/internal/types"
)
// GenericProvider 实现通用 OpenAI 兼容的 Provider 接口
type GenericProvider struct{}
func init() {
Register(&GenericProvider{})
}
// Info 返回通用 provider 的元数据
func (p *GenericProvider) Info() ProviderInfo {
return ProviderInfo{
Name: ProviderGeneric,
DisplayName: "自定义 (OpenAI兼容接口)",
Description: "Generic API endpoint (OpenAI-compatible)",
DefaultURLs: map[types.ModelType]string{}, // 需要用户自行配置填写
ModelTypes: []types.ModelType{
types.ModelTypeKnowledgeQA,
types.ModelTypeEmbedding,
types.ModelTypeRerank,
types.ModelTypeVLLM,
types.ModelTypeASR,
},
RequiresAuth: false, // 可能需要也可能不需要
}
}
// ValidateConfig 验证通用 provider 配置
func (p *GenericProvider) ValidateConfig(config *Config) error {
if config.BaseURL == "" {
return fmt.Errorf("base URL is required for generic provider")
}
if config.ModelName == "" {
return fmt.Errorf("model name is required")
}
return nil
}
+59
View File
@@ -0,0 +1,59 @@
package provider
import (
"fmt"
"github.com/Tencent/WeKnora/internal/types"
)
const (
// GPUStackBaseURL GPUStack API BaseURL (OpenAI 兼容模式)
GPUStackBaseURL = "http://your_gpustack_server_url/v1-openai"
// GPUStackRerankBaseURL GPUStack Rerank API 虽然兼容OpenAI,但路径不同 (/v1/rerank 而非 /v1-openai/rerank)
GPUStackRerankBaseURL = "http://your_gpustack_server_url/v1"
)
// GPUStackProvider 实现 GPUStack 的 Provider 接口
type GPUStackProvider struct{}
func init() {
Register(&GPUStackProvider{})
}
// Info 返回 GPUStack provider 的元数据
func (p *GPUStackProvider) Info() ProviderInfo {
return ProviderInfo{
Name: ProviderGPUStack,
DisplayName: "GPUStack",
Description: "Choose your deployed model on GPUStack",
DefaultURLs: map[types.ModelType]string{
types.ModelTypeKnowledgeQA: GPUStackBaseURL,
types.ModelTypeEmbedding: GPUStackBaseURL,
types.ModelTypeRerank: GPUStackRerankBaseURL,
types.ModelTypeVLLM: GPUStackBaseURL,
types.ModelTypeASR: GPUStackBaseURL,
},
ModelTypes: []types.ModelType{
types.ModelTypeKnowledgeQA,
types.ModelTypeEmbedding,
types.ModelTypeRerank,
types.ModelTypeVLLM,
types.ModelTypeASR,
},
RequiresAuth: true, // GPUStack 需要 API Key
}
}
// ValidateConfig 验证 GPUStack provider 配置
func (p *GPUStackProvider) ValidateConfig(config *Config) error {
if config.BaseURL == "" {
return fmt.Errorf("base URL is required for GPUStack provider")
}
if config.APIKey == "" {
return fmt.Errorf("API key is required for GPUStack provider")
}
if config.ModelName == "" {
return fmt.Errorf("model name is required")
}
return nil
}
+48
View File
@@ -0,0 +1,48 @@
package provider
import (
"fmt"
"github.com/Tencent/WeKnora/internal/types"
)
const (
// HunyuanBaseURL 腾讯混元 API BaseURL (OpenAI 兼容模式)
HunyuanBaseURL = "https://api.hunyuan.cloud.tencent.com/v1"
)
// HunyuanProvider 实现腾讯混元的 Provider 接口
type HunyuanProvider struct{}
func init() {
Register(&HunyuanProvider{})
}
// Info 返回腾讯混元 provider 的元数据
func (p *HunyuanProvider) Info() ProviderInfo {
return ProviderInfo{
Name: ProviderHunyuan,
DisplayName: "腾讯混元 Hunyuan",
Description: "hunyuan-pro, hunyuan-standard, hunyuan-embedding, etc.",
DefaultURLs: map[types.ModelType]string{
types.ModelTypeKnowledgeQA: HunyuanBaseURL,
types.ModelTypeEmbedding: HunyuanBaseURL,
},
ModelTypes: []types.ModelType{
types.ModelTypeKnowledgeQA,
types.ModelTypeEmbedding,
},
RequiresAuth: true,
}
}
// ValidateConfig 验证腾讯混元 provider 配置
func (p *HunyuanProvider) ValidateConfig(config *Config) error {
if config.APIKey == "" {
return fmt.Errorf("API key is required for Hunyuan provider")
}
if config.ModelName == "" {
return fmt.Errorf("model name is required")
}
return nil
}
+44
View File
@@ -0,0 +1,44 @@
package provider
import (
"fmt"
"github.com/Tencent/WeKnora/internal/types"
)
const (
JinaBaseURL = "https://api.jina.ai/v1"
)
// JinaProvider 实现 Jina AI 的 Provider 接口
type JinaProvider struct{}
func init() {
Register(&JinaProvider{})
}
// Info 返回 Jina AI provider 的元数据
func (p *JinaProvider) Info() ProviderInfo {
return ProviderInfo{
Name: ProviderJina,
DisplayName: "Jina",
Description: "jina-clip-v1, jina-embeddings-v2-base-zh, etc.",
DefaultURLs: map[types.ModelType]string{
types.ModelTypeEmbedding: JinaBaseURL,
types.ModelTypeRerank: JinaBaseURL,
},
ModelTypes: []types.ModelType{
types.ModelTypeEmbedding,
types.ModelTypeRerank,
},
RequiresAuth: true,
}
}
// ValidateConfig 验证 Jina AI provider 配置
func (p *JinaProvider) ValidateConfig(config *Config) error {
if config.APIKey == "" {
return fmt.Errorf("API key is required for Jina AI provider")
}
return nil
}
+55
View File
@@ -0,0 +1,55 @@
package provider
import (
"fmt"
"github.com/Tencent/WeKnora/internal/types"
)
const (
// LiteLLMBaseURL is a placeholder the user must replace with a reachable
// LiteLLM proxy. Loopback defaults (localhost:4000) are rejected by SSRF
// unless the host is added to SSRF_WHITELIST; Docker also cannot reach a
// proxy on the host via localhost. The hostname contains "litellm" so
// DetectProvider recognises the catalog default.
LiteLLMBaseURL = "http://your_litellm_proxy/v1"
)
// LiteLLMProvider implements the Provider interface for LiteLLM
// (https://github.com/BerriAI/litellm). LiteLLM exposes a single
// OpenAI-compatible endpoint that routes to 100+ providers (OpenAI, Anthropic,
// Gemini, Bedrock, Vertex, Azure, ...) behind one base URL and key, so it plugs
// into WeKnora's OpenAI-compatible transport like the other gateway providers.
type LiteLLMProvider struct{}
func init() {
Register(&LiteLLMProvider{})
}
// Info returns the metadata for the LiteLLM provider.
func (p *LiteLLMProvider) Info() ProviderInfo {
return ProviderInfo{
Name: ProviderLiteLLM,
DisplayName: "LiteLLM",
Description: "Self-hosted LiteLLM proxy: one OpenAI-compatible endpoint to 100+ providers.",
DefaultURLs: map[types.ModelType]string{
types.ModelTypeKnowledgeQA: LiteLLMBaseURL,
types.ModelTypeEmbedding: LiteLLMBaseURL,
types.ModelTypeVLLM: LiteLLMBaseURL,
},
ModelTypes: []types.ModelType{
types.ModelTypeKnowledgeQA,
types.ModelTypeEmbedding,
types.ModelTypeVLLM,
},
RequiresAuth: true,
}
}
// ValidateConfig validates the LiteLLM provider configuration.
func (p *LiteLLMProvider) ValidateConfig(config *Config) error {
if config.APIKey == "" {
return fmt.Errorf("API key is required for LiteLLM provider")
}
return nil
}
+69
View File
@@ -0,0 +1,69 @@
package provider
import (
"fmt"
"strings"
"github.com/Tencent/WeKnora/internal/types"
)
const (
// LKEAPBaseURL 腾讯云知识引擎原子能力 (LKEAP) 兼容 OpenAI 协议的 BaseURL
LKEAPBaseURL = "https://api.lkeap.cloud.tencent.com/v1"
// LKEAPRerankBaseURL 腾讯云知识引擎原子能力 Rerank API 域名(TC3 签名)
LKEAPRerankBaseURL = "https://lkeap.tencentcloudapi.com"
)
// LKEAPProvider 实现腾讯云 LKEAP 的 Provider 接口
// 支持 DeepSeek-R1, DeepSeek-V3 系列模型,具备思维链能力
type LKEAPProvider struct{}
func init() {
Register(&LKEAPProvider{})
}
// Info 返回 LKEAP provider 的元数据
func (p *LKEAPProvider) Info() ProviderInfo {
return ProviderInfo{
Name: ProviderLKEAP,
DisplayName: "腾讯云 LKEAP",
Description: "DeepSeek-R1, DeepSeek-V3, lke-reranker-base 等",
DefaultURLs: map[types.ModelType]string{
types.ModelTypeKnowledgeQA: LKEAPBaseURL,
types.ModelTypeRerank: LKEAPRerankBaseURL,
},
ModelTypes: []types.ModelType{
types.ModelTypeKnowledgeQA,
types.ModelTypeRerank,
},
RequiresAuth: true,
}
}
// ValidateConfig 验证 LKEAP provider 配置
func (p *LKEAPProvider) ValidateConfig(config *Config) error {
if config.APIKey == "" {
return fmt.Errorf("API key is required for LKEAP provider")
}
if config.ModelName == "" {
return fmt.Errorf("model name is required")
}
return nil
}
// IsLKEAPDeepSeekV3Model 检查是否为 DeepSeek V3.x 系列模型
// V3.x 系列支持通过 Thinking 参数控制思维链开关
func IsLKEAPDeepSeekV3Model(modelName string) bool {
return strings.Contains(strings.ToLower(modelName), "deepseek-v3")
}
// IsLKEAPDeepSeekR1Model 检查是否为 DeepSeek R1 系列模型
// R1 系列默认开启思维链
func IsLKEAPDeepSeekR1Model(modelName string) bool {
return strings.Contains(strings.ToLower(modelName), "deepseek-r1")
}
// IsLKEAPThinkingModel 检查是否为支持思维链的 LKEAP 模型
func IsLKEAPThinkingModel(modelName string) bool {
return IsLKEAPDeepSeekR1Model(modelName) || IsLKEAPDeepSeekV3Model(modelName)
}
+48
View File
@@ -0,0 +1,48 @@
package provider
import (
"fmt"
"github.com/Tencent/WeKnora/internal/types"
)
const (
LongCatBaseURL = "https://api.longcat.chat/openai/v1"
)
// LongCatProvider 实现 LongCat AI 的 Provider 接口
type LongCatProvider struct{}
func init() {
Register(&LongCatProvider{})
}
// Info 返回 LongCat provider 的元数据
func (p *LongCatProvider) Info() ProviderInfo {
return ProviderInfo{
Name: ProviderLongCat,
DisplayName: "LongCat AI",
Description: "LongCat-Flash-Chat, LongCat-Flash-Thinking, etc.",
DefaultURLs: map[types.ModelType]string{
types.ModelTypeKnowledgeQA: LongCatBaseURL,
},
ModelTypes: []types.ModelType{
types.ModelTypeKnowledgeQA,
},
RequiresAuth: true,
}
}
// ValidateConfig 验证 LongCat provider 配置
func (p *LongCatProvider) ValidateConfig(config *Config) error {
if config.BaseURL == "" {
return fmt.Errorf("base URL is required for LongCat provider")
}
if config.APIKey == "" {
return fmt.Errorf("API key is required for LongCat provider")
}
if config.ModelName == "" {
return fmt.Errorf("model name is required")
}
return nil
}
+46
View File
@@ -0,0 +1,46 @@
package provider
import (
"fmt"
"github.com/Tencent/WeKnora/internal/types"
)
const (
// MimoBaseURL 小米 Mimo API BaseURL
MimoBaseURL = "https://api.xiaomimimo.com/v1"
)
// MimoProvider 实现小米 Mimo 的 Provider 接口
type MimoProvider struct{}
func init() {
Register(&MimoProvider{})
}
// Info 返回小米 Mimo provider 的元数据
func (p *MimoProvider) Info() ProviderInfo {
return ProviderInfo{
Name: ProviderMimo,
DisplayName: "小米 MiMo",
Description: "mimo-v2-flash",
DefaultURLs: map[types.ModelType]string{
types.ModelTypeKnowledgeQA: MimoBaseURL,
},
ModelTypes: []types.ModelType{
types.ModelTypeKnowledgeQA,
},
RequiresAuth: true,
}
}
// ValidateConfig 验证小米 Mimo provider 配置
func (p *MimoProvider) ValidateConfig(config *Config) error {
if config.APIKey == "" {
return fmt.Errorf("API key is required for Mimo provider")
}
if config.ModelName == "" {
return fmt.Errorf("model name is required")
}
return nil
}
+48
View File
@@ -0,0 +1,48 @@
package provider
import (
"fmt"
"github.com/Tencent/WeKnora/internal/types"
)
const (
// MiniMaxBaseURL MiniMax 国际版 API BaseURL
MiniMaxBaseURL = "https://api.minimax.io/v1"
// MiniMaxCNBaseURL MiniMax 国内版 API BaseURL
MiniMaxCNBaseURL = "https://api.minimaxi.com/v1"
)
// MiniMaxProvider 实现 MiniMax 的 Provider 接口
type MiniMaxProvider struct{}
func init() {
Register(&MiniMaxProvider{})
}
// Info 返回 MiniMax provider 的元数据
func (p *MiniMaxProvider) Info() ProviderInfo {
return ProviderInfo{
Name: ProviderMiniMax,
DisplayName: "MiniMax",
Description: "MiniMax-M3, MiniMax-M2.7, MiniMax-M2.7-highspeed, etc.",
DefaultURLs: map[types.ModelType]string{
types.ModelTypeKnowledgeQA: MiniMaxCNBaseURL,
},
ModelTypes: []types.ModelType{
types.ModelTypeKnowledgeQA,
},
RequiresAuth: true,
}
}
// ValidateConfig 验证 MiniMax provider 配置
func (p *MiniMaxProvider) ValidateConfig(config *Config) error {
if config.APIKey == "" {
return fmt.Errorf("API key is required for MiniMax provider")
}
if config.ModelName == "" {
return fmt.Errorf("model name is required")
}
return nil
}
+53
View File
@@ -0,0 +1,53 @@
package provider
import (
"fmt"
"github.com/Tencent/WeKnora/internal/types"
)
const (
// ModelScopeBaseURL ModelScope API BaseURL (OpenAI 兼容模式)
ModelScopeBaseURL = "https://api-inference.modelscope.cn/v1"
)
// ModelScopeProvider 实现 ModelScope (魔搭) 的 Provider 接口
type ModelScopeProvider struct{}
func init() {
Register(&ModelScopeProvider{})
}
// Info 返回 ModelScope provider 的元数据
func (p *ModelScopeProvider) Info() ProviderInfo {
return ProviderInfo{
Name: ProviderModelScope,
DisplayName: "魔搭 ModelScope",
Description: "Qwen/Qwen3-8B, Qwen/Qwen3-Embedding-8B, etc.",
DefaultURLs: map[types.ModelType]string{
types.ModelTypeKnowledgeQA: ModelScopeBaseURL,
types.ModelTypeEmbedding: ModelScopeBaseURL,
types.ModelTypeVLLM: ModelScopeBaseURL,
},
ModelTypes: []types.ModelType{
types.ModelTypeKnowledgeQA,
types.ModelTypeEmbedding,
types.ModelTypeVLLM,
},
RequiresAuth: true,
}
}
// ValidateConfig 验证 ModelScope provider 配置
func (p *ModelScopeProvider) ValidateConfig(config *Config) error {
if config.BaseURL == "" {
return fmt.Errorf("base URL is required for ModelScope provider")
}
if config.APIKey == "" {
return fmt.Errorf("API key is required for ModelScope provider")
}
if config.ModelName == "" {
return fmt.Errorf("model name is required")
}
return nil
}
+71
View File
@@ -0,0 +1,71 @@
package provider
import (
"fmt"
"strings"
"github.com/Tencent/WeKnora/internal/types"
)
const (
MoonshotBaseURL = "https://api.moonshot.ai/v1"
)
// IsMoonshotFixedTempModel reports whether the given Moonshot/Kimi model
// only accepts temperature=1. The following models reject any temperature
// value other than 1:
// - moonshot-v1 series (moonshot-v1-8k, moonshot-v1-32k, moonshot-v1-128k)
// - kimi-k2.5 and kimi-k2.6 (no temperature param in API docs)
//
// kimi-k2 / kimi-k2-turbo / kimi-k2-thinking models accept the full [0,1]
// range and are NOT affected.
func IsMoonshotFixedTempModel(modelName string) bool {
name := strings.ToLower(strings.TrimSpace(modelName))
if strings.HasPrefix(name, "moonshot-v1") {
return true
}
// kimi-k2.5, kimi-k2.6 — no temperature parameter supported
if name == "kimi-k2.5" || name == "kimi-k2.6" {
return true
}
return false
}
// MoonshotProvider 实现 Moonshot AI (Kimi) 的 Provider 接口
type MoonshotProvider struct{}
func init() {
Register(&MoonshotProvider{})
}
// Info 返回 Moonshot provider 的元数据
func (p *MoonshotProvider) Info() ProviderInfo {
return ProviderInfo{
Name: ProviderMoonshot,
DisplayName: "月之暗面 Moonshot",
Description: "kimi-k2-turbo-preview, moonshot-v1-8k-vision-preview, etc.",
DefaultURLs: map[types.ModelType]string{
types.ModelTypeKnowledgeQA: MoonshotBaseURL,
types.ModelTypeVLLM: MoonshotBaseURL,
},
ModelTypes: []types.ModelType{
types.ModelTypeKnowledgeQA,
types.ModelTypeVLLM,
},
RequiresAuth: true,
}
}
// ValidateConfig 验证 Moonshot provider 配置
func (p *MoonshotProvider) ValidateConfig(config *Config) error {
if config.BaseURL == "" {
return fmt.Errorf("base URL is required for Moonshot provider")
}
if config.APIKey == "" {
return fmt.Errorf("API key is required for Moonshot provider")
}
if config.ModelName == "" {
return fmt.Errorf("model name is required")
}
return nil
}
+50
View File
@@ -0,0 +1,50 @@
package provider
import (
"fmt"
"github.com/Tencent/WeKnora/internal/types"
)
const (
// NovitaOpenAIBaseURL Novita OpenAI-compatible API BaseURL
NovitaOpenAIBaseURL = "https://api.novita.ai/openai/v1"
)
// NovitaProvider 实现 Novita AI 的 Provider 接口
type NovitaProvider struct{}
func init() {
Register(&NovitaProvider{})
}
// Info 返回 Novita provider 的元数据
func (p *NovitaProvider) Info() ProviderInfo {
return ProviderInfo{
Name: ProviderNovita,
DisplayName: "Novita AI",
Description: "moonshotai/kimi-k2.5, zai-org/glm-5, minimax/minimax-m2.7, qwen/qwen3-embedding-0.6b, etc.",
DefaultURLs: map[types.ModelType]string{
types.ModelTypeKnowledgeQA: NovitaOpenAIBaseURL,
types.ModelTypeEmbedding: NovitaOpenAIBaseURL,
types.ModelTypeVLLM: NovitaOpenAIBaseURL,
},
ModelTypes: []types.ModelType{
types.ModelTypeKnowledgeQA,
types.ModelTypeEmbedding,
types.ModelTypeVLLM,
},
RequiresAuth: true,
}
}
// ValidateConfig 验证 Novita provider 配置
func (p *NovitaProvider) ValidateConfig(config *Config) error {
if config.APIKey == "" {
return fmt.Errorf("API key is required for Novita provider")
}
if config.ModelName == "" {
return fmt.Errorf("model name is required")
}
return nil
}
+54
View File
@@ -0,0 +1,54 @@
package provider
import (
"fmt"
"github.com/Tencent/WeKnora/internal/types"
)
const (
// NvidiaChatBaseURL NVIDIA Chat 的默认 BaseURL
NvidiaChatBaseURL = "https://integrate.api.nvidia.com/v1"
// NvidiaRerankBaseURL NVIDIA Rerank 的默认 BaseURL
NvidiaRerankBaseURL = "https://ai.api.nvidia.com/v1/retrieval/nvidia/reranking"
)
// NvidiaProvider 实现NVIDIA AI 的 Provider 接口
type NvidiaProvider struct{}
func init() {
Register(&NvidiaProvider{})
}
// Info 返回NVIDIA provider 的元数据
func (p *NvidiaProvider) Info() ProviderInfo {
return ProviderInfo{
Name: ProviderNvidia,
DisplayName: "NVIDIA",
Description: "deepseek-ai-deepseek-v3_1, nv-embed-v1, rerank-qa-mistral-4b, etc.",
DefaultURLs: map[types.ModelType]string{
types.ModelTypeKnowledgeQA: NvidiaChatBaseURL,
types.ModelTypeEmbedding: NvidiaChatBaseURL,
types.ModelTypeRerank: NvidiaRerankBaseURL,
types.ModelTypeVLLM: NvidiaChatBaseURL,
},
ModelTypes: []types.ModelType{
types.ModelTypeKnowledgeQA,
types.ModelTypeEmbedding,
types.ModelTypeRerank,
types.ModelTypeVLLM,
},
RequiresAuth: true,
}
}
// ValidateConfig 验证NVIDIA provider 配置
func (p *NvidiaProvider) ValidateConfig(config *Config) error {
if config.APIKey == "" {
return fmt.Errorf("API key is required for NVIDIA")
}
if config.ModelName == "" {
return fmt.Errorf("model name is required")
}
return nil
}
+86
View File
@@ -0,0 +1,86 @@
package provider
import (
"fmt"
"strings"
"github.com/Tencent/WeKnora/internal/types"
)
const (
OpenAIBaseURL = "https://api.openai.com/v1"
)
// OpenAIProvider 实现 OpenAI 的 Provider 接口
type OpenAIProvider struct{}
func init() {
Register(&OpenAIProvider{})
}
// Info 返回 OpenAI provider 的元数据
func (p *OpenAIProvider) Info() ProviderInfo {
return ProviderInfo{
Name: ProviderOpenAI,
DisplayName: "OpenAI",
Description: "gpt-5.2, gpt-5-mini, etc.",
DefaultURLs: map[types.ModelType]string{
types.ModelTypeKnowledgeQA: OpenAIBaseURL,
types.ModelTypeEmbedding: OpenAIBaseURL,
types.ModelTypeRerank: OpenAIBaseURL,
types.ModelTypeVLLM: OpenAIBaseURL,
types.ModelTypeASR: OpenAIBaseURL,
},
ModelTypes: []types.ModelType{
types.ModelTypeKnowledgeQA,
types.ModelTypeEmbedding,
types.ModelTypeRerank,
types.ModelTypeVLLM,
types.ModelTypeASR,
},
RequiresAuth: true,
}
}
// ValidateConfig 验证 OpenAI provider 配置
func (p *OpenAIProvider) ValidateConfig(config *Config) error {
if config.APIKey == "" {
return fmt.Errorf("API key is required for OpenAI provider")
}
if config.ModelName == "" {
return fmt.Errorf("model name is required")
}
return nil
}
// IsOpenAIReasoningOrGPT5Model 判断模型是否为 OpenAI / Azure OpenAI 的
// 推理类(o-series)或 GPT-5 系列模型。
//
// 这些模型在 OpenAI Chat Completions API 中:
// - 不再支持 `max_tokens`,必须使用 `max_completion_tokens`;
// - 仅支持默认的 `temperature=1`、`top_p=1`,且不支持 `frequency_penalty` /
// `presence_penalty` 等采样参数(传非默认值会被拒绝)。
//
// 参考:
// - https://platform.openai.com/docs/api-reference/chat
// - https://learn.microsoft.com/azure/ai-services/openai/how-to/reasoning
//
// 仅基于模型名做启发式匹配;对于 Azure OpenAI,因为模型名实际上是 deployment 名,
// 用户若用了自定义部署名我们无法识别,此时仍会按普通模型处理(保持原行为)。
func IsOpenAIReasoningOrGPT5Model(modelName string) bool {
name := strings.ToLower(strings.TrimSpace(modelName))
if name == "" {
return false
}
if strings.HasPrefix(name, "gpt-5") {
return true
}
// o1 / o1-mini / o1-preview / o3 / o3-mini / o4-mini ...
// 必须精确匹配,避免误命中 "openai-..." 之类。
for _, prefix := range []string{"o1", "o3", "o4"} {
if name == prefix || strings.HasPrefix(name, prefix+"-") {
return true
}
}
return false
}
+46
View File
@@ -0,0 +1,46 @@
package provider
import "testing"
// TestIsOpenAIReasoningOrGPT5Model 验证 GPT-5 / o-series 模型识别逻辑。
// 见 issue #1283:这些模型必须使用 max_completion_tokens 替代 max_tokens。
func TestIsOpenAIReasoningOrGPT5Model(t *testing.T) {
cases := []struct {
name string
input string
want bool
}{
{"empty", "", false},
{"gpt-5", "gpt-5", true},
{"gpt-5-mini", "gpt-5-mini", true},
{"gpt-5.2", "gpt-5.2", true},
{"gpt-5.5-pro", "gpt-5.5-pro", true},
{"gpt-5 mixed case", "GPT-5.4-Mini", true},
{"o1", "o1", true},
{"o1-mini", "o1-mini", true},
{"o1-preview", "o1-preview", true},
{"o3", "o3", true},
{"o3-mini", "o3-mini", true},
{"o4-mini", "o4-mini", true},
{"gpt-4", "gpt-4", false},
{"gpt-4o", "gpt-4o", false},
{"gpt-4o-mini", "gpt-4o-mini", false},
{"gpt-3.5-turbo", "gpt-3.5-turbo", false},
{"name starting with 'o1' but not o-series", "olympus-1", false},
{"name starting with 'openai-'", "openai-gpt-4", false},
{"name starting with 'o3' but not o-series", "o3xtra", false},
{"random", "qwen-max", false},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
got := IsOpenAIReasoningOrGPT5Model(tc.input)
if got != tc.want {
t.Errorf("IsOpenAIReasoningOrGPT5Model(%q) = %v, want %v", tc.input, got, tc.want)
}
})
}
}
+46
View File
@@ -0,0 +1,46 @@
package provider
import (
"fmt"
"github.com/Tencent/WeKnora/internal/types"
)
const (
OpenRouterBaseURL = "https://openrouter.ai/api/v1"
)
// OpenRouterProvider 实现 OpenRouter 的 Provider 接口
type OpenRouterProvider struct{}
func init() {
Register(&OpenRouterProvider{})
}
// Info 返回 OpenRouter provider 的元数据
func (p *OpenRouterProvider) Info() ProviderInfo {
return ProviderInfo{
Name: ProviderOpenRouter,
DisplayName: "OpenRouter",
Description: "openai/gpt-5.2-chat, google/gemini-3-flash-preview, etc.",
DefaultURLs: map[types.ModelType]string{
types.ModelTypeKnowledgeQA: OpenRouterBaseURL,
types.ModelTypeEmbedding: OpenRouterBaseURL,
types.ModelTypeVLLM: OpenRouterBaseURL,
},
ModelTypes: []types.ModelType{
types.ModelTypeKnowledgeQA,
types.ModelTypeEmbedding,
types.ModelTypeVLLM,
},
RequiresAuth: true,
}
}
// ValidateConfig 验证 OpenRouter provider 配置
func (p *OpenRouterProvider) ValidateConfig(config *Config) error {
if config.APIKey == "" {
return fmt.Errorf("API key is required for OpenRouter provider")
}
return nil
}
+313
View File
@@ -0,0 +1,313 @@
// Package provider defines the unified interface and registry for multi-vendor model API adapters.
package provider
import (
"fmt"
"strings"
"sync"
"github.com/Tencent/WeKnora/internal/types"
)
// ProviderName 模型服务商名称
type ProviderName string
const (
// OpenAI
ProviderOpenAI ProviderName = "openai"
// Anthropic Claude
ProviderAnthropic ProviderName = "anthropic"
// 阿里云 DashScope
ProviderAliyun ProviderName = "aliyun"
// 智谱AI (GLM 系列)
ProviderZhipu ProviderName = "zhipu"
// OpenRouter
ProviderOpenRouter ProviderName = "openrouter"
// ProviderLiteLLM is the LiteLLM self-hosted proxy (OpenAI-compatible gateway to 100+ providers).
ProviderLiteLLM ProviderName = "litellm"
// Requesty
ProviderRequesty ProviderName = "requesty"
// 硅基流动
ProviderSiliconFlow ProviderName = "siliconflow"
// Jina AI (Embedding and Rerank)
ProviderJina ProviderName = "jina"
// Generic 兼容OpenAI (自定义部署)
ProviderGeneric ProviderName = "generic"
// DeepSeek
ProviderDeepSeek ProviderName = "deepseek"
// Google Gemini
ProviderGemini ProviderName = "gemini"
// 火山引擎 Ark
ProviderVolcengine ProviderName = "volcengine"
// 腾讯混元
ProviderHunyuan ProviderName = "hunyuan"
// MiniMax
ProviderMiniMax ProviderName = "minimax"
// 小米 Mimo
ProviderMimo ProviderName = "mimo"
// GPUStack (私有化部署)
ProviderGPUStack ProviderName = "gpustack"
// 月之暗面 Moonshot (Kimi)
ProviderMoonshot ProviderName = "moonshot"
// 魔搭 ModelScope
ProviderModelScope ProviderName = "modelscope"
// 百度千帆
ProviderQianfan ProviderName = "qianfan"
// 七牛云
ProviderQiniu ProviderName = "qiniu"
// 美团 LongCat AI
ProviderLongCat ProviderName = "longcat"
// 腾讯云 LKEAP (知识引擎原子能力)
ProviderLKEAP ProviderName = "lkeap"
// NVIDIA
ProviderNvidia ProviderName = "nvidia"
// Novita AI
ProviderNovita ProviderName = "novita"
// Azure OpenAI
ProviderAzureOpenAI ProviderName = "azure_openai"
)
// AllProviders 返回所有注册的提供者名称
func AllProviders() []ProviderName {
return []ProviderName{
ProviderGeneric,
ProviderWeKnoraCloud,
ProviderAliyun,
ProviderZhipu,
ProviderVolcengine,
ProviderHunyuan,
ProviderSiliconFlow,
ProviderDeepSeek,
ProviderMiniMax,
ProviderMoonshot,
ProviderModelScope,
ProviderQianfan,
ProviderQiniu,
ProviderOpenAI,
ProviderAnthropic,
ProviderGemini,
ProviderOpenRouter,
ProviderLiteLLM,
ProviderRequesty,
ProviderJina,
ProviderMimo,
ProviderLongCat,
ProviderLKEAP,
ProviderGPUStack,
ProviderNvidia,
ProviderNovita,
ProviderAzureOpenAI,
}
}
// ProviderInfo 包含提供者的元数据
type ProviderInfo struct {
Name ProviderName // 提供者标识
DisplayName string // 可读名称
Description string // 提供者描述
DefaultURLs map[types.ModelType]string // 按模型类型区分的默认 BaseURL
ModelTypes []types.ModelType // 支持的模型类型
RequiresAuth bool // 是否需要 API key
ExtraFields []ExtraFieldConfig // 额外配置字段
}
// GetDefaultURL 获取指定模型类型的默认 URL
func (p ProviderInfo) GetDefaultURL(modelType types.ModelType) string {
if url, ok := p.DefaultURLs[modelType]; ok {
return url
}
// 回退到 Chat URL
if url, ok := p.DefaultURLs[types.ModelTypeKnowledgeQA]; ok {
return url
}
return ""
}
// ExtraFieldConfig 定义提供者的额外配置字段
type ExtraFieldConfig struct {
Key string `json:"key"`
Label string `json:"label"`
Type string `json:"type"` // "string", "number", "boolean", "select"
Required bool `json:"required"`
Default string `json:"default"`
Placeholder string `json:"placeholder"`
Options []struct {
Label string `json:"label"`
Value string `json:"value"`
} `json:"options,omitempty"`
}
// Config 表示模型提供者的配置
type Config struct {
Provider ProviderName `json:"provider"`
BaseURL string `json:"base_url"`
APIKey string `json:"api_key"`
ModelName string `json:"model_name"`
ModelID string `json:"model_id"`
Extra map[string]any `json:"extra,omitempty"`
}
type Provider interface {
// Info 返回服务商的元数据
Info() ProviderInfo
// ValidateConfig 验证服务商的配置
ValidateConfig(config *Config) error
}
// registry 存储所有注册的提供者
var (
registryMu sync.RWMutex
registry = make(map[ProviderName]Provider)
)
// Register 添加一个提供者到全局注册表
func Register(p Provider) {
registryMu.Lock()
defer registryMu.Unlock()
registry[p.Info().Name] = p
}
// Get 通过名称从注册表中获取提供者
func Get(name ProviderName) (Provider, bool) {
registryMu.RLock()
defer registryMu.RUnlock()
p, ok := registry[name]
return p, ok
}
// GetOrDefault 通过名称从注册表中获取提供者,如果未找到则返回默认提供者
func GetOrDefault(name ProviderName) Provider {
p, ok := Get(name)
if ok {
return p
}
// 如果未找到则返回默认提供者
p, _ = Get(ProviderGeneric)
return p
}
// List 返回所有注册的提供者(按 AllProviders 定义的顺序)
func List() []ProviderInfo {
registryMu.RLock()
defer registryMu.RUnlock()
result := make([]ProviderInfo, 0, len(registry))
for _, name := range AllProviders() {
if p, ok := registry[name]; ok {
result = append(result, p.Info())
}
}
return result
}
// ListByModelType 返回所有支持指定模型类型的提供者(按 AllProviders 定义的顺序)
func ListByModelType(modelType types.ModelType) []ProviderInfo {
registryMu.RLock()
defer registryMu.RUnlock()
result := make([]ProviderInfo, 0)
for _, name := range AllProviders() {
if p, ok := registry[name]; ok {
info := p.Info()
for _, t := range info.ModelTypes {
if t == modelType {
result = append(result, info)
break
}
}
}
}
return result
}
// DetectProvider 通过 BaseURL 检测服务商
func DetectProvider(baseURL string) ProviderName {
switch {
case containsAny(baseURL, "dashscope.aliyuncs.com"):
return ProviderAliyun
case containsAny(baseURL, "open.bigmodel.cn", "zhipu"):
return ProviderZhipu
case containsAny(baseURL, "openrouter.ai"):
return ProviderOpenRouter
// Hostname/path containing "litellm" (including the catalog placeholder
// your_litellm_proxy). Loopback URLs such as localhost:4000 stay generic
// because they are SSRF-blocked unless explicitly whitelisted.
case containsAny(baseURL, "litellm"):
return ProviderLiteLLM
case containsAny(baseURL, "router.requesty.ai", "requesty.ai"):
return ProviderRequesty
case containsAny(baseURL, "siliconflow.cn"):
return ProviderSiliconFlow
case containsAny(baseURL, "api.jina.ai"):
return ProviderJina
case containsAny(baseURL, "openai.azure.com"):
return ProviderAzureOpenAI
case containsAny(baseURL, "api.openai.com"):
return ProviderOpenAI
case containsAny(baseURL, "api.anthropic.com"):
return ProviderAnthropic
case containsAny(baseURL, "api.deepseek.com"):
return ProviderDeepSeek
case containsAny(baseURL, "generativelanguage.googleapis.com"):
return ProviderGemini
case containsAny(baseURL, "volces.com", "volcengine"):
return ProviderVolcengine
case containsAny(baseURL, "hunyuan.cloud.tencent.com"):
return ProviderHunyuan
case containsAny(baseURL, "minimax.io", "minimaxi.com"):
return ProviderMiniMax
case containsAny(baseURL, "xiaomimimo.com"):
return ProviderMimo
case containsAny(baseURL, "gpustack"):
return ProviderGPUStack
case containsAny(baseURL, "modelscope.cn"):
return ProviderModelScope
case containsAny(baseURL, "qiniuapi.com", "qiniu"):
return ProviderQiniu
case containsAny(baseURL, "moonshot.ai"):
return ProviderMoonshot
case containsAny(baseURL, "qianfan.baidubce.com", "baidubce.com"):
return ProviderQianfan
case containsAny(baseURL, "longcat.chat"):
return ProviderLongCat
case containsAny(baseURL, "lkeap.cloud.tencent.com", "api.lkeap", "lkeap.tencentcloudapi.com"):
return ProviderLKEAP
case containsAny(baseURL, "nvidia.com"):
return ProviderNvidia
case containsAny(baseURL, "api.novita.ai", "novita.ai"):
return ProviderNovita
case containsAny(baseURL, "weknora.weixin.qq.com"):
return ProviderWeKnoraCloud
default:
return ProviderGeneric
}
}
func containsAny(s string, substrs ...string) bool {
for _, sub := range substrs {
if strings.Contains(s, sub) {
return true
}
}
return false
}
func NewConfigFromModel(model *types.Model) (*Config, error) {
if model == nil {
return nil, fmt.Errorf("model is nil")
}
providerName := ProviderName(model.Parameters.Provider)
if providerName == "" {
providerName = DetectProvider(model.Parameters.BaseURL)
}
return &Config{
Provider: providerName,
BaseURL: model.Parameters.BaseURL,
APIKey: model.Parameters.APIKey,
ModelName: model.Name,
ModelID: model.ID,
}, nil
}
+381
View File
@@ -0,0 +1,381 @@
package provider
import (
"testing"
"github.com/Tencent/WeKnora/internal/types"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestProviderRegistry(t *testing.T) {
// Test that all default providers are registered
t.Run("default providers registered", func(t *testing.T) {
providers := List()
assert.NotEmpty(t, providers, "should have registered providers")
// Check specific providers exist
for _, name := range []ProviderName{ProviderOpenAI, ProviderAliyun, ProviderZhipu, ProviderGeneric} {
p, ok := Get(name)
assert.True(t, ok, "provider %s should be registered", name)
assert.NotNil(t, p, "provider %s should not be nil", name)
}
})
t.Run("GetOrDefault fallback", func(t *testing.T) {
// Non-existent provider should fall back to generic
p := GetOrDefault("nonexistent")
require.NotNil(t, p)
assert.Equal(t, ProviderGeneric, p.Info().Name)
})
}
func TestDetectProvider(t *testing.T) {
tests := []struct {
url string
expected ProviderName
}{
{"https://api.openai.com/v1", ProviderOpenAI},
{"https://api.anthropic.com/v1", ProviderAnthropic},
{"https://openrouter.ai/api/v1", ProviderOpenRouter},
{"https://litellm.example.com/v1", ProviderLiteLLM},
{LiteLLMBaseURL, ProviderLiteLLM},
{"http://localhost:4000/v1", ProviderGeneric},
{"https://router.requesty.ai/v1", ProviderRequesty},
{"https://dashscope.aliyuncs.com/compatible-mode/v1", ProviderAliyun},
{"https://open.bigmodel.cn/api/paas/v4", ProviderZhipu},
{"https://api.deepseek.com/v1", ProviderDeepSeek},
{"https://generativelanguage.googleapis.com/v1beta/openai", ProviderGemini},
{"https://ark.cn-beijing.volces.com/api/v3", ProviderVolcengine},
{"https://api.hunyuan.cloud.tencent.com/v1", ProviderHunyuan},
{"https://api.minimaxi.com/v1", ProviderMiniMax},
{"https://api.minimax.io/v1", ProviderMiniMax},
{"https://api.xiaomimimo.com/v1", ProviderMimo},
{"https://custom-endpoint.example.com/v1", ProviderGeneric},
{"http://localhost:11434/v1", ProviderGeneric},
{"https://integrate.api.nvidia.com/v1", ProviderNvidia},
{"https://ai.api.nvidia.com/v1/retrieval/nvidia/reranking", ProviderNvidia},
}
for _, tt := range tests {
t.Run(tt.url, func(t *testing.T) {
result := DetectProvider(tt.url)
assert.Equal(t, tt.expected, result)
})
}
}
func TestAnthropicProviderValidation(t *testing.T) {
p := &AnthropicProvider{}
t.Run("valid config", func(t *testing.T) {
config := &Config{
APIKey: "sk-ant-test",
ModelName: "claude-sonnet-4-5",
}
err := p.ValidateConfig(config)
assert.NoError(t, err)
})
t.Run("missing API key", func(t *testing.T) {
config := &Config{
ModelName: "claude-sonnet-4-5",
}
err := p.ValidateConfig(config)
assert.Error(t, err)
assert.Contains(t, err.Error(), "API key")
})
t.Run("info", func(t *testing.T) {
info := p.Info()
assert.Equal(t, ProviderAnthropic, info.Name)
assert.Equal(t, AnthropicBaseURL, info.GetDefaultURL(types.ModelTypeKnowledgeQA))
assert.Contains(t, info.ModelTypes, types.ModelTypeKnowledgeQA)
assert.True(t, info.RequiresAuth)
})
}
func TestOpenAIProviderValidation(t *testing.T) {
p := &OpenAIProvider{}
t.Run("valid config", func(t *testing.T) {
config := &Config{
APIKey: "sk-test",
ModelName: "gpt-4",
}
err := p.ValidateConfig(config)
assert.NoError(t, err)
})
t.Run("missing API key", func(t *testing.T) {
config := &Config{
ModelName: "gpt-4",
}
err := p.ValidateConfig(config)
assert.Error(t, err)
assert.Contains(t, err.Error(), "API key")
})
t.Run("missing model name", func(t *testing.T) {
config := &Config{
APIKey: "sk-test",
}
err := p.ValidateConfig(config)
assert.Error(t, err)
assert.Contains(t, err.Error(), "model name")
})
}
func TestAliyunProviderValidation(t *testing.T) {
p := &AliyunProvider{}
t.Run("valid config", func(t *testing.T) {
config := &Config{
APIKey: "sk-test",
ModelName: "qwen-max",
}
err := p.ValidateConfig(config)
assert.NoError(t, err)
})
t.Run("info", func(t *testing.T) {
info := p.Info()
assert.Equal(t, ProviderAliyun, info.Name)
assert.Contains(t, info.ModelTypes, types.ModelTypeKnowledgeQA)
assert.Contains(t, info.ModelTypes, types.ModelTypeEmbedding)
assert.Contains(t, info.ModelTypes, types.ModelTypeRerank)
})
}
func TestAliyunModelDetection(t *testing.T) {
t.Run("Qwen3 model detection", func(t *testing.T) {
assert.True(t, IsQwen3Model("qwen3-32b"))
assert.True(t, IsQwen3Model("qwen3-72b"))
assert.False(t, IsQwen3Model("qwen-max"))
assert.False(t, IsQwen3Model("qwen2.5-72b"))
})
t.Run("DeepSeek model detection", func(t *testing.T) {
assert.True(t, IsDeepSeekModel("deepseek-chat"))
assert.True(t, IsDeepSeekModel("deepseek-v3.1"))
assert.True(t, IsDeepSeekModel("DeepSeek-Chat"))
assert.False(t, IsDeepSeekModel("qwen-max"))
})
}
func TestMiniMaxProviderValidation(t *testing.T) {
p := &MiniMaxProvider{}
t.Run("valid config", func(t *testing.T) {
config := &Config{
APIKey: "test-key",
ModelName: "MiniMax-M2.7",
}
err := p.ValidateConfig(config)
assert.NoError(t, err)
})
t.Run("missing API key", func(t *testing.T) {
config := &Config{
ModelName: "MiniMax-M2.7",
}
err := p.ValidateConfig(config)
assert.Error(t, err)
assert.Contains(t, err.Error(), "API key")
})
t.Run("missing model name", func(t *testing.T) {
config := &Config{
APIKey: "test-key",
}
err := p.ValidateConfig(config)
assert.Error(t, err)
assert.Contains(t, err.Error(), "model name")
})
t.Run("info", func(t *testing.T) {
info := p.Info()
assert.Equal(t, ProviderMiniMax, info.Name)
assert.Equal(t, "MiniMax", info.DisplayName)
assert.Contains(t, info.ModelTypes, types.ModelTypeKnowledgeQA)
assert.True(t, info.RequiresAuth)
assert.Contains(t, info.Description, "M2.7")
})
}
func TestZhipuProviderValidation(t *testing.T) {
p := &ZhipuProvider{}
t.Run("valid config", func(t *testing.T) {
config := &Config{
APIKey: "test-key",
ModelName: "glm-4",
}
err := p.ValidateConfig(config)
assert.NoError(t, err)
})
t.Run("info", func(t *testing.T) {
info := p.Info()
assert.Equal(t, ProviderZhipu, info.Name)
assert.Equal(t, ZhipuChatBaseURL, info.GetDefaultURL(types.ModelTypeKnowledgeQA))
assert.Equal(t, ZhipuEmbeddingBaseURL, info.GetDefaultURL(types.ModelTypeEmbedding))
})
}
func TestRequestyProviderValidation(t *testing.T) {
p := &RequestyProvider{}
t.Run("valid config", func(t *testing.T) {
config := &Config{
APIKey: "test-key",
ModelName: "openai/gpt-4o-mini",
}
err := p.ValidateConfig(config)
assert.NoError(t, err)
})
t.Run("missing API key", func(t *testing.T) {
config := &Config{
ModelName: "openai/gpt-4o-mini",
}
err := p.ValidateConfig(config)
assert.Error(t, err)
assert.Contains(t, err.Error(), "API key")
})
t.Run("info", func(t *testing.T) {
info := p.Info()
assert.Equal(t, ProviderRequesty, info.Name)
assert.Equal(t, "Requesty", info.DisplayName)
assert.Equal(t, RequestyBaseURL, info.GetDefaultURL(types.ModelTypeKnowledgeQA))
assert.Equal(t, RequestyBaseURL, info.GetDefaultURL(types.ModelTypeEmbedding))
assert.Contains(t, info.ModelTypes, types.ModelTypeKnowledgeQA)
assert.True(t, info.RequiresAuth)
})
}
func TestListByModelType(t *testing.T) {
t.Run("chat models", func(t *testing.T) {
providers := ListByModelType(types.ModelTypeKnowledgeQA)
assert.NotEmpty(t, providers)
// Multiple providers support chat
assert.GreaterOrEqual(t, len(providers), 9)
})
t.Run("rerank models", func(t *testing.T) {
providers := ListByModelType(types.ModelTypeRerank)
assert.NotEmpty(t, providers)
// Check that Aliyun supports rerank
foundAliyun := false
foundLKEAP := false
foundVolcengine := false
for _, p := range providers {
if p.Name == ProviderAliyun {
foundAliyun = true
}
if p.Name == ProviderLKEAP {
foundLKEAP = true
assert.Equal(t, LKEAPRerankBaseURL, p.GetDefaultURL(types.ModelTypeRerank))
}
if p.Name == ProviderVolcengine {
foundVolcengine = true
assert.Equal(t, VolcengineRerankBaseURL, p.GetDefaultURL(types.ModelTypeRerank))
}
}
assert.True(t, foundAliyun, "Aliyun should support rerank")
assert.True(t, foundLKEAP, "LKEAP should support rerank")
assert.True(t, foundVolcengine, "Volcengine should support rerank")
})
t.Run("embedding models include openrouter", func(t *testing.T) {
providers := ListByModelType(types.ModelTypeEmbedding)
assert.NotEmpty(t, providers)
found := false
for _, p := range providers {
if p.Name == ProviderOpenRouter {
found = true
assert.Equal(t, OpenRouterBaseURL, p.GetDefaultURL(types.ModelTypeEmbedding))
break
}
}
assert.True(t, found, "OpenRouter should support embedding")
})
t.Run("embedding models include gemini", func(t *testing.T) {
providers := ListByModelType(types.ModelTypeEmbedding)
assert.NotEmpty(t, providers)
found := false
for _, p := range providers {
if p.Name == ProviderGemini {
found = true
assert.Equal(t, GeminiBaseURL, p.GetDefaultURL(types.ModelTypeEmbedding))
break
}
}
assert.True(t, found, "Gemini should support embedding via the native Gemini API")
})
}
func TestLiteLLMProviderValidation(t *testing.T) {
p := &LiteLLMProvider{}
t.Run("valid config", func(t *testing.T) {
config := &Config{
APIKey: "test-key",
ModelName: "gpt-4.1-mini",
}
err := p.ValidateConfig(config)
assert.NoError(t, err)
})
t.Run("missing API key", func(t *testing.T) {
config := &Config{
ModelName: "gpt-4.1-mini",
}
err := p.ValidateConfig(config)
assert.Error(t, err)
assert.Contains(t, err.Error(), "API key")
})
t.Run("info", func(t *testing.T) {
info := p.Info()
assert.Equal(t, ProviderLiteLLM, info.Name)
assert.Equal(t, "LiteLLM", info.DisplayName)
assert.True(t, info.RequiresAuth)
assert.Equal(t, LiteLLMBaseURL, info.DefaultURLs[types.ModelTypeKnowledgeQA])
assert.Equal(t, LiteLLMBaseURL, info.DefaultURLs[types.ModelTypeEmbedding])
assert.Equal(t, LiteLLMBaseURL, info.DefaultURLs[types.ModelTypeVLLM])
assert.Contains(t, info.ModelTypes, types.ModelTypeKnowledgeQA)
assert.Contains(t, info.ModelTypes, types.ModelTypeEmbedding)
assert.Contains(t, info.ModelTypes, types.ModelTypeVLLM)
})
t.Run("registered and listed", func(t *testing.T) {
got, ok := Get(ProviderLiteLLM)
require.True(t, ok)
require.NotNil(t, got)
foundChat := false
for _, info := range ListByModelType(types.ModelTypeKnowledgeQA) {
if info.Name == ProviderLiteLLM {
foundChat = true
break
}
}
assert.True(t, foundChat, "LiteLLM should appear for chat models")
foundEmbed := false
for _, info := range ListByModelType(types.ModelTypeEmbedding) {
if info.Name == ProviderLiteLLM {
foundEmbed = true
break
}
}
assert.True(t, foundEmbed, "LiteLLM should appear for embedding models")
})
}
+54
View File
@@ -0,0 +1,54 @@
package provider
import (
"fmt"
"github.com/Tencent/WeKnora/internal/types"
)
const (
QianfanBaseURL = "https://qianfan.baidubce.com/v2"
)
// QianfanProvider 实现百度千帆的 Provider 接口
type QianfanProvider struct{}
func init() {
Register(&QianfanProvider{})
}
// Info 返回百度千帆 provider 的元数据
func (p *QianfanProvider) Info() ProviderInfo {
return ProviderInfo{
Name: ProviderQianfan,
DisplayName: "百度千帆 Baidu Cloud",
Description: "ernie-5.0-thinking-preview, embedding-v1, bce-reranker-base, etc.",
DefaultURLs: map[types.ModelType]string{
types.ModelTypeKnowledgeQA: QianfanBaseURL,
types.ModelTypeEmbedding: QianfanBaseURL,
types.ModelTypeRerank: QianfanBaseURL,
types.ModelTypeVLLM: QianfanBaseURL,
},
ModelTypes: []types.ModelType{
types.ModelTypeKnowledgeQA,
types.ModelTypeEmbedding,
types.ModelTypeRerank,
types.ModelTypeVLLM,
},
RequiresAuth: true,
}
}
// ValidateConfig 验证百度千帆 provider 配置
func (p *QianfanProvider) ValidateConfig(config *Config) error {
if config.BaseURL == "" {
return fmt.Errorf("base URL is required for Qianfan provider")
}
if config.APIKey == "" {
return fmt.Errorf("API key is required for Qianfan provider")
}
if config.ModelName == "" {
return fmt.Errorf("model name is required")
}
return nil
}
+49
View File
@@ -0,0 +1,49 @@
package provider
import (
"fmt"
"github.com/Tencent/WeKnora/internal/types"
)
const (
// QiniuBaseURL 七牛云 API BaseURL (OpenAI 兼容模式)
QiniuBaseURL = "https://api.qnaigc.com/v1"
)
// QiniuProvider 实现七牛云的 Provider 接口
type QiniuProvider struct{}
func init() {
Register(&QiniuProvider{})
}
// Info 返回七牛云 provider 的元数据
func (p *QiniuProvider) Info() ProviderInfo {
return ProviderInfo{
Name: ProviderQiniu,
DisplayName: "七牛云 Qiniu",
Description: "deepseek/deepseek-v3.2-251201, z-ai/glm-4.7, etc.",
DefaultURLs: map[types.ModelType]string{
types.ModelTypeKnowledgeQA: QiniuBaseURL,
},
ModelTypes: []types.ModelType{
types.ModelTypeKnowledgeQA,
},
RequiresAuth: true,
}
}
// ValidateConfig 验证七牛云 provider 配置
func (p *QiniuProvider) ValidateConfig(config *Config) error {
if config.BaseURL == "" {
return fmt.Errorf("base URL is required for Qiniu provider")
}
if config.APIKey == "" {
return fmt.Errorf("API key is required for Qiniu provider")
}
if config.ModelName == "" {
return fmt.Errorf("model name is required")
}
return nil
}
+46
View File
@@ -0,0 +1,46 @@
package provider
import (
"fmt"
"github.com/Tencent/WeKnora/internal/types"
)
const (
RequestyBaseURL = "https://router.requesty.ai/v1"
)
// RequestyProvider 实现 Requesty 的 Provider 接口
type RequestyProvider struct{}
func init() {
Register(&RequestyProvider{})
}
// Info 返回 Requesty provider 的元数据
func (p *RequestyProvider) Info() ProviderInfo {
return ProviderInfo{
Name: ProviderRequesty,
DisplayName: "Requesty",
Description: "openai/gpt-4o-mini, anthropic/claude-sonnet-4-5, etc.",
DefaultURLs: map[types.ModelType]string{
types.ModelTypeKnowledgeQA: RequestyBaseURL,
types.ModelTypeEmbedding: RequestyBaseURL,
types.ModelTypeVLLM: RequestyBaseURL,
},
ModelTypes: []types.ModelType{
types.ModelTypeKnowledgeQA,
types.ModelTypeEmbedding,
types.ModelTypeVLLM,
},
RequiresAuth: true,
}
}
// ValidateConfig 验证 Requesty provider 配置
func (p *RequestyProvider) ValidateConfig(config *Config) error {
if config.APIKey == "" {
return fmt.Errorf("API key is required for Requesty provider")
}
return nil
}
+50
View File
@@ -0,0 +1,50 @@
package provider
import (
"fmt"
"github.com/Tencent/WeKnora/internal/types"
)
const (
SiliconFlowBaseURL = "https://api.siliconflow.cn/v1"
)
// SiliconFlowProvider 实现硅基流动的 Provider 接口
type SiliconFlowProvider struct{}
func init() {
Register(&SiliconFlowProvider{})
}
// Info 返回硅基流动 provider 的元数据
func (p *SiliconFlowProvider) Info() ProviderInfo {
return ProviderInfo{
Name: ProviderSiliconFlow,
DisplayName: "硅基流动 SiliconFlow",
Description: "deepseek-ai/DeepSeek-V3.1, etc.",
DefaultURLs: map[types.ModelType]string{
types.ModelTypeKnowledgeQA: SiliconFlowBaseURL,
types.ModelTypeEmbedding: SiliconFlowBaseURL,
types.ModelTypeRerank: SiliconFlowBaseURL,
types.ModelTypeVLLM: SiliconFlowBaseURL,
types.ModelTypeASR: SiliconFlowBaseURL,
},
ModelTypes: []types.ModelType{
types.ModelTypeKnowledgeQA,
types.ModelTypeEmbedding,
types.ModelTypeRerank,
types.ModelTypeVLLM,
types.ModelTypeASR,
},
RequiresAuth: true,
}
}
// ValidateConfig 验证硅基流动 provider 配置
func (p *SiliconFlowProvider) ValidateConfig(config *Config) error {
if config.APIKey == "" {
return fmt.Errorf("API key is required for SiliconFlow provider")
}
return nil
}
+56
View File
@@ -0,0 +1,56 @@
package provider
import (
"fmt"
"github.com/Tencent/WeKnora/internal/types"
)
const (
// VolcengineChatBaseURL 火山引擎 Ark Chat API BaseURL (OpenAI 兼容模式)
VolcengineChatBaseURL = "https://ark.cn-beijing.volces.com/api/v3"
// VolcengineEmbeddingBaseURL 火山引擎 Ark Multimodal Embedding API BaseURL
VolcengineEmbeddingBaseURL = "https://ark.cn-beijing.volces.com/api/v3/embeddings/multimodal"
// VolcengineRerankBaseURL 火山引擎知识库托管 Rerank API BaseURL
VolcengineRerankBaseURL = "https://api-knowledgebase.mlp.cn-beijing.volces.com"
)
// VolcengineProvider 实现火山引擎 Ark 的 Provider 接口
type VolcengineProvider struct{}
func init() {
Register(&VolcengineProvider{})
}
// Info 返回火山引擎 provider 的元数据
func (p *VolcengineProvider) Info() ProviderInfo {
return ProviderInfo{
Name: ProviderVolcengine,
DisplayName: "火山引擎 Volcengine",
Description: "doubao-1-5-pro-32k-250115, doubao-embedding-vision-250615, doubao-seed-rerank, etc.",
DefaultURLs: map[types.ModelType]string{
types.ModelTypeKnowledgeQA: VolcengineChatBaseURL,
types.ModelTypeEmbedding: VolcengineEmbeddingBaseURL,
types.ModelTypeRerank: VolcengineRerankBaseURL,
types.ModelTypeVLLM: VolcengineChatBaseURL,
},
ModelTypes: []types.ModelType{
types.ModelTypeKnowledgeQA,
types.ModelTypeEmbedding,
types.ModelTypeRerank,
types.ModelTypeVLLM,
},
RequiresAuth: true,
}
}
// ValidateConfig 验证火山引擎 provider 配置
func (p *VolcengineProvider) ValidateConfig(config *Config) error {
if config.APIKey == "" {
return fmt.Errorf("API key is required for Volcengine Ark provider")
}
if config.ModelName == "" {
return fmt.Errorf("model name is required")
}
return nil
}

Some files were not shown because too many files have changed in this diff Show More