knowledge-governance/internal/platform/weknora/client.go

package weknora

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"mime/multipart"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"
)

const defaultMaxResponseBytes int64 = 2 * 1024 * 1024

var ErrConflict = errors.New("WeKnora resource conflict")

type HTTPError struct {
	StatusCode int
}

func (e *HTTPError) Error() string {
	return fmt.Sprintf("WeKnora HTTP status %d", e.StatusCode)
}

func IsRetryable(err error) bool {
	var httpErr *HTTPError
	if errors.As(err, &httpErr) {
		return httpErr.StatusCode == http.StatusTooManyRequests || httpErr.StatusCode >= 500
	}
	if errors.Is(err, context.Canceled) {
		return false
	}
	if errors.Is(err, context.DeadlineExceeded) {
		return true
	}
	var networkErr net.Error
	return errors.As(err, &networkErr)
}

type ClientConfig struct {
	BaseURL          string
	APIKey           string
	Timeout          time.Duration
	MaxResponseBytes int64
	AllowDockerHTTP  bool
}

type Client struct {
	baseURL          *url.URL
	apiKey           string
	httpClient       *http.Client
	maxResponseBytes int64
}

func NewClient(config ClientConfig) (*Client, error) {
	baseURL, err := url.Parse(strings.TrimSpace(config.BaseURL))
	if err != nil {
		return nil, errors.New("invalid WeKnora base URL")
	}
	if baseURL.Scheme == "" || baseURL.Host == "" || baseURL.User != nil || baseURL.RawQuery != "" || baseURL.Fragment != "" {
		return nil, errors.New("invalid WeKnora base URL")
	}
	if baseURL.Scheme != "https" && !(baseURL.Scheme == "http" && (isLoopbackHost(baseURL.Hostname()) || (config.AllowDockerHTTP && isSingleLabelDockerHost(baseURL.Hostname())))) {
		return nil, errors.New("WeKnora base URL must use HTTPS")
	}
	baseURL.Path = strings.TrimRight(baseURL.Path, "/")
	if strings.TrimSpace(config.APIKey) == "" {
		return nil, errors.New("WeKnora API key is required")
	}
	if config.Timeout <= 0 {
		return nil, errors.New("WeKnora HTTP timeout must be positive")
	}
	if config.MaxResponseBytes <= 0 {
		config.MaxResponseBytes = defaultMaxResponseBytes
	}
	return &Client{
		baseURL: baseURL,
		apiKey:  config.APIKey,
		httpClient: &http.Client{
			Timeout: config.Timeout,
			CheckRedirect: func(*http.Request, []*http.Request) error {
				return http.ErrUseLastResponse
			},
		},
		maxResponseBytes: config.MaxResponseBytes,
	}, nil
}

func isLoopbackHost(host string) bool {
	if strings.EqualFold(host, "localhost") {
		return true
	}
	ip := net.ParseIP(host)
	return ip != nil && ip.IsLoopback()
}

func isSingleLabelDockerHost(host string) bool {
	if host == "" || strings.Contains(host, ".") || net.ParseIP(host) != nil || len(host) > 63 {
		return false
	}
	for index, char := range host {
		if (char >= 'a' && char <= 'z') || (char >= 'A' && char <= 'Z') || (char >= '0' && char <= '9') {
			continue
		}
		if char != '-' || index == 0 || index == len(host)-1 {
			return false
		}
	}
	return true
}

func (c *Client) GetKnowledge(ctx context.Context, knowledgeID string) (Knowledge, error) {
	if !validPathSegment(knowledgeID) {
		return Knowledge{}, errors.New("invalid knowledge ID")
	}
	var response knowledgeResponse
	if err := c.doJSON(ctx, http.MethodGet, "/api/v1/knowledge/"+knowledgeID, nil, &response); err != nil {
		return Knowledge{}, err
	}
	if !response.Success || response.Data.ID == "" {
		return Knowledge{}, errors.New("WeKnora returned an unsuccessful knowledge response")
	}
	return response.Data, nil
}

func (c *Client) DeleteKnowledge(ctx context.Context, knowledgeID string) error {
	if !validPathSegment(knowledgeID) {
		return errors.New("invalid knowledge ID")
	}
	var response struct {
		Success bool `json:"success"`
	}
	if err := c.doJSON(ctx, http.MethodDelete, "/api/v1/knowledge/"+knowledgeID, nil, &response); err != nil {
		return err
	}
	if !response.Success {
		return errors.New("WeKnora returned an unsuccessful delete response")
	}
	return nil
}

