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

internal/application/repository/retriever/opensearch/retrieve.go

@@ -0,0 +1,205 @@
+package opensearch
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+
+	osapi "github.com/opensearch-project/opensearch-go/v4/opensearchapi"
+
+	"github.com/Tencent/WeKnora/internal/logger"
+	"github.com/Tencent/WeKnora/internal/types"
+)
+
+// Retrieve dispatches to vectorRetrieve or keywordsRetrieve based on
+// params.RetrieverType. Returns exactly one *RetrieveResult — the
+// service-layer multi-store dispatcher depends on this one-bundle-per-
+// driver protocol. Score in result is RAW [0, 1] passthrough (the
+// OpenSearch k-NN plugin's SpaceType.COSINESIMIL.scoreTranslation has
+// already mapped to (1 + cos) / 2 before the score reaches us).
+//
+// Dim resolution order:
+//  1. params.AdditionalParams["dim"] if present and positive
+//  2. len(params.Embedding) if positive
+//  3. cross-dim multi-index search for keyword-only paths via
+//     "<base>_*" wildcard pattern
+//
+// Vector retrieve always requires a positive dim (it's intrinsic to k-NN).
+func (r *Repository) Retrieve(
+	ctx context.Context,
+	params types.RetrieveParams,
+) ([]*types.RetrieveResult, error) {
+	dim, multiIndex := resolveDim(params)
+
+	switch params.RetrieverType {
+	case types.VectorRetrieverType:
+		if dim == 0 {
+			return nil, fmt.Errorf(
+				"opensearch: vector retrieve requires embedding or AdditionalParams[\"dim\"]: %w",
+				ErrDimensionMismatch)
+		}
+		if err := r.ensureReady(ctx, dim); err != nil {
+			return nil, err
+		}
+		return r.vectorRetrieve(ctx, params, dim)
+	case types.KeywordsRetrieverType:
+		if multiIndex {
+			// BM25-only path without dim — search across "<base>_*"
+			// wildcard pattern. Slower than a single-index search but
+			// correct for admin / keyword-only callers.
+			return r.keywordsRetrieve(ctx, params, r.baseIndex+"_*")
+		}
+		if err := r.ensureReady(ctx, dim); err != nil {
+			return nil, err
+		}
+		return r.keywordsRetrieve(ctx, params, r.indexAlias(dim))
+	default:
+		return nil, fmt.Errorf("opensearch: unsupported retriever type %q", params.RetrieverType)
+	}
+}
+
+// resolveDim returns (dim, multiIndex). multiIndex=true means dim is
+// unknown and the caller should use a wildcard alias pattern (keyword
+// path only — vector retrieve cannot be dim-less).
+func resolveDim(p types.RetrieveParams) (int, bool) {
+	if p.AdditionalParams != nil {
+		if v, ok := p.AdditionalParams["dim"].(int); ok && v > 0 {
+			return v, false
+		}
+	}
+	if len(p.Embedding) > 0 {
+		return len(p.Embedding), false
+	}
+	return 0, true // unknown — multi-index fallback for keyword
+}
+
+func (r *Repository) vectorRetrieve(
+	ctx context.Context, p types.RetrieveParams, dim int,
+) ([]*types.RetrieveResult, error) {
+	body, err := buildKNNQuery(p.Embedding, effectiveTopK(ctx, p), p.Threshold, fromParams(p))
+	if err != nil {
+		return nil, err
+	}
+	hits, err := r.search(ctx, r.indexAlias(dim), body)
+	if err != nil {
+		return nil, err
+	}
+	return wrapResults(ctx, hits, types.VectorRetrieverType, types.MatchTypeEmbedding), nil
+}
+
+func (r *Repository) keywordsRetrieve(
+	ctx context.Context, p types.RetrieveParams, indexPattern string,
+) ([]*types.RetrieveResult, error) {
+	body, err := buildKeywordQuery(p.Query, effectiveTopK(ctx, p), p.Threshold, fromParams(p))
+	if err != nil {
+		return nil, err
+	}
+	hits, err := r.search(ctx, indexPattern, body)
+	if err != nil {
+		return nil, err
+	}
+	return wrapResults(ctx, hits, types.KeywordsRetrieverType, types.MatchTypeKeywords), nil
+}
+
+// effectiveTopK returns the caller's TopK or 10 if unset. The 10 is a
+// defensive default — every real caller sets TopK, so a TopK=10 in
+// production logs indicates a caller bug. Caps at 10000 to stay below
+// OpenSearch's index.max_result_window default.
+func effectiveTopK(ctx context.Context, p types.RetrieveParams) int {
+	if p.TopK <= 0 {
+		logger.GetLogger(ctx).Warnf("[OpenSearch] Retrieve called with TopK<=0; defaulting to 10 (caller bug?)")
+		return 10
+	}
+	if p.TopK > 10000 {
+		return 10000
+	}
+	return p.TopK
+}
+
+func (r *Repository) search(ctx context.Context, indexPattern string, body []byte) ([]hit, error) {
+	req := osapi.SearchReq{Indices: []string{indexPattern}, Body: bytes.NewReader(body)}
+	resp, err := r.client.Search(ctx, &req)
+	if err != nil {
+		if isNotFound(err) {
+			return nil, fmt.Errorf("opensearch: index %s missing: %w", indexPattern, ErrIndexNotFound)
+		}
+		return nil, wrapTransport(err)
+	}
+	defer func() {
+		if resp != nil {
+			drainAndClose(resp.Inspect().Response.Body)
+		}
+	}()
+	// Cap response body at 16MB (search results bounded by max_result_window * doc_size).
+	return parseSearchHits(io.LimitReader(resp.Inspect().Response.Body, 16<<20))
+}
+
+// hit is the structural shape we pull from each OpenSearch search hit.
+// Field-by-field decode (vs map[string]any) keeps the JSON shape
+// pinned at compile time.
+type hit struct {
+	ID     string  `json:"_id"` // equals chunk_id per the indexing invariant
+	Score  float64 `json:"_score"`
+	Source struct {
+		Content         string `json:"content"`
+		ChunkID         string `json:"chunk_id"`
+		KnowledgeID     string `json:"knowledge_id"`
+		KnowledgeBaseID string `json:"knowledge_base_id"`
+		SourceID        string `json:"source_id"`
+		SourceType      int    `json:"source_type"` // integer, not stringified
+		TagID           string `json:"tag_id"`
+		IsEnabled       bool   `json:"is_enabled"`
+	} `json:"_source"`
+}
+
+func parseSearchHits(body io.Reader) ([]hit, error) {
+	var resp struct {
+		Hits struct {
+			Hits []hit `json:"hits"`
+		} `json:"hits"`
+	}
+	if err := json.NewDecoder(body).Decode(&resp); err != nil {
+		return nil, fmt.Errorf("opensearch: parse search response: %w", ErrTransport)
+	}
+	return resp.Hits.Hits, nil
+}
+
+// wrapResults converts OpenSearch hits to the upstream IndexWithScore
+// shape. Always returns exactly one *RetrieveResult — the service-layer
+// multi-store dispatcher depends on this "one engine result per driver"
+// contract.
+//
+// Defensive: warns if hit._id != _source.chunk_id (we set _id=chunk_id
+// on every write, so a mismatch means external tooling has been writing
+// into the index) but does not fail the query — the result is still
+// usable for ranking, just suspicious.
+func wrapResults(ctx context.Context, hits []hit, rt types.RetrieverType, mt types.MatchType) []*types.RetrieveResult {
+	log := logger.GetLogger(ctx)
+	out := make([]*types.IndexWithScore, 0, len(hits))
+	for _, h := range hits {
+		if h.ID != h.Source.ChunkID {
+			log.Warnf("[OpenSearch] hit._id=%q != _source.chunk_id=%q (D12 invariant violation)",
+				h.ID, h.Source.ChunkID)
+		}
+		out = append(out, &types.IndexWithScore{
+			ID:              h.ID,
+			ChunkID:         h.Source.ChunkID,
+			KnowledgeID:     h.Source.KnowledgeID,
+			KnowledgeBaseID: h.Source.KnowledgeBaseID,
+			SourceID:        h.Source.SourceID,
+			SourceType:      types.SourceType(h.Source.SourceType),
+			TagID:           h.Source.TagID,
+			Content:         h.Source.Content,
+			Score:           h.Score,
+			MatchType:       mt,
+			IsEnabled:       h.Source.IsEnabled,
+		})
+	}
+	return []*types.RetrieveResult{{
+		Results:             out,
+		RetrieverEngineType: types.OpenSearchRetrieverEngineType,
+		RetrieverType:       rt,
+	}}
+}