WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

internal/infrastructure/web_search/zhipu.go

@@ -0,0 +1,271 @@
+package web_search
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"strings"
+	"time"
+	"unicode/utf8"
+
+	"github.com/Tencent/WeKnora/internal/logger"
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/Tencent/WeKnora/internal/types/interfaces"
+)
+
+const (
+	defaultZhipuSearchURL    = "https://open.bigmodel.cn/api/paas/v4/web_search"
+	defaultZhipuTimeout      = 15 * time.Second
+	defaultZhipuResults      = 10
+	maxZhipuResults          = 50
+	maxZhipuQueryRunes       = 70
+	maxZhipuResponseBytes    = 2 << 20
+	defaultZhipuSearchEngine = "search_std"
+	defaultZhipuContentSize  = "medium"
+)
+
+var validZhipuSearchEngines = map[string]struct{}{
+	"search_std":       {},
+	"search_pro":       {},
+	"search_pro_sogou": {},
+	"search_pro_quark": {},
+}
+
+var validZhipuContentSizes = map[string]struct{}{
+	"medium": {},
+	"high":   {},
+}
+
+// ZhipuProvider implements web search using Zhipu AI's standalone Web Search API.
+type ZhipuProvider struct {
+	client       *http.Client
+	baseURL      string
+	apiKey       string
+	searchEngine string
+	contentSize  string
+}
+
+// NewZhipuProvider creates a Zhipu AI web search provider from persisted parameters.
+func NewZhipuProvider(params types.WebSearchProviderParameters) (interfaces.WebSearchProvider, error) {
+	if err := ValidateZhipuParameters(params); err != nil {
+		return nil, err
+	}
+	client, err := NewSearchHTTPClient(defaultZhipuTimeout, params.ProxyURL)
+	if err != nil {
+		return nil, err
+	}
+	searchEngine, contentSize := zhipuOptions(params.ExtraConfig)
+	return &ZhipuProvider{
+		client:       client,
+		baseURL:      defaultZhipuSearchURL,
+		apiKey:       strings.TrimSpace(params.APIKey),
+		searchEngine: searchEngine,
+		contentSize:  contentSize,
+	}, nil
+}
+
+// ValidateZhipuParameters validates credentials and provider-specific options.
+func ValidateZhipuParameters(params types.WebSearchProviderParameters) error {
+	if strings.TrimSpace(params.APIKey) == "" {
+		return fmt.Errorf("API key is required for Zhipu provider")
+	}
+	searchEngine, contentSize := zhipuOptions(params.ExtraConfig)
+	if _, ok := validZhipuSearchEngines[searchEngine]; !ok {
+		return fmt.Errorf("invalid Zhipu search engine: %s", searchEngine)
+	}
+	if _, ok := validZhipuContentSizes[contentSize]; !ok {
+		return fmt.Errorf("invalid Zhipu content size: %s", contentSize)
+	}
+	return nil
+}
+
+func zhipuOptions(extraConfig map[string]string) (searchEngine, contentSize string) {
+	searchEngine = defaultZhipuSearchEngine
+	contentSize = defaultZhipuContentSize
+	if value := strings.TrimSpace(extraConfig["search_engine"]); value != "" {
+		searchEngine = value
+	}
+	if value := strings.TrimSpace(extraConfig["content_size"]); value != "" {
+		contentSize = value
+	}
+	return searchEngine, contentSize
+}
+
+// Name returns the provider name.
+func (p *ZhipuProvider) Name() string {
+	return "zhipu"
+}
+
+// Search performs a web search using Zhipu AI's standalone Web Search API.
+func (p *ZhipuProvider) Search(
+	ctx context.Context,
+	query string,
+	maxResults int,
+	includeDate bool,
+) ([]*types.WebSearchResult, error) {
+	preparedQuery := normalizeZhipuQuery(query)
+	if preparedQuery == "" {
+		return nil, fmt.Errorf("query is empty")
+	}
+	if utf8.RuneCountInString(strings.TrimSpace(query)) > maxZhipuQueryRunes {
+		logger.Infof(ctx, "[WebSearch][Zhipu] truncated query to %d characters", maxZhipuQueryRunes)
+	}
+	if maxResults <= 0 {
+		maxResults = defaultZhipuResults
+	}
+	if maxResults > maxZhipuResults {
+		maxResults = maxZhipuResults
+	}
+
+	requestBody := zhipuSearchRequest{
+		SearchQuery:  preparedQuery,
+		SearchEngine: p.searchEngine,
+		SearchIntent: false,
+		Count:        maxResults,
+		ContentSize:  p.contentSize,
+	}
+	body, err := json.Marshal(requestBody)
+	if err != nil {
+		return nil, fmt.Errorf("failed to marshal Zhipu request: %w", err)
+	}
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, p.baseURL, bytes.NewReader(body))
+	if err != nil {
+		return nil, fmt.Errorf("failed to create Zhipu request: %w", err)
+	}
+	req.Header.Set("Authorization", "Bearer "+p.apiKey)
+	req.Header.Set("Content-Type", "application/json")
+
+	logger.Infof(ctx, "[WebSearch][Zhipu] query=%q maxResults=%d engine=%s", preparedQuery, maxResults, p.searchEngine)
+	resp, err := p.client.Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("failed to execute Zhipu request: %w", err)
+	}
+	defer resp.Body.Close()
+
+	respBody, err := readZhipuResponseBody(resp.Body)
+	if err != nil {
+		return nil, err
+	}
+	if resp.StatusCode != http.StatusOK {
+		return nil, zhipuHTTPError(resp.StatusCode, respBody)
+	}
+
+	var response zhipuSearchResponse
+	if err := json.Unmarshal(respBody, &response); err != nil {
+		return nil, fmt.Errorf("failed to unmarshal Zhipu response: %w", err)
+	}
+	if response.Error.Message != "" || response.Error.Code != "" {
+		return nil, fmt.Errorf("Zhipu API error (%s): %s", response.Error.Code, response.Error.Message)
+	}
+
+	results := make([]*types.WebSearchResult, 0, len(response.SearchResult))
+	for _, item := range response.SearchResult {
+		if strings.TrimSpace(item.Title) == "" && strings.TrimSpace(item.Link) == "" {
+			continue
+		}
+		result := &types.WebSearchResult{
+			Title:   item.Title,
+			URL:     item.Link,
+			Snippet: item.Content,
+			Source:  "zhipu",
+		}
+		if includeDate {
+			if publishedAt, ok := parseZhipuDate(item.PublishDate); ok {
+				result.PublishedAt = &publishedAt
+			}
+		}
+		results = append(results, result)
+		if len(results) >= maxResults {
+			break
+		}
+	}
+	logger.Infof(ctx, "[WebSearch][Zhipu] returned %d results", len(results))
+	return results, nil
+}
+
+func normalizeZhipuQuery(query string) string {
+	query = strings.TrimSpace(query)
+	if utf8.RuneCountInString(query) <= maxZhipuQueryRunes {
+		return query
+	}
+	runes := []rune(query)
+	return string(runes[:maxZhipuQueryRunes])
+}
+
+func parseZhipuDate(value string) (time.Time, bool) {
+	value = strings.TrimSpace(value)
+	if value == "" {
+		return time.Time{}, false
+	}
+	for _, layout := range []string{
+		time.RFC3339Nano,
+		"2006-01-02 15:04:05",
+		"2006-01-02 15:04",
+		"2006-01-02",
+	} {
+		if parsed, err := time.Parse(layout, value); err == nil {
+			return parsed, true
+		}
+	}
+	return time.Time{}, false
+}
+
+func readZhipuResponseBody(reader io.Reader) ([]byte, error) {
+	body, err := io.ReadAll(io.LimitReader(reader, maxZhipuResponseBytes+1))
+	if err != nil {
+		return nil, fmt.Errorf("failed to read Zhipu response: %w", err)
+	}
+	if len(body) > maxZhipuResponseBytes {
+		return nil, fmt.Errorf("Zhipu response exceeds %d bytes", maxZhipuResponseBytes)
+	}
+	return body, nil
+}
+
+func zhipuHTTPError(statusCode int, body []byte) error {
+	var response zhipuSearchResponse
+	if err := json.Unmarshal(body, &response); err == nil && (response.Error.Code != "" || response.Error.Message != "") {
+		return fmt.Errorf("Zhipu API returned status %d (%s): %s", statusCode, response.Error.Code, response.Error.Message)
+	}
+	detail := strings.TrimSpace(string(body))
+	if len(detail) > 4096 {
+		detail = detail[:4096]
+	}
+	if detail == "" {
+		return fmt.Errorf("Zhipu API returned status %d", statusCode)
+	}
+	return fmt.Errorf("Zhipu API returned status %d: %s", statusCode, detail)
+}
+
+type zhipuSearchRequest struct {
+	SearchQuery  string `json:"search_query"`
+	SearchEngine string `json:"search_engine"`
+	SearchIntent bool   `json:"search_intent"`
+	Count        int    `json:"count"`
+	ContentSize  string `json:"content_size"`
+}
+
+type zhipuSearchResponse struct {
+	ID           string              `json:"id"`
+	RequestID    string              `json:"request_id"`
+	SearchResult []zhipuSearchResult `json:"search_result"`
+	Error        zhipuError          `json:"error"`
+}
+
+type zhipuSearchResult struct {
+	Title       string `json:"title"`
+	Content     string `json:"content"`
+	Link        string `json:"link"`
+	Media       string `json:"media"`
+	Icon        string `json:"icon"`
+	Refer       string `json:"refer"`
+	PublishDate string `json:"publish_date"`
+}
+
+type zhipuError struct {
+	Code    string `json:"code"`
+	Message string `json:"message"`
+}