func (c *Client) UploadFile(ctx context.Context, request UploadRequest) (Knowledge, error) {
	if !validPathSegment(request.KnowledgeBaseID) {
		return Knowledge{}, errors.New("invalid knowledge base ID")
	}
	if strings.TrimSpace(request.FileName) == "" || strings.ContainsAny(request.FileName, "/\\\x00") {
		return Knowledge{}, errors.New("invalid upload file name")
	}
	if request.Content == nil {
		return Knowledge{}, errors.New("upload content is required")
	}

	pipeReader, pipeWriter := io.Pipe()
	writer := multipart.NewWriter(pipeWriter)
	writeResult := make(chan error, 1)
	go func() {
		err := writeMultipartUpload(writer, request)
		if err == nil {
			err = writer.Close()
		}
		_ = pipeWriter.CloseWithError(err)
		writeResult <- err
	}()

	endpoint := c.endpoint("/api/v1/knowledge-bases/" + request.KnowledgeBaseID + "/knowledge/file")
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, pipeReader)
	if err != nil {
		_ = pipeReader.CloseWithError(err)
		<-writeResult
		return Knowledge{}, errors.New("create WeKnora upload request")
	}
	req.Header.Set("Accept", "application/json")
	req.Header.Set("Content-Type", writer.FormDataContentType())
	req.Header.Set("X-API-Key", c.apiKey)
	resp, err := c.httpClient.Do(req)
	if err != nil {
		_ = pipeReader.CloseWithError(err)
		<-writeResult
		return Knowledge{}, fmt.Errorf("send WeKnora upload request: %w", err)
	}
	if writeErr := <-writeResult; writeErr != nil {
		resp.Body.Close()
		return Knowledge{}, errors.New("stream WeKnora upload body")
	}

	var response knowledgeResponse
	if resp.StatusCode == http.StatusConflict {
		if err := c.decodeBody(resp, &response); err != nil {
			return Knowledge{}, err
		}
		return response.Data, ErrConflict
	}
	if err := c.decodeResponse(resp, &response); err != nil {
		return Knowledge{}, err
	}
	if !response.Success || response.Data.ID == "" {
		return Knowledge{}, errors.New("WeKnora returned an unsuccessful upload response")
	}
	return response.Data, nil
}

func (c *Client) ListKnowledge(ctx context.Context, knowledgeBaseID string, page, pageSize int) ([]Knowledge, int64, error) {
	if !validPathSegment(knowledgeBaseID) || !validPage(page, pageSize) {
		return nil, 0, errors.New("invalid knowledge list request")
	}
	query := url.Values{}
	query.Set("page", fmt.Sprint(page))
	query.Set("page_size", fmt.Sprint(pageSize))
	var response knowledgeListResponse
	if err := c.doJSONQuery(ctx, http.MethodGet, "/api/v1/knowledge-bases/"+knowledgeBaseID+"/knowledge", query, nil, &response); err != nil {
		return nil, 0, err
	}
	if !response.Success || response.Total < 0 {
		return nil, 0, errors.New("WeKnora returned an unsuccessful knowledge list")
	}
	return response.Data, response.Total, nil
}

// ListKnowledgeBases lists the knowledge bases visible to this API key so the
// UI can bind routing targets without humans copying UUIDs.
func (c *Client) ListKnowledgeBases(ctx context.Context) ([]KnowledgeBase, error) {
	var response knowledgeBaseListResponse
	if err := c.doJSON(ctx, http.MethodGet, "/api/v1/knowledge-bases", nil, &response); err != nil {
		return nil, err
	}
	if !response.Success {
		return nil, errors.New("WeKnora returned an unsuccessful knowledge base list")
	}
	return response.Data, nil
}

// GetKnowledgeBase reads a single knowledge base by id. A knowledge base that is
// not there surfaces as *HTTPError with status 404, so a caller can tell
// "definitely absent" from "could not ask".
func (c *Client) GetKnowledgeBase(ctx context.Context, knowledgeBaseID string) (KnowledgeBase, error) {
	if !validPathSegment(knowledgeBaseID) {
		return KnowledgeBase{}, errors.New("invalid knowledge base ID")
	}
	var response knowledgeBaseResponse
	if err := c.doJSON(ctx, http.MethodGet, "/api/v1/knowledge-bases/"+knowledgeBaseID, nil, &response); err != nil {
		return KnowledgeBase{}, err
	}
	if !response.Success || response.Data.ID == "" {
		return KnowledgeBase{}, errors.New("WeKnora returned an unsuccessful knowledge base response")
	}
	return response.Data, nil
}

func (c *Client) ListChunks(ctx context.Context, knowledgeID string, page, pageSize int, chunkTypes []string) ([]Chunk, int64, error) {
	if !validPathSegment(knowledgeID) || !validPage(page, pageSize) {
		return nil, 0, errors.New("invalid chunk list request")
	}
	query := url.Values{}
	query.Set("page", fmt.Sprint(page))
	query.Set("page_size", fmt.Sprint(pageSize))
	for _, chunkType := range chunkTypes {
		if !validPathSegment(chunkType) {
			return nil, 0, errors.New("invalid chunk type")
		}
		query.Add("chunk_type", chunkType)
	}
	var response chunkListResponse
	if err := c.doJSONQuery(ctx, http.MethodGet, "/api/v1/chunks/"+knowledgeID, query, nil, &response); err != nil {
		return nil, 0, err
	}
	if !response.Success || response.Total < 0 {
		return nil, 0, errors.New("WeKnora returned an unsuccessful chunk list")
	}
	return response.Data, response.Total, nil
}

