Add configurable io.writer

hupe1980 · Jul 27, 2023 · fd7c658 · fd7c658
1 parent 58d6370
commit fd7c658
Show file tree

Hide file tree

Showing 11 changed files with 209 additions and 180 deletions.
diff --git a/callback/callback.go b/callback/callback.go
@@ -1,87 +1,2 @@
 // Package callback provides utilities for implementing callbacks.
 package callback
-
-import (
-	"context"
-
-	"github.com/hupe1980/golc/schema"
-)
-
-// Compile time check to ensure handler satisfies the Callback interface.
-var _ schema.Callback = (*handler)(nil)
-
-type handler struct{}
-
-func (h *handler) AlwaysVerbose() bool {
-	return false
-}
-
-func (h *handler) RaiseError() bool {
-	return false
-}
-
-func (h *handler) OnLLMStart(ctx context.Context, input *schema.LLMStartInput) error {
-	return nil
-}
-
-func (h *handler) OnChatModelStart(ctx context.Context, input *schema.ChatModelStartInput) error {
-	return nil
-}
-
-func (h *handler) OnModelNewToken(ctx context.Context, input *schema.ModelNewTokenInput) error {
-	return nil
-}
-
-func (h *handler) OnModelEnd(ctx context.Context, input *schema.ModelEndInput) error {
-	return nil
-}
-
-func (h *handler) OnModelError(ctx context.Context, input *schema.ModelErrorInput) error {
-	return nil
-}
-
-func (h *handler) OnChainStart(ctx context.Context, input *schema.ChainStartInput) error {
-	return nil
-}
-
-func (h *handler) OnChainEnd(ctx context.Context, input *schema.ChainEndInput) error {
-	return nil
-}
-
-func (h *handler) OnChainError(ctx context.Context, input *schema.ChainErrorInput) error {
-	return nil
-}
-
-func (h *handler) OnAgentAction(ctx context.Context, input *schema.AgentActionInput) error {
-	return nil
-}
-
-func (h *handler) OnAgentFinish(ctx context.Context, input *schema.AgentFinishInput) error {
-	return nil
-}
-
-func (h *handler) OnToolStart(ctx context.Context, input *schema.ToolStartInput) error {
-	return nil
-}
-func (h *handler) OnToolEnd(ctx context.Context, input *schema.ToolEndInput) error {
-	return nil
-}
-func (h *handler) OnToolError(ctx context.Context, input *schema.ToolErrorInput) error {
-	return nil
-}
-
-func (h *handler) OnText(ctx context.Context, input *schema.TextInput) error {
-	return nil
-}
-
-func (h *handler) OnRetrieverStart(ctx context.Context, input *schema.RetrieverStartInput) error {
-	return nil
-}
-
-func (h *handler) OnRetrieverEnd(ctx context.Context, input *schema.RetrieverEndInput) error {
-	return nil
-}
-
-func (h *handler) OnRetrieverError(ctx context.Context, input *schema.RetrieverErrorInput) error {
-	return nil
-}
diff --git a/callback/manager.go b/callback/manager.go
@@ -43,8 +43,8 @@ func newManager(runID string, inheritableCallbacks, localCallbacks []schema.Call
 	}
 
 	callbacks := append(inheritableCallbacks, localCallbacks...)
