feat(search): add search provider endpoints and data structures
- Introduced new endpoints for managing search providers, including GET, PUT, and POST methods for testing and updating providers. - Added data structures for search provider configuration, presets, and tool assignments to support the new functionality. - Enhanced the OpenAPI settings to accommodate the new search-related features.
This commit is contained in:
Max
parent
1efa87e50a
commit
c3040559e6
5 changed files with 1077 additions and 0 deletions
622
openapi/setting/search.go
Normal file
622
openapi/setting/search.go
Normal file
|
|
@ -0,0 +1,622 @@
|
||||||
|
package setting
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
_ "embed"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"net/http"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/gin-gonic/gin"
|
||||||
|
"github.com/yaoapp/yao/openapi/oauth/authorized"
|
||||||
|
oauthTypes "github.com/yaoapp/yao/openapi/oauth/types"
|
||||||
|
"github.com/yaoapp/yao/openapi/response"
|
||||||
|
"github.com/yaoapp/yao/setting"
|
||||||
|
"gopkg.in/yaml.v3"
|
||||||
|
)
|
||||||
|
|
||||||
|
//go:embed search_presets.yml
|
||||||
|
var searchPresetsYML []byte
|
||||||
|
|
||||||
|
var searchPresets []SearchProviderPreset
|
||||||
|
|
||||||
|
func init() {
|
||||||
|
if err := yaml.Unmarshal(searchPresetsYML, &searchPresets); err != nil {
|
||||||
|
searchPresets = nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func searchFindPreset(key string) *SearchProviderPreset {
|
||||||
|
for i := range searchPresets {
|
||||||
|
if searchPresets[i].Key == key {
|
||||||
|
return &searchPresets[i]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func searchScope(info *oauthTypes.AuthorizedInfo) setting.ScopeID {
|
||||||
|
if info.TeamID != "" {
|
||||||
|
return setting.ScopeID{Scope: setting.ScopeTeam, TeamID: info.TeamID}
|
||||||
|
}
|
||||||
|
return setting.ScopeID{Scope: setting.ScopeUser, UserID: info.UserID}
|
||||||
|
}
|
||||||
|
|
||||||
|
func searchProviderNS(key string) string {
|
||||||
|
return "search.providers." + key
|
||||||
|
}
|
||||||
|
|
||||||
|
const searchAssignmentNS = "search.tool_assignment"
|
||||||
|
|
||||||
|
func searchPasswordFields(preset *SearchProviderPreset) map[string]bool {
|
||||||
|
m := make(map[string]bool)
|
||||||
|
for _, f := range preset.Fields {
|
||||||
|
if f.Type == "password" {
|
||||||
|
m[f.Key] = true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return m
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// GET /setting/search
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
func handleSearchGet(c *gin.Context) {
|
||||||
|
info := authorized.GetInfo(c)
|
||||||
|
|
||||||
|
providers := make([]SearchProviderConfig, 0, len(searchPresets))
|
||||||
|
for _, preset := range searchPresets {
|
||||||
|
cfg := SearchProviderConfig{
|
||||||
|
PresetKey: preset.Key,
|
||||||
|
Enabled: false,
|
||||||
|
FieldValues: map[string]string{},
|
||||||
|
Status: "unconfigured",
|
||||||
|
}
|
||||||
|
|
||||||
|
if preset.IsCloud {
|
||||||
|
var cloudSaved map[string]interface{}
|
||||||
|
if setting.Global != nil {
|
||||||
|
cloudSaved, _ = setting.Global.GetMerged(info.UserID, info.TeamID, cloudNS)
|
||||||
|
}
|
||||||
|
if cloudSaved != nil {
|
||||||
|
if st, ok := cloudSaved["status"].(string); ok && st == "connected" {
|
||||||
|
cfg.Enabled = true
|
||||||
|
cfg.Status = "connected"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else if setting.Global != nil {
|
||||||
|
saved, _ := setting.Global.GetMerged(info.UserID, info.TeamID, searchProviderNS(preset.Key))
|
||||||
|
if saved != nil {
|
||||||
|
if v, ok := saved["enabled"].(bool); ok {
|
||||||
|
cfg.Enabled = v
|
||||||
|
}
|
||||||
|
if v, ok := saved["status"].(string); ok && v != "" {
|
||||||
|
cfg.Status = v
|
||||||
|
}
|
||||||
|
pwFields := searchPasswordFields(&preset)
|
||||||
|
if fv, ok := saved["field_values"].(map[string]interface{}); ok {
|
||||||
|
for k, v := range fv {
|
||||||
|
s, _ := v.(string)
|
||||||
|
if pwFields[k] && s != "" {
|
||||||
|
cfg.FieldValues[k] = cloudMaskKey(cloudDecrypt(s))
|
||||||
|
} else {
|
||||||
|
cfg.FieldValues[k] = s
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
providers = append(providers, cfg)
|
||||||
|
}
|
||||||
|
|
||||||
|
var assignment SearchToolAssignment
|
||||||
|
if setting.Global != nil {
|
||||||
|
saved, _ := setting.Global.GetMerged(info.UserID, info.TeamID, searchAssignmentNS)
|
||||||
|
if saved != nil {
|
||||||
|
if v, ok := saved["web_search"].(string); ok && v != "" {
|
||||||
|
assignment.WebSearch = &v
|
||||||
|
}
|
||||||
|
if v, ok := saved["web_scrape"].(string); ok && v != "" {
|
||||||
|
assignment.WebScrape = &v
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
response.RespondWithSuccess(c, http.StatusOK, SearchPageData{
|
||||||
|
Presets: searchPresets,
|
||||||
|
Providers: providers,
|
||||||
|
ToolAssignment: assignment,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// PUT /setting/search/providers/:key
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
func handleSearchProviderUpdate(c *gin.Context) {
|
||||||
|
if !guardOwner(c) {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
key := c.Param("key")
|
||||||
|
if key == "cloud" {
|
||||||
|
respondError(c, http.StatusBadRequest, "cloud provider is managed by cloud service settings")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
preset := searchFindPreset(key)
|
||||||
|
if preset == nil {
|
||||||
|
respondError(c, http.StatusBadRequest, fmt.Sprintf("unknown provider: %s", key))
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
var body struct {
|
||||||
|
FieldValues map[string]string `json:"field_values"`
|
||||||
|
}
|
||||||
|
if err := c.ShouldBindJSON(&body); err != nil {
|
||||||
|
respondError(c, http.StatusBadRequest, "invalid request body")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
if setting.Global == nil {
|
||||||
|
respondError(c, http.StatusInternalServerError, "setting registry not initialized")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
info := authorized.GetInfo(c)
|
||||||
|
scope := searchScope(info)
|
||||||
|
|
||||||
|
existing, _ := setting.Global.Get(scope, searchProviderNS(key))
|
||||||
|
m := make(map[string]interface{})
|
||||||
|
for k, v := range existing {
|
||||||
|
m[k] = v
|
||||||
|
}
|
||||||
|
|
||||||
|
validFields := make(map[string]bool)
|
||||||
|
for _, f := range preset.Fields {
|
||||||
|
validFields[f.Key] = true
|
||||||
|
}
|
||||||
|
|
||||||
|
pwFields := searchPasswordFields(preset)
|
||||||
|
existingFV := map[string]interface{}{}
|
||||||
|
if fv, ok := m["field_values"].(map[string]interface{}); ok {
|
||||||
|
existingFV = fv
|
||||||
|
}
|
||||||
|
|
||||||
|
newFV := make(map[string]interface{})
|
||||||
|
for k, v := range existingFV {
|
||||||
|
newFV[k] = v
|
||||||
|
}
|
||||||
|
|
||||||
|
for k, v := range body.FieldValues {
|
||||||
|
if !validFields[k] {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if pwFields[k] {
|
||||||
|
if v == "" {
|
||||||
|
continue // keep existing
|
||||||
|
}
|
||||||
|
newFV[k] = cloudEncrypt(v)
|
||||||
|
} else {
|
||||||
|
newFV[k] = v
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
m["field_values"] = newFV
|
||||||
|
if _, ok := m["enabled"]; !ok {
|
||||||
|
m["enabled"] = false
|
||||||
|
}
|
||||||
|
if _, ok := m["status"]; !ok {
|
||||||
|
m["status"] = "unconfigured"
|
||||||
|
}
|
||||||
|
|
||||||
|
if _, err := setting.Global.Set(scope, searchProviderNS(key), m); err != nil {
|
||||||
|
respondError(c, http.StatusInternalServerError, err.Error())
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
cfg := SearchProviderConfig{
|
||||||
|
PresetKey: key,
|
||||||
|
Enabled: false,
|
||||||
|
FieldValues: map[string]string{},
|
||||||
|
Status: "unconfigured",
|
||||||
|
}
|
||||||
|
if v, ok := m["enabled"].(bool); ok {
|
||||||
|
cfg.Enabled = v
|
||||||
|
}
|
||||||
|
if v, ok := m["status"].(string); ok && v != "" {
|
||||||
|
cfg.Status = v
|
||||||
|
}
|
||||||
|
if fv, ok := m["field_values"].(map[string]interface{}); ok {
|
||||||
|
for k, v := range fv {
|
||||||
|
s, _ := v.(string)
|
||||||
|
if pwFields[k] && s != "" {
|
||||||
|
cfg.FieldValues[k] = cloudMaskKey(cloudDecrypt(s))
|
||||||
|
} else {
|
||||||
|
cfg.FieldValues[k] = s
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
response.RespondWithSuccess(c, http.StatusOK, cfg)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// PUT /setting/search/providers/:key/toggle
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
func handleSearchProviderToggle(c *gin.Context) {
|
||||||
|
if !guardOwner(c) {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
key := c.Param("key")
|
||||||
|
if key == "cloud" {
|
||||||
|
respondError(c, http.StatusBadRequest, "cloud provider is managed by cloud service settings")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
preset := searchFindPreset(key)
|
||||||
|
if preset == nil {
|
||||||
|
respondError(c, http.StatusBadRequest, fmt.Sprintf("unknown provider: %s", key))
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
var body struct {
|
||||||
|
Enabled bool `json:"enabled"`
|
||||||
|
}
|
||||||
|
if err := c.ShouldBindJSON(&body); err != nil {
|
||||||
|
respondError(c, http.StatusBadRequest, "invalid request body")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
if setting.Global == nil {
|
||||||
|
respondError(c, http.StatusInternalServerError, "setting registry not initialized")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
info := authorized.GetInfo(c)
|
||||||
|
scope := searchScope(info)
|
||||||
|
|
||||||
|
existing, _ := setting.Global.Get(scope, searchProviderNS(key))
|
||||||
|
m := make(map[string]interface{})
|
||||||
|
for k, v := range existing {
|
||||||
|
m[k] = v
|
||||||
|
}
|
||||||
|
m["enabled"] = body.Enabled
|
||||||
|
|
||||||
|
if _, err := setting.Global.Set(scope, searchProviderNS(key), m); err != nil {
|
||||||
|
respondError(c, http.StatusInternalServerError, err.Error())
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// When disabling, clear tool_assignment references
|
||||||
|
if !body.Enabled {
|
||||||
|
assignData, _ := setting.Global.Get(scope, searchAssignmentNS)
|
||||||
|
if assignData != nil {
|
||||||
|
changed := false
|
||||||
|
if v, ok := assignData["web_search"].(string); ok && v == key {
|
||||||
|
assignData["web_search"] = ""
|
||||||
|
changed = true
|
||||||
|
}
|
||||||
|
if v, ok := assignData["web_scrape"].(string); ok && v == key {
|
||||||
|
assignData["web_scrape"] = ""
|
||||||
|
changed = true
|
||||||
|
}
|
||||||
|
if changed {
|
||||||
|
setting.Global.Set(scope, searchAssignmentNS, assignData)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
cfg := SearchProviderConfig{
|
||||||
|
PresetKey: key,
|
||||||
|
Enabled: body.Enabled,
|
||||||
|
FieldValues: map[string]string{},
|
||||||
|
Status: "unconfigured",
|
||||||
|
}
|
||||||
|
if v, ok := m["status"].(string); ok && v != "" {
|
||||||
|
cfg.Status = v
|
||||||
|
}
|
||||||
|
pwFields := searchPasswordFields(preset)
|
||||||
|
if fv, ok := m["field_values"].(map[string]interface{}); ok {
|
||||||
|
for k, v := range fv {
|
||||||
|
s, _ := v.(string)
|
||||||
|
if pwFields[k] && s != "" {
|
||||||
|
cfg.FieldValues[k] = cloudMaskKey(cloudDecrypt(s))
|
||||||
|
} else {
|
||||||
|
cfg.FieldValues[k] = s
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
response.RespondWithSuccess(c, http.StatusOK, cfg)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// POST /setting/search/providers/:key/test
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
func handleSearchProviderTest(c *gin.Context) {
|
||||||
|
if !guardOwner(c) {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
key := c.Param("key")
|
||||||
|
if key == "cloud" {
|
||||||
|
respondError(c, http.StatusBadRequest, "cloud provider status is determined by cloud service configuration")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
preset := searchFindPreset(key)
|
||||||
|
if preset == nil {
|
||||||
|
respondError(c, http.StatusBadRequest, fmt.Sprintf("unknown provider: %s", key))
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
var body struct {
|
||||||
|
FieldValues map[string]string `json:"field_values"`
|
||||||
|
}
|
||||||
|
c.ShouldBindJSON(&body)
|
||||||
|
|
||||||
|
info := authorized.GetInfo(c)
|
||||||
|
scope := searchScope(info)
|
||||||
|
|
||||||
|
// Resolve API key: prefer body, fall back to saved
|
||||||
|
apiKey := ""
|
||||||
|
if body.FieldValues != nil {
|
||||||
|
apiKey = body.FieldValues["api_key"]
|
||||||
|
}
|
||||||
|
if apiKey == "" && setting.Global != nil {
|
||||||
|
saved, _ := setting.Global.Get(scope, searchProviderNS(key))
|
||||||
|
if saved != nil {
|
||||||
|
if fv, ok := saved["field_values"].(map[string]interface{}); ok {
|
||||||
|
if v, ok := fv["api_key"].(string); ok {
|
||||||
|
apiKey = cloudDecrypt(v)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if apiKey == "" {
|
||||||
|
response.RespondWithSuccess(c, http.StatusOK, SearchTestResult{
|
||||||
|
Success: false,
|
||||||
|
Message: "API key is required",
|
||||||
|
})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
start := time.Now()
|
||||||
|
var testErr error
|
||||||
|
|
||||||
|
zone := ""
|
||||||
|
if body.FieldValues != nil {
|
||||||
|
zone = body.FieldValues["zone"]
|
||||||
|
}
|
||||||
|
if zone == "" && setting.Global != nil {
|
||||||
|
saved, _ := setting.Global.Get(scope, searchProviderNS(key))
|
||||||
|
if saved != nil {
|
||||||
|
if fv, ok := saved["field_values"].(map[string]interface{}); ok {
|
||||||
|
if v, ok := fv["zone"].(string); ok {
|
||||||
|
zone = v
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
switch key {
|
||||||
|
case "tavily":
|
||||||
|
testErr = searchTestTavily(apiKey)
|
||||||
|
case "serper":
|
||||||
|
testErr = searchTestSerper(apiKey)
|
||||||
|
case "brightdata":
|
||||||
|
testErr = searchTestBrightdata(apiKey, zone)
|
||||||
|
default:
|
||||||
|
respondError(c, http.StatusBadRequest, fmt.Sprintf("test not supported for provider: %s", key))
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
latency := time.Since(start).Milliseconds()
|
||||||
|
|
||||||
|
if testErr != nil {
|
||||||
|
// Update status to disconnected
|
||||||
|
if setting.Global != nil {
|
||||||
|
saved, _ := setting.Global.Get(scope, searchProviderNS(key))
|
||||||
|
if saved == nil {
|
||||||
|
saved = map[string]interface{}{}
|
||||||
|
}
|
||||||
|
saved["status"] = "disconnected"
|
||||||
|
setting.Global.Set(scope, searchProviderNS(key), saved)
|
||||||
|
}
|
||||||
|
response.RespondWithSuccess(c, http.StatusOK, SearchTestResult{
|
||||||
|
Success: false,
|
||||||
|
Message: testErr.Error(),
|
||||||
|
})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// Update status to connected
|
||||||
|
if setting.Global != nil {
|
||||||
|
saved, _ := setting.Global.Get(scope, searchProviderNS(key))
|
||||||
|
if saved == nil {
|
||||||
|
saved = map[string]interface{}{}
|
||||||
|
}
|
||||||
|
saved["status"] = "connected"
|
||||||
|
setting.Global.Set(scope, searchProviderNS(key), saved)
|
||||||
|
}
|
||||||
|
|
||||||
|
response.RespondWithSuccess(c, http.StatusOK, SearchTestResult{
|
||||||
|
Success: true,
|
||||||
|
Message: "Connection successful",
|
||||||
|
LatencyMs: latency,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func searchTestTavily(apiKey string) error {
|
||||||
|
payload, _ := json.Marshal(map[string]interface{}{
|
||||||
|
"api_key": apiKey,
|
||||||
|
"query": "test",
|
||||||
|
})
|
||||||
|
client := &http.Client{Timeout: 15 * time.Second}
|
||||||
|
resp, err := client.Post("https://api.tavily.com/search", "application/json", bytes.NewReader(payload))
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("connection failed: %s", err.Error())
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden {
|
||||||
|
return fmt.Errorf("invalid API key (HTTP %d)", resp.StatusCode)
|
||||||
|
}
|
||||||
|
if resp.StatusCode != http.StatusOK {
|
||||||
|
return fmt.Errorf("server returned HTTP %d", resp.StatusCode)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func searchTestSerper(apiKey string) error {
|
||||||
|
payload, _ := json.Marshal(map[string]string{"q": "test"})
|
||||||
|
client := &http.Client{Timeout: 15 * time.Second}
|
||||||
|
req, err := http.NewRequest("POST", "https://google.serper.dev/search", bytes.NewReader(payload))
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("failed to build request: %s", err.Error())
|
||||||
|
}
|
||||||
|
req.Header.Set("X-API-KEY", apiKey)
|
||||||
|
req.Header.Set("Content-Type", "application/json")
|
||||||
|
resp, err := client.Do(req)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("connection failed: %s", err.Error())
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden {
|
||||||
|
return fmt.Errorf("invalid API key (HTTP %d)", resp.StatusCode)
|
||||||
|
}
|
||||||
|
if resp.StatusCode != http.StatusOK {
|
||||||
|
return fmt.Errorf("server returned HTTP %d", resp.StatusCode)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func searchTestBrightdata(apiKey, zone string) error {
|
||||||
|
if zone == "" {
|
||||||
|
return fmt.Errorf("Zone is required")
|
||||||
|
}
|
||||||
|
client := &http.Client{Timeout: 10 * time.Second}
|
||||||
|
req, err := http.NewRequest("GET", "https://api.brightdata.com/zone/status?zone="+zone, nil)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("failed to build request: %s", err.Error())
|
||||||
|
}
|
||||||
|
req.Header.Set("Authorization", "Bearer "+apiKey)
|
||||||
|
resp, err := client.Do(req)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("connection failed: %s", err.Error())
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden {
|
||||||
|
return fmt.Errorf("invalid API key (HTTP %d)", resp.StatusCode)
|
||||||
|
}
|
||||||
|
if resp.StatusCode == http.StatusNotFound {
|
||||||
|
return fmt.Errorf("zone '%s' not found", zone)
|
||||||
|
}
|
||||||
|
if resp.StatusCode != http.StatusOK {
|
||||||
|
return fmt.Errorf("server returned HTTP %d", resp.StatusCode)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// PUT /setting/search/tool-assignment
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
func handleSearchToolAssignment(c *gin.Context) {
|
||||||
|
if !guardOwner(c) {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
var body struct {
|
||||||
|
WebSearch *string `json:"web_search"`
|
||||||
|
WebScrape *string `json:"web_scrape"`
|
||||||
|
}
|
||||||
|
if err := c.ShouldBindJSON(&body); err != nil {
|
||||||
|
respondError(c, http.StatusBadRequest, "invalid request body")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
if setting.Global == nil {
|
||||||
|
respondError(c, http.StatusInternalServerError, "setting registry not initialized")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
info := authorized.GetInfo(c)
|
||||||
|
scope := searchScope(info)
|
||||||
|
|
||||||
|
// Validate: provider must be enabled and support the tool
|
||||||
|
validateAssignment := func(providerKey *string, toolType string) error {
|
||||||
|
if providerKey == nil || *providerKey == "" {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
preset := searchFindPreset(*providerKey)
|
||||||
|
if preset == nil {
|
||||||
|
return fmt.Errorf("unknown provider: %s", *providerKey)
|
||||||
|
}
|
||||||
|
|
||||||
|
hasTools := false
|
||||||
|
for _, t := range preset.Tools {
|
||||||
|
if t == toolType {
|
||||||
|
hasTools = true
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !hasTools {
|
||||||
|
return fmt.Errorf("provider %s does not support %s", *providerKey, toolType)
|
||||||
|
}
|
||||||
|
|
||||||
|
if preset.IsCloud {
|
||||||
|
return nil // cloud provider enablement is implicit
|
||||||
|
}
|
||||||
|
|
||||||
|
saved, _ := setting.Global.Get(scope, searchProviderNS(*providerKey))
|
||||||
|
if saved != nil {
|
||||||
|
if v, ok := saved["enabled"].(bool); ok && v {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return fmt.Errorf("provider %s is not enabled", *providerKey)
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := validateAssignment(body.WebSearch, "web_search"); err != nil {
|
||||||
|
respondError(c, http.StatusBadRequest, err.Error())
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if err := validateAssignment(body.WebScrape, "web_scrape"); err != nil {
|
||||||
|
respondError(c, http.StatusBadRequest, err.Error())
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
m := make(map[string]interface{})
|
||||||
|
if body.WebSearch != nil {
|
||||||
|
m["web_search"] = *body.WebSearch
|
||||||
|
} else {
|
||||||
|
m["web_search"] = ""
|
||||||
|
}
|
||||||
|
if body.WebScrape != nil {
|
||||||
|
m["web_scrape"] = *body.WebScrape
|
||||||
|
} else {
|
||||||
|
m["web_scrape"] = ""
|
||||||
|
}
|
||||||
|
|
||||||
|
if _, err := setting.Global.Set(scope, searchAssignmentNS, m); err != nil {
|
||||||
|
respondError(c, http.StatusInternalServerError, err.Error())
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
result := SearchToolAssignment{}
|
||||||
|
if v, ok := m["web_search"].(string); ok && v != "" {
|
||||||
|
result.WebSearch = &v
|
||||||
|
}
|
||||||
|
if v, ok := m["web_scrape"].(string); ok && v != "" {
|
||||||
|
result.WebScrape = &v
|
||||||
|
}
|
||||||
|
|
||||||
|
response.RespondWithSuccess(c, http.StatusOK, result)
|
||||||
|
}
|
||||||
61
openapi/setting/search_presets.yml
Normal file
61
openapi/setting/search_presets.yml
Normal file
|
|
@ -0,0 +1,61 @@
|
||||||
|
# Search & Scrape provider presets.
|
||||||
|
# Embedded at compile time via go:embed in search.go.
|
||||||
|
|
||||||
|
- key: cloud
|
||||||
|
name: Yao Agents
|
||||||
|
description:
|
||||||
|
zh-CN: "云服务提供的搜索与抓取能力，凭证来自云服务配置页"
|
||||||
|
en-US: "Search & scrape capabilities from cloud service, credentials from cloud config"
|
||||||
|
website: "https://yaoagents.com"
|
||||||
|
tools: [web_search, web_scrape]
|
||||||
|
tool_labels:
|
||||||
|
- { zh-CN: "网页搜索", en-US: "Web Search" }
|
||||||
|
- { zh-CN: "网页抓取", en-US: "Web Scrape" }
|
||||||
|
fields: []
|
||||||
|
is_cloud: true
|
||||||
|
|
||||||
|
- key: tavily
|
||||||
|
name: Tavily
|
||||||
|
description:
|
||||||
|
zh-CN: "AI 优化的搜索 API，返回结构化结果，适合 Agent 使用"
|
||||||
|
en-US: "AI-optimized search API with structured results, ideal for agents"
|
||||||
|
website: "https://tavily.com"
|
||||||
|
tools: [web_search]
|
||||||
|
tool_labels:
|
||||||
|
- { zh-CN: "网页搜索", en-US: "Web Search" }
|
||||||
|
fields:
|
||||||
|
- key: api_key
|
||||||
|
label: { zh-CN: "API Key", en-US: "API Key" }
|
||||||
|
type: password
|
||||||
|
|
||||||
|
- key: serper
|
||||||
|
name: "Serper (Google)"
|
||||||
|
description:
|
||||||
|
zh-CN: "基于 Google 搜索的 API，价格实惠，结果质量高"
|
||||||
|
en-US: "Google Search API with affordable pricing and high-quality results"
|
||||||
|
website: "https://serper.dev"
|
||||||
|
tools: [web_search]
|
||||||
|
tool_labels:
|
||||||
|
- { zh-CN: "网页搜索", en-US: "Web Search" }
|
||||||
|
fields:
|
||||||
|
- key: api_key
|
||||||
|
label: { zh-CN: "API Key", en-US: "API Key" }
|
||||||
|
type: password
|
||||||
|
|
||||||
|
- key: brightdata
|
||||||
|
name: Brightdata
|
||||||
|
description:
|
||||||
|
zh-CN: "部分网站有访问限制，启用代理可提升抓取成功率。需开通 Web Access API (Web Unlocker)。"
|
||||||
|
en-US: "Some websites have access restrictions. Enabling proxy improves scraping success rate. Requires Web Access API (Web Unlocker)."
|
||||||
|
website: "https://brightdata.com"
|
||||||
|
tools: [web_scrape]
|
||||||
|
tool_labels:
|
||||||
|
- { zh-CN: "网页抓取", en-US: "Web Scrape" }
|
||||||
|
fields:
|
||||||
|
- key: api_key
|
||||||
|
label: { zh-CN: "API Key", en-US: "API Key" }
|
||||||
|
type: password
|
||||||
|
- key: zone
|
||||||
|
label: { zh-CN: "Zone", en-US: "Zone" }
|
||||||
|
type: text
|
||||||
|
hint: { zh-CN: "Web Unlocker API 的 Zone 名称", en-US: "Zone name of your Web Unlocker API" }
|
||||||
|
|
@ -45,6 +45,13 @@ func Attach(group *gin.RouterGroup, oauth oauthTypes.OAuth) {
|
||||||
llm.PUT("/providers/:key", handleLLMProviderUpdate)
|
llm.PUT("/providers/:key", handleLLMProviderUpdate)
|
||||||
llm.DELETE("/providers/:key", handleLLMProviderDelete)
|
llm.DELETE("/providers/:key", handleLLMProviderDelete)
|
||||||
llm.POST("/providers/:key/test", handleLLMProviderTest)
|
llm.POST("/providers/:key/test", handleLLMProviderTest)
|
||||||
|
|
||||||
|
search := group.Group("/search")
|
||||||
|
search.GET("", handleSearchGet)
|
||||||
|
search.PUT("/providers/:key", handleSearchProviderUpdate)
|
||||||
|
search.PUT("/providers/:key/toggle", handleSearchProviderToggle)
|
||||||
|
search.POST("/providers/:key/test", handleSearchProviderTest)
|
||||||
|
search.PUT("/tool-assignment", handleSearchToolAssignment)
|
||||||
}
|
}
|
||||||
|
|
||||||
// requireOwner checks that the current user is the team owner.
|
// requireOwner checks that the current user is the team owner.
|
||||||
|
|
|
||||||
|
|
@ -92,3 +92,51 @@ type LLMPageData struct {
|
||||||
Roles map[string]interface{} `json:"roles"`
|
Roles map[string]interface{} `json:"roles"`
|
||||||
PresetProviders []interface{} `json:"preset_providers"`
|
PresetProviders []interface{} `json:"preset_providers"`
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Search & Scrape
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
type SearchProviderField struct {
|
||||||
|
Key string `json:"key" yaml:"key"`
|
||||||
|
Label map[string]string `json:"label" yaml:"label"`
|
||||||
|
Type string `json:"type" yaml:"type"`
|
||||||
|
Default string `json:"default,omitempty" yaml:"default"`
|
||||||
|
Placeholder string `json:"placeholder,omitempty" yaml:"placeholder"`
|
||||||
|
Hint map[string]string `json:"hint,omitempty" yaml:"hint"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type SearchProviderPreset struct {
|
||||||
|
Key string `json:"key" yaml:"key"`
|
||||||
|
Name string `json:"name" yaml:"name"`
|
||||||
|
Description map[string]string `json:"description,omitempty" yaml:"description"`
|
||||||
|
Website string `json:"website,omitempty" yaml:"website"`
|
||||||
|
Tools []string `json:"tools" yaml:"tools"`
|
||||||
|
ToolLabels []map[string]string `json:"tool_labels" yaml:"tool_labels"`
|
||||||
|
Fields []SearchProviderField `json:"fields" yaml:"fields"`
|
||||||
|
IsCloud bool `json:"is_cloud,omitempty" yaml:"is_cloud"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type SearchProviderConfig struct {
|
||||||
|
PresetKey string `json:"preset_key"`
|
||||||
|
Enabled bool `json:"enabled"`
|
||||||
|
FieldValues map[string]string `json:"field_values"`
|
||||||
|
Status string `json:"status"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type SearchToolAssignment struct {
|
||||||
|
WebSearch *string `json:"web_search"`
|
||||||
|
WebScrape *string `json:"web_scrape"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type SearchPageData struct {
|
||||||
|
Presets []SearchProviderPreset `json:"presets"`
|
||||||
|
Providers []SearchProviderConfig `json:"providers"`
|
||||||
|
ToolAssignment SearchToolAssignment `json:"tool_assignment"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type SearchTestResult struct {
|
||||||
|
Success bool `json:"success"`
|
||||||
|
Message string `json:"message"`
|
||||||
|
LatencyMs int64 `json:"latency_ms,omitempty"`
|
||||||
|
}
|
||||||
|
|
|
||||||
339
openapi/tests/setting/search_test.go
Normal file
339
openapi/tests/setting/search_test.go
Normal file
|
|
@ -0,0 +1,339 @@
|
||||||
|
package setting_test
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"encoding/json"
|
||||||
|
"net/http"
|
||||||
|
"os"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/stretchr/testify/assert"
|
||||||
|
"github.com/yaoapp/yao/openapi/tests/testutils"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Functional tests (system:root token)
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
func TestSearchGet(t *testing.T) {
|
||||||
|
serverURL := testutils.Prepare(t)
|
||||||
|
defer testutils.Clean()
|
||||||
|
initSettingRegistry(t)
|
||||||
|
token := obtainToken(t, serverURL)
|
||||||
|
|
||||||
|
req, err := http.NewRequest("GET", serverURL+baseURL()+"/setting/search", nil)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
req.Header.Set("Authorization", "Bearer "+token)
|
||||||
|
|
||||||
|
resp, err := http.DefaultClient.Do(req)
|
||||||
|
if !assert.NoError(t, err) || !assert.NotNil(t, resp) {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
assert.Equal(t, http.StatusOK, resp.StatusCode)
|
||||||
|
|
||||||
|
var body map[string]interface{}
|
||||||
|
err = json.NewDecoder(resp.Body).Decode(&body)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
|
||||||
|
assert.Contains(t, body, "presets")
|
||||||
|
assert.Contains(t, body, "providers")
|
||||||
|
assert.Contains(t, body, "tool_assignment")
|
||||||
|
|
||||||
|
presets, ok := body["presets"].([]interface{})
|
||||||
|
assert.True(t, ok)
|
||||||
|
assert.Equal(t, 4, len(presets), "should have 4 presets: cloud, tavily, serper, brightdata")
|
||||||
|
|
||||||
|
providers, ok := body["providers"].([]interface{})
|
||||||
|
assert.True(t, ok)
|
||||||
|
assert.Equal(t, 4, len(providers), "should have 4 provider configs")
|
||||||
|
|
||||||
|
// Cloud provider should be first
|
||||||
|
first, _ := providers[0].(map[string]interface{})
|
||||||
|
assert.Equal(t, "cloud", first["preset_key"])
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSearchGetUnauthenticated(t *testing.T) {
|
||||||
|
serverURL := testutils.Prepare(t)
|
||||||
|
defer testutils.Clean()
|
||||||
|
|
||||||
|
req, err := http.NewRequest("GET", serverURL+baseURL()+"/setting/search", nil)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
|
||||||
|
resp, err := http.DefaultClient.Do(req)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
defer resp.Body.Close()
|
||||||
|
assert.Equal(t, http.StatusUnauthorized, resp.StatusCode)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSearchProviderUpdate(t *testing.T) {
|
||||||
|
serverURL := testutils.Prepare(t)
|
||||||
|
defer testutils.Clean()
|
||||||
|
initSettingRegistry(t)
|
||||||
|
token := obtainToken(t, serverURL)
|
||||||
|
|
||||||
|
payload := map[string]interface{}{
|
||||||
|
"field_values": map[string]string{
|
||||||
|
"api_key": "tvly-test-key-12345",
|
||||||
|
},
|
||||||
|
}
|
||||||
|
raw, _ := json.Marshal(payload)
|
||||||
|
req, err := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/providers/tavily", bytes.NewReader(raw))
|
||||||
|
assert.NoError(t, err)
|
||||||
|
req.Header.Set("Authorization", "Bearer "+token)
|
||||||
|
req.Header.Set("Content-Type", "application/json")
|
||||||
|
|
||||||
|
resp, err := http.DefaultClient.Do(req)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
defer resp.Body.Close()
|
||||||
|
assert.Equal(t, http.StatusOK, resp.StatusCode)
|
||||||
|
|
||||||
|
var body map[string]interface{}
|
||||||
|
json.NewDecoder(resp.Body).Decode(&body)
|
||||||
|
assert.Equal(t, "tavily", body["preset_key"])
|
||||||
|
|
||||||
|
// api_key should be masked in response
|
||||||
|
fv, _ := body["field_values"].(map[string]interface{})
|
||||||
|
maskedKey, _ := fv["api_key"].(string)
|
||||||
|
assert.True(t, strings.Contains(maskedKey, "..."), "api_key should be masked")
|
||||||
|
|
||||||
|
// GET should also return masked key
|
||||||
|
req2, _ := http.NewRequest("GET", serverURL+baseURL()+"/setting/search", nil)
|
||||||
|
req2.Header.Set("Authorization", "Bearer "+token)
|
||||||
|
resp2, err := http.DefaultClient.Do(req2)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
defer resp2.Body.Close()
|
||||||
|
|
||||||
|
var getData map[string]interface{}
|
||||||
|
json.NewDecoder(resp2.Body).Decode(&getData)
|
||||||
|
providers, _ := getData["providers"].([]interface{})
|
||||||
|
for _, p := range providers {
|
||||||
|
pm, _ := p.(map[string]interface{})
|
||||||
|
if pm["preset_key"] == "tavily" {
|
||||||
|
tfv, _ := pm["field_values"].(map[string]interface{})
|
||||||
|
assert.True(t, strings.Contains(tfv["api_key"].(string), "..."), "GET should return masked key")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSearchProviderUpdateCloud(t *testing.T) {
|
||||||
|
serverURL := testutils.Prepare(t)
|
||||||
|
defer testutils.Clean()
|
||||||
|
initSettingRegistry(t)
|
||||||
|
token := obtainToken(t, serverURL)
|
||||||
|
|
||||||
|
payload := map[string]interface{}{
|
||||||
|
"field_values": map[string]string{},
|
||||||
|
}
|
||||||
|
raw, _ := json.Marshal(payload)
|
||||||
|
req, err := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/providers/cloud", bytes.NewReader(raw))
|
||||||
|
assert.NoError(t, err)
|
||||||
|
req.Header.Set("Authorization", "Bearer "+token)
|
||||||
|
req.Header.Set("Content-Type", "application/json")
|
||||||
|
|
||||||
|
resp, err := http.DefaultClient.Do(req)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
defer resp.Body.Close()
|
||||||
|
assert.Equal(t, http.StatusBadRequest, resp.StatusCode, "cloud provider should be rejected")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSearchProviderToggle(t *testing.T) {
|
||||||
|
serverURL := testutils.Prepare(t)
|
||||||
|
defer testutils.Clean()
|
||||||
|
initSettingRegistry(t)
|
||||||
|
token := obtainToken(t, serverURL)
|
||||||
|
|
||||||
|
// Save tavily first
|
||||||
|
savePayload := map[string]interface{}{
|
||||||
|
"field_values": map[string]string{"api_key": "tvly-toggle-key"},
|
||||||
|
}
|
||||||
|
raw, _ := json.Marshal(savePayload)
|
||||||
|
req, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/providers/tavily", bytes.NewReader(raw))
|
||||||
|
req.Header.Set("Authorization", "Bearer "+token)
|
||||||
|
req.Header.Set("Content-Type", "application/json")
|
||||||
|
resp, err := http.DefaultClient.Do(req)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
resp.Body.Close()
|
||||||
|
|
||||||
|
// Enable tavily
|
||||||
|
enablePayload := map[string]interface{}{"enabled": true}
|
||||||
|
raw, _ = json.Marshal(enablePayload)
|
||||||
|
req2, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/providers/tavily/toggle", bytes.NewReader(raw))
|
||||||
|
req2.Header.Set("Authorization", "Bearer "+token)
|
||||||
|
req2.Header.Set("Content-Type", "application/json")
|
||||||
|
resp2, err := http.DefaultClient.Do(req2)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
defer resp2.Body.Close()
|
||||||
|
assert.Equal(t, http.StatusOK, resp2.StatusCode)
|
||||||
|
|
||||||
|
var body map[string]interface{}
|
||||||
|
json.NewDecoder(resp2.Body).Decode(&body)
|
||||||
|
assert.Equal(t, true, body["enabled"])
|
||||||
|
|
||||||
|
// Assign tavily to web_search
|
||||||
|
assignPayload := map[string]interface{}{"web_search": "tavily"}
|
||||||
|
raw, _ = json.Marshal(assignPayload)
|
||||||
|
req3, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/tool-assignment", bytes.NewReader(raw))
|
||||||
|
req3.Header.Set("Authorization", "Bearer "+token)
|
||||||
|
req3.Header.Set("Content-Type", "application/json")
|
||||||
|
resp3, err := http.DefaultClient.Do(req3)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
resp3.Body.Close()
|
||||||
|
|
||||||
|
// Disable tavily -- should clear tool_assignment
|
||||||
|
disablePayload := map[string]interface{}{"enabled": false}
|
||||||
|
raw, _ = json.Marshal(disablePayload)
|
||||||
|
req4, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/providers/tavily/toggle", bytes.NewReader(raw))
|
||||||
|
req4.Header.Set("Authorization", "Bearer "+token)
|
||||||
|
req4.Header.Set("Content-Type", "application/json")
|
||||||
|
resp4, err := http.DefaultClient.Do(req4)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
resp4.Body.Close()
|
||||||
|
|
||||||
|
// Verify tool_assignment cleared
|
||||||
|
req5, _ := http.NewRequest("GET", serverURL+baseURL()+"/setting/search", nil)
|
||||||
|
req5.Header.Set("Authorization", "Bearer "+token)
|
||||||
|
resp5, err := http.DefaultClient.Do(req5)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
defer resp5.Body.Close()
|
||||||
|
|
||||||
|
var getData map[string]interface{}
|
||||||
|
json.NewDecoder(resp5.Body).Decode(&getData)
|
||||||
|
ta, _ := getData["tool_assignment"].(map[string]interface{})
|
||||||
|
assert.Nil(t, ta["web_search"], "web_search should be cleared after disabling tavily")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSearchToolAssignment(t *testing.T) {
|
||||||
|
serverURL := testutils.Prepare(t)
|
||||||
|
defer testutils.Clean()
|
||||||
|
initSettingRegistry(t)
|
||||||
|
token := obtainToken(t, serverURL)
|
||||||
|
|
||||||
|
// Save and enable tavily
|
||||||
|
savePayload := map[string]interface{}{"field_values": map[string]string{"api_key": "tvly-assign-key"}}
|
||||||
|
raw, _ := json.Marshal(savePayload)
|
||||||
|
req, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/providers/tavily", bytes.NewReader(raw))
|
||||||
|
req.Header.Set("Authorization", "Bearer "+token)
|
||||||
|
req.Header.Set("Content-Type", "application/json")
|
||||||
|
resp, _ := http.DefaultClient.Do(req)
|
||||||
|
resp.Body.Close()
|
||||||
|
|
||||||
|
enablePayload := map[string]interface{}{"enabled": true}
|
||||||
|
raw, _ = json.Marshal(enablePayload)
|
||||||
|
req2, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/providers/tavily/toggle", bytes.NewReader(raw))
|
||||||
|
req2.Header.Set("Authorization", "Bearer "+token)
|
||||||
|
req2.Header.Set("Content-Type", "application/json")
|
||||||
|
resp2, _ := http.DefaultClient.Do(req2)
|
||||||
|
resp2.Body.Close()
|
||||||
|
|
||||||
|
// Assign tavily to web_search
|
||||||
|
assignPayload := map[string]interface{}{"web_search": "tavily"}
|
||||||
|
raw, _ = json.Marshal(assignPayload)
|
||||||
|
req3, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/tool-assignment", bytes.NewReader(raw))
|
||||||
|
req3.Header.Set("Authorization", "Bearer "+token)
|
||||||
|
req3.Header.Set("Content-Type", "application/json")
|
||||||
|
resp3, err := http.DefaultClient.Do(req3)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
defer resp3.Body.Close()
|
||||||
|
assert.Equal(t, http.StatusOK, resp3.StatusCode)
|
||||||
|
|
||||||
|
var body map[string]interface{}
|
||||||
|
json.NewDecoder(resp3.Body).Decode(&body)
|
||||||
|
assert.Equal(t, "tavily", body["web_search"])
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSearchToolAssignmentDisabledProvider(t *testing.T) {
|
||||||
|
serverURL := testutils.Prepare(t)
|
||||||
|
defer testutils.Clean()
|
||||||
|
initSettingRegistry(t)
|
||||||
|
token := obtainToken(t, serverURL)
|
||||||
|
|
||||||
|
// Try to assign a provider that isn't enabled
|
||||||
|
assignPayload := map[string]interface{}{"web_search": "serper"}
|
||||||
|
raw, _ := json.Marshal(assignPayload)
|
||||||
|
req, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/tool-assignment", bytes.NewReader(raw))
|
||||||
|
req.Header.Set("Authorization", "Bearer "+token)
|
||||||
|
req.Header.Set("Content-Type", "application/json")
|
||||||
|
resp, err := http.DefaultClient.Do(req)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
defer resp.Body.Close()
|
||||||
|
assert.Equal(t, http.StatusBadRequest, resp.StatusCode, "should reject assignment to disabled provider")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSearchProviderTest(t *testing.T) {
|
||||||
|
apiKey := os.Getenv("TAVILY_API_KEY")
|
||||||
|
if apiKey == "" {
|
||||||
|
t.Skip("TAVILY_API_KEY not set, skipping search provider test")
|
||||||
|
}
|
||||||
|
|
||||||
|
serverURL := testutils.Prepare(t)
|
||||||
|
defer testutils.Clean()
|
||||||
|
initSettingRegistry(t)
|
||||||
|
token := obtainToken(t, serverURL)
|
||||||
|
|
||||||
|
payload := map[string]interface{}{
|
||||||
|
"field_values": map[string]string{"api_key": apiKey},
|
||||||
|
}
|
||||||
|
raw, _ := json.Marshal(payload)
|
||||||
|
req, err := http.NewRequest("POST", serverURL+baseURL()+"/setting/search/providers/tavily/test", bytes.NewReader(raw))
|
||||||
|
assert.NoError(t, err)
|
||||||
|
req.Header.Set("Authorization", "Bearer "+token)
|
||||||
|
req.Header.Set("Content-Type", "application/json")
|
||||||
|
|
||||||
|
resp, err := http.DefaultClient.Do(req)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
defer resp.Body.Close()
|
||||||
|
assert.Equal(t, http.StatusOK, resp.StatusCode)
|
||||||
|
|
||||||
|
var body map[string]interface{}
|
||||||
|
json.NewDecoder(resp.Body).Decode(&body)
|
||||||
|
assert.Equal(t, true, body["success"])
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSearchProviderTestCloud(t *testing.T) {
|
||||||
|
serverURL := testutils.Prepare(t)
|
||||||
|
defer testutils.Clean()
|
||||||
|
initSettingRegistry(t)
|
||||||
|
token := obtainToken(t, serverURL)
|
||||||
|
|
||||||
|
req, _ := http.NewRequest("POST", serverURL+baseURL()+"/setting/search/providers/cloud/test", nil)
|
||||||
|
req.Header.Set("Authorization", "Bearer "+token)
|
||||||
|
resp, err := http.DefaultClient.Do(req)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
defer resp.Body.Close()
|
||||||
|
assert.Equal(t, http.StatusBadRequest, resp.StatusCode)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// ACL permission tests
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
func TestSearchACL_ReadOnlyScopeCannotWrite(t *testing.T) {
|
||||||
|
serverURL := testutils.Prepare(t)
|
||||||
|
defer testutils.Clean()
|
||||||
|
initSettingRegistry(t)
|
||||||
|
|
||||||
|
readToken := obtainRestrictedToken(t, serverURL, "setting:search:read:all")
|
||||||
|
|
||||||
|
// GET should work
|
||||||
|
req, _ := http.NewRequest("GET", serverURL+baseURL()+"/setting/search", nil)
|
||||||
|
req.Header.Set("Authorization", "Bearer "+readToken)
|
||||||
|
resp, err := http.DefaultClient.Do(req)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
defer resp.Body.Close()
|
||||||
|
assert.Equal(t, http.StatusOK, resp.StatusCode, "read-only scope should allow GET")
|
||||||
|
|
||||||
|
// PUT should be denied
|
||||||
|
payload := map[string]interface{}{
|
||||||
|
"field_values": map[string]string{"api_key": "tvly-acl-test"},
|
||||||
|
}
|
||||||
|
raw, _ := json.Marshal(payload)
|
||||||
|
req2, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/providers/tavily", bytes.NewReader(raw))
|
||||||
|
req2.Header.Set("Authorization", "Bearer "+readToken)
|
||||||
|
req2.Header.Set("Content-Type", "application/json")
|
||||||
|
resp2, err := http.DefaultClient.Do(req2)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
defer resp2.Body.Close()
|
||||||
|
assert.Equal(t, http.StatusForbidden, resp2.StatusCode, "read-only scope should deny PUT")
|
||||||
|
}
|
||||||
Loading…
Add table
Reference in a new issue