feat: web search plugin support add usage to result

Author: zijiren233

core/model/batch.go

@@ -404,7 +404,7 @@ func updateGroupSummaryData(group string, tokenName string, modelName string, cr
 		groupSummary.MaxTPS = groupModelTokenRate.TPS
 	}
 
-	groupSummary.Usage.Add(&usage)
+	groupSummary.Usage.Add(usage)
 	if code != http.StatusOK {
 		groupSummary.ExceptionCount++
 	}
@@ -444,7 +444,7 @@ func updateSummaryData(channelID int, modelName string, createAt time.Time, code
 		summary.MaxTPS = channelModelRate.TPS
 	}
 
-	summary.Usage.Add(&usage)
+	summary.Usage.Add(usage)
 	if code != http.StatusOK {
 		summary.ExceptionCount++
 	}

core/model/log.go

@@ -117,10 +117,7 @@ type Usage struct {
 	WebSearchCount      ZeroNullInt64 `json:"web_search_count,omitempty"`
 }
 
-func (u *Usage) Add(other *Usage) {
-	if other == nil {
-		return
-	}
+func (u *Usage) Add(other Usage) {
 	u.InputTokens += other.InputTokens
 	u.ImageInputTokens += other.ImageInputTokens
 	u.OutputTokens += other.OutputTokens

core/relay/adaptor/interface.go

@@ -3,6 +3,7 @@ package adaptor
 import (
 	"encoding/json"
 	"errors"
+	"fmt"
 	"io"
 	"net/http"
 
@@ -50,6 +51,7 @@ type ConvertRequestResult struct {
 
 type Error interface {
 	json.Marshaler
+	error
 	StatusCode() int
 }
 
@@ -66,6 +68,10 @@ func (e ErrorImpl[T]) StatusCode() int {
 	return e.statusCode
 }
 
+func (e ErrorImpl[T]) Error() string {
+	return fmt.Sprintf("status code: %d, error: %v", e.statusCode, e.error)
+}
+
 var ErrGetBalanceNotImplemented = errors.New("get balance not implemented")
 
 type Balancer interface {

core/relay/controller/dohelper.go

@@ -46,15 +46,7 @@ func (rw *responseWriter) Write(b []byte) (int, error) {
 }
 
 func (rw *responseWriter) WriteString(s string) (int, error) {
-	if rw.firstByteAt.IsZero() {
-		rw.firstByteAt = time.Now()
-	}
-	if rw.body.Len()+len(s) <= maxBufferSize {
-		rw.body.WriteString(s)
-	} else {
-		rw.body.WriteString(s[:maxBufferSize-rw.body.Len()])
-	}
-	return rw.ResponseWriter.WriteString(s)
+	return rw.Write(conv.StringToBytes(s))
 }
 
 var bufferPool = sync.Pool{

core/relay/meta/meta.go

@@ -181,6 +181,15 @@ func (m *Meta) GetInt64(key string) int64 {
 	return i
 }
 
+func (m *Meta) GetInt(key string) int {
+	v, ok := m.Get(key)
+	if !ok {
+		return 0
+	}
+	i, _ := v.(int)
+	return i
+}
+
 func GetMappedModelName(modelName string, mapping map[string]string) (string, bool) {
 	if len(modelName) == 0 {
 		return modelName, false

core/relay/plugin/web-search/README.md

@@ -17,7 +17,9 @@ The Web Search Plugin is a plugin that provides real-time web search capabilitie
             "enable_plugin": true,
             "default_enable": true,
             "search_rewrite": {
-                "enable": true
+                "enable": true,
+                "add_rewrite_usage": true,
+                "rewrite_usage_field": "rewrite_usage"
             },
             "need_reference": true,
             "search_from": [
@@ -67,6 +69,8 @@ The Web Search Plugin is a plugin that provides real-time web search capabilitie
 | `model_name` | string | No | - | Model name for query rewriting, uses current request model if empty |
 | `timeout_millisecond` | uint32 | No | 10000 | Rewrite request timeout (milliseconds) |
 | `max_count` | int | No | 3 | Maximum number of rewritten queries |
+| `add_rewrite_usage` | bool | No | false | Whether to add rewrite usage statistics to the response |
+| `rewrite_usage_field` | string | No | "rewrite_usage" | Rewrite usage statistics field name |
 
 #### Search Engine Configuration (`search_from`)
 

core/relay/plugin/web-search/README.zh.md

@@ -17,7 +17,9 @@ Web Search Plugin 是一个为 AI 模型提供实时网络搜索能力的插件
             "enable_plugin": true,
             "default_enable": true,
             "search_rewrite": {
-                "enable": true
+                "enable": true,
+                "add_rewrite_usage": true,
+                "rewrite_usage_field": "rewrite_usage"
             },
             "need_reference": true,
             "search_from": [
@@ -67,6 +69,8 @@ Web Search Plugin 是一个为 AI 模型提供实时网络搜索能力的插件
 | `model_name` | string | 否 | - | 用于重写查询的模型名称，为空时使用当前请求的模型 |
 | `timeout_millisecond` | uint32 | 否 | 10000 | 重写请求超时时间（毫秒） |
 | `max_count` | int | 否 | 3 | 最大重写查询数量 |
+| `add_rewrite_usage` | bool | 否 | false | 是否在响应中添加重写使用统计信息 |
+| `rewrite_usage_field` | string | 否 | "rewrite_usage" | 重写使用统计信息字段名称 |
 
 #### 搜索引擎配置 (`search_from`)
 

core/relay/plugin/web-search/config.go

@@ -25,6 +25,8 @@ type SearchRewrite struct {
 	ModelName          string `json:"model_name"`
 	TimeoutMillisecond uint32 `json:"timeout_millisecond"`
 	MaxCount           int    `json:"max_count"`
+	AddRewriteUsage    bool   `json:"add_rewrite_usage"`
+	RewriteUsageField  string `json:"rewrite_usage_field"`
 }
 
 type EngineConfig struct {