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
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:
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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,
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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"))
|
||||
}
|
||||
@@ -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{}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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`)
|
||||
})
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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"])
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
@@ -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)))
|
||||
}
|
||||
@@ -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"
|
||||
}
|
||||
}
|
||||
@@ -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"},
|
||||
}))
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -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},
|
||||
)
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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()
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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")
|
||||
})
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
Reference in New Issue
Block a user