feat: refactor knowledge retrieval to introduce RetrieveContext method and streamline context handling

internal/ai/runtime/internal/impl/retrievers/knowledge_retriever.go

@@ -2,16 +2,28 @@ package retrievers
 
 import (
 	"context"
+	"strings"
 
 	"cs-agent/internal/ai/rag"
 	"cs-agent/internal/models"
 	"cs-agent/internal/pkg/utils"
 )
 
+const defaultRuntimeKnowledgeContextTokens = 4000
+
 type KnowledgeRetriever struct {
 	AIAgent *models.AIAgent
 }
 
+type KnowledgeRetrieveResult struct {
+	KnowledgeBaseIDs []int64
+	Query            string
+	Hits             []rag.RetrieveResult
+	ContextResults   []rag.RetrieveResult
+	ContextText      string
+	Trace            *rag.RetrieveTrace
+}
+
 func NewKnowledgeRetriever(aiAgent *models.AIAgent) *KnowledgeRetriever {
 	return &KnowledgeRetriever{AIAgent: aiAgent}
 }
@@ -30,3 +42,24 @@ func (r *KnowledgeRetriever) Retrieve(ctx context.Context, query string) ([]rag.
 		KnowledgeBaseIDs: ids,
 	})
 }
+
+func (r *KnowledgeRetriever) RetrieveContext(ctx context.Context, query string) (*KnowledgeRetrieveResult, error) {
+	query = strings.TrimSpace(query)
+	knowledgeBaseIDs := r.KnowledgeBaseIDs()
+	ret := &KnowledgeRetrieveResult{
+		KnowledgeBaseIDs: append([]int64(nil), knowledgeBaseIDs...),
+		Query:            query,
+	}
+	if query == "" || len(knowledgeBaseIDs) == 0 {
+		return ret, nil
+	}
+	results, trace, err := r.Retrieve(ctx, query)
+	if err != nil {
+		return nil, err
+	}
+	ret.Hits = append([]rag.RetrieveResult(nil), results...)
+	ret.Trace = trace
+	ret.ContextResults = rag.Retrieve.SelectContextResults(results, defaultRuntimeKnowledgeContextTokens)
+	ret.ContextText = strings.TrimSpace(rag.Retrieve.BuildContext(ctx, results, defaultRuntimeKnowledgeContextTokens))
+	return ret, nil
+}