-	if verbose && !containsStdOutCallbackHandler(callbacks) {
-		callbacks = append(callbacks, NewStdOutHandler())
+	if verbose && !containsWriterCallbackHandler(callbacks) {
+		callbacks = append(callbacks, NewWriterHandler())
 	}
 
 	return &manager{
@@ -384,9 +384,9 @@ func (m *manager) OnRetrieverError(ctx context.Context, input *schema.RetrieverE
 	return nil
 }
 
-func containsStdOutCallbackHandler(handlers []schema.Callback) bool {
+func containsWriterCallbackHandler(handlers []schema.Callback) bool {
 	for _, handler := range handlers {
-		if _, ok := handler.(*StdOutHandler); ok {
+		if _, ok := handler.(*WriterHandler); ok {
 			return true
 		}
 	}

diff --git a/callback/noop_handler.go b/callback/noop_handler.go
@@ -0,0 +1,86 @@
+package callback
+
+import (
+	"context"
+
+	"github.com/hupe1980/golc/schema"
+)
+
+// Compile time check to ensure NoopHandler satisfies the Callback interface.
+var _ schema.Callback = (*NoopHandler)(nil)
+
+type NoopHandler struct{}
+
+func (h *NoopHandler) AlwaysVerbose() bool {
+	return false
+}
+
+func (h *NoopHandler) RaiseError() bool {
+	return false
+}
+
+func (h *NoopHandler) OnLLMStart(ctx context.Context, input *schema.LLMStartInput) error {
+	return nil
+}
+
+func (h *NoopHandler) OnChatModelStart(ctx context.Context, input *schema.ChatModelStartInput) error {
+	return nil
+}
+
+func (h *NoopHandler) OnModelNewToken(ctx context.Context, input *schema.ModelNewTokenInput) error {
+	return nil
+}
+
+func (h *NoopHandler) OnModelEnd(ctx context.Context, input *schema.ModelEndInput) error {
+	return nil
+}
+
+func (h *NoopHandler) OnModelError(ctx context.Context, input *schema.ModelErrorInput) error {
+	return nil
+}
+
+func (h *NoopHandler) OnChainStart(ctx context.Context, input *schema.ChainStartInput) error {
+	return nil
+}
+
+func (h *NoopHandler) OnChainEnd(ctx context.Context, input *schema.ChainEndInput) error {
+	return nil
+}
+
+func (h *NoopHandler) OnChainError(ctx context.Context, input *schema.ChainErrorInput) error {
+	return nil
+}
+
+func (h *NoopHandler) OnAgentAction(ctx context.Context, input *schema.AgentActionInput) error {
+	return nil
+}
+
+func (h *NoopHandler) OnAgentFinish(ctx context.Context, input *schema.AgentFinishInput) error {
+	return nil
+}
+
+func (h *NoopHandler) OnToolStart(ctx context.Context, input *schema.ToolStartInput) error {
+	return nil
+}
+func (h *NoopHandler) OnToolEnd(ctx context.Context, input *schema.ToolEndInput) error {
+	return nil
+}
+func (h *NoopHandler) OnToolError(ctx context.Context, input *schema.ToolErrorInput) error {
+	return nil
+}
+
+func (h *NoopHandler) OnText(ctx context.Context, input *schema.TextInput) error {
+	return nil
+}
+
+func (h *NoopHandler) OnRetrieverStart(ctx context.Context, input *schema.RetrieverStartInput) error {
+	return nil
+}
+
+func (h *NoopHandler) OnRetrieverEnd(ctx context.Context, input *schema.RetrieverEndInput) error {
+	return nil
+}
+
+func (h *NoopHandler) OnRetrieverError(ctx context.Context, input *schema.RetrieverErrorInput) error {
+	return nil
+}
diff --git a/callback/openai_info.go b/callback/openai_info.go
@@ -39,7 +39,7 @@ var modelCostMapping = map[string]float64{
 var _ schema.Callback = (*OpenAIHandler)(nil)
 
 type OpenAIHandler struct {
-	handler
+	NoopHandler
 	totalTokens        int
 	promptTokens       int
 	completionTokens   int

diff --git a/callback/prompt_layer.go b/callback/prompt_layer.go
@@ -21,7 +21,7 @@ type PromptLayerHandlerOptions struct {
 }
 
 type PromptLayerHandler struct {
-	handler
+	NoopHandler
 	apiKey  string
 	client  *promptlayer.Client
 	runInfo map[string]map[string]any

diff --git a/callback/stdout.go b/callback/stdout.go
diff --git a/callback/stream_writer_handler.go b/callback/stream_writer_handler.go
@@ -0,0 +1,47 @@
+package callback
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"os"
+
+	"github.com/hupe1980/golc/schema"
+)
+
+// Compile time check to ensure StreamWriterHandler satisfies the Callback interface.
+var _ schema.Callback = (*StreamWriterHandler)(nil)
+
+type StreamWriterHandlerOptions struct {
+	Writer io.Writer
+}
+
+type StreamWriterHandler struct {
+	NoopHandler
+	writer io.Writer
+	opts   StreamWriterHandlerOptions
+}
+
+func NewStreamWriterHandler(optFns ...func(o *StreamWriterHandlerOptions)) *StreamWriterHandler {
+	opts := StreamWriterHandlerOptions{
+		Writer: os.Stdout,
+	}
+
+	for _, fn := range optFns {
+		fn(&opts)
+	}
+
+	return &StreamWriterHandler{
+		writer: opts.Writer,
+		opts:   opts,
+	}
+}
+
+func (cb *StreamWriterHandler) AlwaysVerbose() bool {
+	return true
+}
+
+func (cb *StreamWriterHandler) OnModelNewToken(ctx context.Context, input *schema.ModelNewTokenInput) error {
+	fmt.Fprint(cb.writer, input.Token)
+	return nil
+}
diff --git a/callback/streaming_stdout.go b/callback/streaming_stdout.go