// ListWikiPages lists the wiki pages of a knowledge base. Used by release
// verification to confirm the production wiki matches the released snapshot.
func (c *Client) ListWikiPages(ctx context.Context, knowledgeBaseID string, page, pageSize int) ([]WikiPageSummary, int64, error) {
	if !validPathSegment(knowledgeBaseID) || !validPage(page, pageSize) {
		return nil, 0, errors.New("invalid wiki page list request")
	}
	query := url.Values{}
	query.Set("page", fmt.Sprint(page))
	query.Set("page_size", fmt.Sprint(pageSize))
	var response wikiPageListResponse
	if err := c.doJSONQuery(ctx, http.MethodGet, "/api/v1/knowledgebase/"+knowledgeBaseID+"/wiki/pages", query, nil, &response); err != nil {
		return nil, 0, err
	}
	if response.Total < 0 {
		return nil, 0, errors.New("WeKnora returned an unsuccessful wiki page list")
	}
	return response.Pages, response.Total, nil
}

func (c *Client) SearchKnowledge(ctx context.Context, request SearchRequest) ([]SearchResult, error) {
	if strings.TrimSpace(request.Query) == "" || len(request.KnowledgeIDs) == 0 {
		return nil, errors.New("search query and knowledge IDs are required")
	}
	for _, knowledgeID := range request.KnowledgeIDs {
		if !validPathSegment(knowledgeID) {
			return nil, errors.New("invalid search knowledge ID")
		}
	}
	body, err := json.Marshal(request)
	if err != nil {
		return nil, errors.New("encode WeKnora search request")
	}
	var response searchResponse
	if err := c.doJSON(ctx, http.MethodPost, "/api/v1/knowledge-search", bytes.NewReader(body), &response); err != nil {
		return nil, err
	}
	if !response.Success {
		return nil, errors.New("WeKnora returned an unsuccessful search response")
	}
	return response.Data, nil
}

func writeMultipartUpload(writer *multipart.Writer, request UploadRequest) error {
	part, err := writer.CreateFormFile("file", request.FileName)
	if err != nil {
		return err
	}
	if _, err := io.Copy(part, request.Content); err != nil {
		return err
	}
	if request.Metadata != nil {
		value, err := json.Marshal(request.Metadata)
		if err != nil {
			return err
		}
		if err := writer.WriteField("metadata", string(value)); err != nil {
			return err
		}
	}
	if err := writer.WriteField("fileName", request.FileName); err != nil {
		return err
	}
	if request.Channel != "" {
		if err := writer.WriteField("channel", request.Channel); err != nil {
			return err
		}
	}
	if request.ProcessConfig != nil {
		value, err := json.Marshal(request.ProcessConfig)
		if err != nil {
			return err
		}
		if err := writer.WriteField("process_config", string(value)); err != nil {
			return err
		}
	}
	return nil
}

func validPathSegment(value string) bool {
	return strings.TrimSpace(value) != "" && !strings.ContainsAny(value, "/?#\\\x00")
}

func validPage(page, pageSize int) bool {
	return page > 0 && pageSize > 0 && pageSize <= 100
}

func (c *Client) endpoint(path string) string {
	endpoint := *c.baseURL
	endpoint.Path = strings.TrimRight(c.baseURL.Path, "/") + path
	return endpoint.String()
}

func (c *Client) doJSON(ctx context.Context, method, path string, body io.Reader, target any) error {
	return c.doJSONQuery(ctx, method, path, nil, body, target)
}

func (c *Client) doJSONQuery(ctx context.Context, method, path string, query url.Values, body io.Reader, target any) error {
	endpoint := c.endpoint(path)
	if len(query) > 0 {
		endpoint += "?" + query.Encode()
	}
	req, err := http.NewRequestWithContext(ctx, method, endpoint, body)
	if err != nil {
		return errors.New("create WeKnora request")
	}
	req.Header.Set("Accept", "application/json")
	if body != nil {
		req.Header.Set("Content-Type", "application/json")
	}
	req.Header.Set("X-API-Key", c.apiKey)
	resp, err := c.httpClient.Do(req)
	if err != nil {
		return fmt.Errorf("send WeKnora request: %w", err)
	}
	return c.decodeResponse(resp, target)
}

func (c *Client) decodeResponse(resp *http.Response, target any) error {
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		resp.Body.Close()
		return &HTTPError{StatusCode: resp.StatusCode}
	}
	return c.decodeBody(resp, target)
}

func (c *Client) decodeBody(resp *http.Response, target any) error {
	defer resp.Body.Close()
	limited := io.LimitReader(resp.Body, c.maxResponseBytes+1)
	payload, err := io.ReadAll(limited)
	if err != nil {
		return errors.New("read WeKnora response")
	}
	if int64(len(payload)) > c.maxResponseBytes {
		return errors.New("WeKnora response exceeds size limit")
	}
	if err := json.Unmarshal(payload, target); err != nil {
		return errors.New("decode WeKnora response")
	}
	return nil
}