From 31774e4f850f82dc7b5089ae02e5989cd79b672c Mon Sep 17 00:00:00 2001
From: Grant Linville <grant@acorn.io>
Date: Mon, 25 Mar 2024 09:36:33 -0400
Subject: [PATCH] enhance: add green color to tool call names

Signed-off-by: Grant Linville <grant@acorn.io>
---
 go.mod                  | 2 +-
 pkg/cli/gptscript.go    | 6 ++++++
 pkg/types/completion.go | 4 +++-
 3 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/go.mod b/go.mod
index 9ad35830..0c118fff 100644
--- a/go.mod
+++ b/go.mod
@@ -8,6 +8,7 @@ require (
 	github.com/acorn-io/broadcaster v0.0.0-20240105011354-bfadd4a7b45d
 	github.com/acorn-io/cmd v0.0.0-20240203032901-e9e631185ddb
 	github.com/adrg/xdg v0.4.0
+	github.com/fatih/color v1.16.0
 	github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510
 	github.com/hexops/autogold/v2 v2.1.0
 	github.com/jaytaylor/html2text v0.0.0-20230321000545-74c2419ad056
@@ -34,7 +35,6 @@ require (
 	github.com/connesc/cipherio v0.2.1 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/dsnet/compress v0.0.1 // indirect
-	github.com/fatih/color v1.16.0 // indirect
 	github.com/go-logr/logr v1.4.1 // indirect
 	github.com/golang/snappy v0.0.4 // indirect
 	github.com/google/go-cmp v0.6.0 // indirect
diff --git a/pkg/cli/gptscript.go b/pkg/cli/gptscript.go
index f09f661c..7ea2469e 100644
--- a/pkg/cli/gptscript.go
+++ b/pkg/cli/gptscript.go
@@ -7,6 +7,7 @@ import (
 	"strings"
 
 	"github.com/acorn-io/cmd"
+	"github.com/fatih/color"
 	"github.com/gptscript-ai/gptscript/pkg/assemble"
 	"github.com/gptscript-ai/gptscript/pkg/builtin"
 	"github.com/gptscript-ai/gptscript/pkg/cache"
@@ -34,6 +35,7 @@ type GPTScript struct {
 	CacheOptions
 	OpenAIOptions
 	DisplayOptions
+	Color         *bool  `usage:"Use color in output (default true)" default:"true"`
 	Confirm       bool   `usage:"Prompt before running potentially dangerous commands"`
 	Debug         bool   `usage:"Enable debug logging"`
 	Quiet         *bool  `usage:"No output logging (set --quiet=false to force on even when there is no TTY)" short:"q"`
@@ -113,6 +115,10 @@ func (r *GPTScript) Pre(*cobra.Command, []string) error {
 }
 
 func (r *GPTScript) Run(cmd *cobra.Command, args []string) error {
+	if r.Color != nil {
+		color.NoColor = !*r.Color
+	}
+
 	gptOpt := gptscript.Options{
 		Cache:   cache.Options(r.CacheOptions),
 		OpenAI:  openai.Options(r.OpenAIOptions),
diff --git a/pkg/types/completion.go b/pkg/types/completion.go
index b78a1aa4..e7a0d418 100644
--- a/pkg/types/completion.go
+++ b/pkg/types/completion.go
@@ -3,6 +3,8 @@ package types
 import (
 	"fmt"
 	"strings"
+
+	"github.com/fatih/color"
 )
 
 type CompletionRequest struct {
@@ -81,7 +83,7 @@ func (in CompletionMessage) String() string {
 		}
 		buf.WriteString(content.Text)
 		if content.ToolCall != nil {
-			buf.WriteString(fmt.Sprintf("tool call %s -> %s", content.ToolCall.Function.Name, content.ToolCall.Function.Arguments))
+			buf.WriteString(fmt.Sprintf("tool call %s -> %s", color.GreenString(content.ToolCall.Function.Name), content.ToolCall.Function.Arguments))
 		}
 	}
 	return buf.String()