ben/coze-studio
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: support EMBEDDING_MAX_BATCH_SIZE (#311)

Browse Source
This commit is contained in:
N3ko
2025-07-30 15:31:54 +08:00
committed by GitHub
parent f93f26fc48
commit bb74272385
9 changed files with 85 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files

113
backend/infra/impl/embedding/http/http.go
View File

@@ -25,6 +25,7 @@ import (
"time"
opt "github.com/cloudwego/eino/components/embedding"
"github.com/coze-dev/coze-studio/backend/pkg/lang/slices"
"github.com/coze-dev/coze-studio/backend/infra/contract/embedding"
)
@@ -41,88 +42,90 @@ type embedResp struct {
Sparse []map[int]float64 `json:"sparse"`
}
func NewEmbedding(addr string, dims int64) (embedding.Embedder, error) {
func NewEmbedding(addr string, dims int64, batchSize int) (embedding.Embedder, error) {
cli := &http.Client{Timeout: time.Second * 30}
return &embedder{
cli: cli,
addr: addr,
dim: dims,
cli: cli,
addr: addr,
dim: dims,
batchSize: batchSize,
}, nil
}
type embedder struct {
cli *http.Client
addr string
dim int64
cli *http.Client
addr string
dim int64
batchSize int
}
func (e *embedder) EmbedStrings(ctx context.Context, texts []string, opts ...opt.Option) ([][]float64, error) {
rb, err := json.Marshal(&embedReq{
Texts: texts,
NeedSparse: false,
})
if err != nil {
return nil, err
dense := make([][]float64, 0, len(texts))
for _, part := range slices.Chunks(texts, e.batchSize) {
rb, err := json.Marshal(&embedReq{
Texts: part,
NeedSparse: false,
})
if err != nil {
return nil, err
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, e.addr+pathEmbed, bytes.NewReader(rb))
if err != nil {
return nil, err
}
req.Header.Set("Content-Type", "application/json; charset=utf-8")
resp, err := e.do(req)
if err != nil {
return nil, err
}
dense = append(dense, resp.Dense...)
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, e.addr+pathEmbed, bytes.NewReader(rb))
if err != nil {
return nil, err
}
req.Header.Set("Content-Type", "application/json; charset=utf-8")
resp, err := e.cli.Do(req)
if err != nil {
return nil, err
}
defer resp.Body.Close()
b, err := io.ReadAll(resp.Body)
if err != nil {
return nil, err
}
r := &embedResp{}
if err = json.Unmarshal(b, r); err != nil {
return nil, err
}
return r.Dense, nil
return dense, nil
}
func (e *embedder) EmbedStringsHybrid(ctx context.Context, texts []string, opts ...opt.Option) ([][]float64, []map[int]float64, error) {
rb, err := json.Marshal(&embedReq{
Texts: texts,
NeedSparse: true,
})
if err != nil {
return nil, nil, err
dense := make([][]float64, 0, len(texts))
sparse := make([]map[int]float64, 0, len(texts))
for _, part := range slices.Chunks(texts, e.batchSize) {
rb, err := json.Marshal(&embedReq{
Texts: part,
NeedSparse: true,
})
if err != nil {
return nil, nil, err
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, e.addr+pathEmbed, bytes.NewReader(rb))
if err != nil {
return nil, nil, err
}
req.Header.Set("Content-Type", "application/json; charset=utf-8")
resp, err := e.do(req)
if err != nil {
return nil, nil, err
}
dense = append(dense, resp.Dense...)
sparse = append(sparse, resp.Sparse...)
}
return dense, sparse, nil
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, e.addr+pathEmbed, bytes.NewReader(rb))
if err != nil {
return nil, nil, err
}
req.Header.Set("Content-Type", "application/json; charset=utf-8")
func (e *embedder) do(req *http.Request) (*embedResp, error) {
resp, err := e.cli.Do(req)
if err != nil {
return nil, nil, err
return nil, err
}
defer resp.Body.Close()
b, err := io.ReadAll(resp.Body)
if err != nil {
return nil, nil, err
return nil, err
}
r := &embedResp{}
if err = json.Unmarshal(b, r); err != nil {
return nil, nil, err
return nil, err
}
return r.Dense, r.Sparse, nil
return r, nil
}
func (e *embedder) Dimensions() int64 {

2
backend/infra/impl/embedding/http/http_test.go
View File

@@ -31,7 +31,7 @@ func TestHTTPEmbedding(t *testing.T) {
}
ctx := context.Background()
emb, err := NewEmbedding("http://127.0.0.1:6543", 1024)
emb, err := NewEmbedding("http://127.0.0.1:6543", 1024, 10)
assert.NoError(t, err)
texts := []string{
"hello",