organize to base_prompt_scorer

Author: abhishekg999

pkg/scorers/api_scorers/prompt_scorer/base_prompt_scorer.go

@@ -2,6 +2,7 @@ package prompt_scorer
 
 import (
 	"fmt"
+	"strconv"
 
 	"github.com/JudgmentLabs/judgeval-go/pkg/data"
 	"github.com/JudgmentLabs/judgeval-go/pkg/env"
@@ -18,6 +19,25 @@ type BasePromptScorer struct {
 	organizationID string
 }
 
+type ScorerOptions struct {
+	APIKey         string
+	OrganizationID string
+}
+
+type ScorerOption func(*ScorerOptions)
+
+func WithAPIKey(apiKey string) ScorerOption {
+	return func(opts *ScorerOptions) {
+		opts.APIKey = apiKey
+	}
+}
+
+func WithOrganizationID(orgID string) ScorerOption {
+	return func(opts *ScorerOptions) {
+		opts.OrganizationID = orgID
+	}
+}
+
 func NewBasePromptScorer(
 	scoreType data.APIScorerType,
 	name string,
@@ -38,6 +58,26 @@ func NewBasePromptScorer(
 	}
 }
 
+func parseScorerOptions(options interface{}) map[string]float64 {
+	result := make(map[string]float64)
+	if options == nil {
+		return result
+	}
+
+	if optionsMap, ok := options.(map[string]interface{}); ok {
+		for k, v := range optionsMap {
+			if num, ok := v.(float64); ok {
+				result[k] = num
+			} else if str, ok := v.(string); ok {
+				if num, err := strconv.ParseFloat(str, 64); err == nil {
+					result[k] = num
+				}
+			}
+		}
+	}
+	return result
+}
+
 func ScorerExists(name, judgmentAPIKey, organizationID string) (bool, error) {
 	client := api.NewClient(env.JudgmentAPIURL, judgmentAPIKey, organizationID)
 	request := &models.ScorerExistsRequest{
@@ -119,3 +159,22 @@ func (bps *BasePromptScorer) GetOptions() map[string]float64 {
 	}
 	return result
 }
+
+func (bps *BasePromptScorer) GetScorerConfig() models.ScorerConfig {
+	config := bps.APIScorer.GetScorerConfig()
+
+	kwargs := make(map[string]interface{})
+	kwargs["prompt"] = bps.GetPrompt()
+	if bps.GetOptions() != nil {
+		kwargs["options"] = bps.GetOptions()
+	}
+
+	if bps.APIScorer.AdditionalProperties != nil {
+		for k, v := range bps.APIScorer.AdditionalProperties {
+			kwargs[k] = v
+		}
+	}
+
+	config.Kwargs = kwargs
+	return config
+}

pkg/scorers/api_scorers/prompt_scorer/prompt_scorer.go

@@ -2,23 +2,26 @@ package prompt_scorer
 
 import (
 	"fmt"
-	"strconv"
 
 	"github.com/JudgmentLabs/judgeval-go/pkg/data"
 	"github.com/JudgmentLabs/judgeval-go/pkg/env"
-	"github.com/JudgmentLabs/judgeval-go/pkg/internal/api/models"
 )
 
 type PromptScorer struct {
 	*BasePromptScorer
 }
 
-func Get(name string) (*PromptScorer, error) {
-	return GetWithCredentials(name, env.JudgmentAPIKey, env.JudgmentOrgID)
-}
+func Get(name string, opts ...ScorerOption) (*PromptScorer, error) {
+	options := &ScorerOptions{
+		APIKey:         env.JudgmentAPIKey,
+		OrganizationID: env.JudgmentOrgID,
+	}
+
+	for _, opt := range opts {
+		opt(options)
+	}
 
-func GetWithCredentials(name, judgmentAPIKey, organizationID string) (*PromptScorer, error) {
-	scorerConfig, err := FetchPromptScorer(name, judgmentAPIKey, organizationID)
+	scorerConfig, err := FetchPromptScorer(name, options.APIKey, options.OrganizationID)
 	if err != nil {
 		return nil, err
 	}
@@ -27,21 +30,7 @@ func GetWithCredentials(name, judgmentAPIKey, organizationID string) (*PromptSco
 		return nil, fmt.Errorf("scorer with name %s is not a PromptScorer", name)
 	}
 
-	options := make(map[string]float64)
-	if scorerConfig.Options != nil {
-		if optionsMap, ok := scorerConfig.Options.(map[string]interface{}); ok {
-			for k, v := range optionsMap {
-				if num, ok := v.(float64); ok {
-					options[k] = num
-				} else if str, ok := v.(string); ok {
-					if num, err := strconv.ParseFloat(str, 64); err == nil {
-						options[k] = num
-					}
-				}
-			}
-		}
-	}
-
+	scorerOptions := parseScorerOptions(scorerConfig.Options)
 	threshold := 0.5
 	if scorerConfig.Threshold != 0 {
 		threshold = scorerConfig.Threshold
@@ -53,28 +42,13 @@ func GetWithCredentials(name, judgmentAPIKey, organizationID string) (*PromptSco
 			name,
 			scorerConfig.Prompt,
 			threshold,
-			options,
-			judgmentAPIKey,
-			organizationID,
+			scorerOptions,
+			options.APIKey,
+			options.OrganizationID,
 		),
 	}, nil
 }
 
-func (ps *PromptScorer) GetScorerConfig() models.ScorerConfig {
-	config := ps.BasePromptScorer.APIScorer.GetScorerConfig()
-
-	kwargs := make(map[string]interface{})
-	kwargs["prompt"] = ps.GetPrompt()
-	if ps.GetOptions() != nil {
-		kwargs["options"] = ps.GetOptions()
-	}
-
-	if ps.APIScorer.AdditionalProperties != nil {
-		for k, v := range ps.APIScorer.AdditionalProperties {
-			kwargs[k] = v
-		}
-	}
-
-	config.Kwargs = kwargs
-	return config
+func (ps *PromptScorer) IsTrace() bool {
+	return false
 }

pkg/scorers/api_scorers/prompt_scorer/trace_prompt_scorer.go

@@ -0,0 +1,54 @@
+package prompt_scorer
+
+import (
+	"fmt"
+
+	"github.com/JudgmentLabs/judgeval-go/pkg/data"
+	"github.com/JudgmentLabs/judgeval-go/pkg/env"
+)
+
+type TracePromptScorer struct {
+	*BasePromptScorer
+}
+
+func GetTrace(name string, opts ...ScorerOption) (*TracePromptScorer, error) {
+	options := &ScorerOptions{
+		APIKey:         env.JudgmentAPIKey,
+		OrganizationID: env.JudgmentOrgID,
+	}
+
+	for _, opt := range opts {
+		opt(options)
+	}
+
+	scorerConfig, err := FetchPromptScorer(name, options.APIKey, options.OrganizationID)
+	if err != nil {
+		return nil, err
+	}
+
+	if !scorerConfig.IsTrace {
+		return nil, fmt.Errorf("scorer with name %s is not a TracePromptScorer", name)
+	}
+
+	scorerOptions := parseScorerOptions(scorerConfig.Options)
+	threshold := 0.5
+	if scorerConfig.Threshold != 0 {
+		threshold = scorerConfig.Threshold
+	}
+
+	return &TracePromptScorer{
+		BasePromptScorer: NewBasePromptScorer(
+			data.TracePromptScorer,
+			name,
+			scorerConfig.Prompt,
+			threshold,
+			scorerOptions,
+			options.APIKey,
+			options.OrganizationID,
+		),
+	}, nil
+}
+
+func (tps *TracePromptScorer) IsTrace() bool {
+	return true
+}