pluginsdk/plugin/wasmguest/caps/rag.go

package caps

import (
	"context"
	"sort"

	abiv1 "git.dev.alexdunmow.com/block/pluginsdk/abi/v1"
	"git.dev.alexdunmow.com/block/pluginsdk/plugin"
	"github.com/google/uuid"
)

// RAGStub implements plugin.RAGService. Query and OnContentChanged marshal out
// to rag.* capability calls, while RegisterContentFetcher records the fetcher
// guest-side: the host re-indexes by inverting the call as a HOOK_RAG_FETCH
// callback, so the wasm shim reads the recorded fetchers to dispatch it. The
// stub is exported so the shim can reach FetcherTypes/Fetcher.
type RAGStub struct {
	base
	fetchers map[string]plugin.ContentFetcher
}

var _ plugin.RAGService = (*RAGStub)(nil)

// NewRAGStub builds a RAG stub with the given transport. A nil transport
// (DESCRIBE probe, native build) still records fetchers so their content
// types can be captured into the manifest.
func NewRAGStub(call CallFunc) *RAGStub {
	return &RAGStub{
		base:     base{family: "rag", call: call},
		fetchers: make(map[string]plugin.ContentFetcher),
	}
}

func (s *RAGStub) RegisterContentFetcher(contentType string, fetcher plugin.ContentFetcher) {
	s.fetchers[contentType] = fetcher
}

func (s *RAGStub) Query(ctx context.Context, query string, limit int) ([]plugin.RAGResult, error) {
	req := &abiv1.RagQueryRequest{Query: query, Limit: int32(limit)}
	resp := &abiv1.RagQueryResponse{}
	if err := s.invoke(ctx, "query", req, resp); err != nil {
		return nil, err
	}
	results := make([]plugin.RAGResult, 0, len(resp.GetResults()))
	for _, r := range resp.GetResults() {
		results = append(results, plugin.RAGResult{
			Content:  r.GetContent(),
			Score:    r.GetScore(),
			Metadata: r.GetMetadata(),
		})
	}
	return results, nil
}

// OnContentChanged has no error channel; a transport failure is dropped (the
// host will re-index on its own schedule regardless).
func (s *RAGStub) OnContentChanged(ctx context.Context, contentType string, contentID uuid.UUID) {
	req := &abiv1.RagOnContentChangedRequest{ContentType: contentType, ContentId: contentID.String()}
	_ = s.invoke(ctx, "on_content_changed", req, &abiv1.RagOnContentChangedResponse{})
}

// Fetcher returns the content fetcher registered for a content type, for
// HOOK_RAG_FETCH dispatch by the wasm shim.
func (s *RAGStub) Fetcher(contentType string) (plugin.ContentFetcher, bool) {
	f, ok := s.fetchers[contentType]
	return f, ok
}

// FetcherTypes lists the registered content-fetcher types in sorted order
// (deterministic manifest capture).
func (s *RAGStub) FetcherTypes() []string {
	types := make([]string, 0, len(s.fetchers))
	for k := range s.fetchers {
		types = append(types, k)
	}
	sort.Strings(types)
	return types
}