Refactor HuggingFaceInjectionDetector

hupe1980 · Sep 15, 2023 · c98f171 · c98f171
1 parent ed5f0f0
commit c98f171
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 6 deletions.
diff --git a/examples/hugging_face_injection_detector/main.go b/examples/hugging_face_injection_detector/main.go
@@ -6,8 +6,6 @@ import (
 	"log"
 	"os"
 
-	huggingface "github.com/hupe1980/go-huggingface"
-
 	"github.com/hupe1980/golc"
 	"github.com/hupe1980/golc/agent"
 	"github.com/hupe1980/golc/model/chatmodel"
@@ -23,10 +21,8 @@ func main() {
 		log.Fatal(err)
 	}
 
-	client := huggingface.NewInferenceClient(os.Getenv("HUGGINGFACEHUB_API_TOKEN"))
-
 	agent, err := agent.NewOpenAIFunctions(openai, []schema.Tool{
-		tool.NewHuggingFaceInjectionDetector(client),
+		tool.NewHuggingFaceInjectionDetector(os.Getenv("HUGGINGFACEHUB_API_TOKEN")),
 	})
 	if err != nil {
 		log.Fatal(err)

diff --git a/tool/hugging_face_injection_detector.go b/tool/hugging_face_injection_detector.go
@@ -33,7 +33,13 @@ type HuggingFaceInjectionDetector struct {
 }
 
 // NewHuggingFaceInjectionDetector creates a new instance of the HuggingFaceInjectionDetector tool.
-func NewHuggingFaceInjectionDetector(client HuggingFaceInjectionDetectorClient, optFns ...func(o *HuggingFaceInjectionDetectorOptions)) *HuggingFaceInjectionDetector {
+func NewHuggingFaceInjectionDetector(token string, optFns ...func(o *HuggingFaceInjectionDetectorOptions)) *HuggingFaceInjectionDetector {
+	client := huggingface.NewInferenceClient(token)
+	return NewHuggingFaceInjectionDetectorFromClient(client, optFns...)
+}
+
+// NewHuggingFaceInjectionDetectorFromClient creates a new instance of the HuggingFaceInjectionDetector tool.
+func NewHuggingFaceInjectionDetectorFromClient(client HuggingFaceInjectionDetectorClient, optFns ...func(o *HuggingFaceInjectionDetectorOptions)) *HuggingFaceInjectionDetector {
 	opts := HuggingFaceInjectionDetectorOptions{
 		Model:     "deepset/deberta-v3-base-injection",
 		Threshold: 0.8,