From 33174dc8423305d227d381a355dafce0f50169cc Mon Sep 17 00:00:00 2001 From: Simun Strukan Date: Tue, 31 Oct 2023 20:35:34 +0100 Subject: [PATCH] Added comments --- callbacks/agent_final_stream.go | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/callbacks/agent_final_stream.go b/callbacks/agent_final_stream.go index 8a40ee4c6..a841d5473 100644 --- a/callbacks/agent_final_stream.go +++ b/callbacks/agent_final_stream.go @@ -5,6 +5,8 @@ import ( "strings" ) +// DefaultKeywords is map of the agents final out prefix keywords. +// //nolint:all var DefaultKeywords = []string{"Final Answer:", "Final:", "AI:"} @@ -21,6 +23,14 @@ type AgentFinalStreamHandler struct { var _ Handler = &AgentFinalStreamHandler{} +// NewFinalStreamHandler creates a new instance of the AgentFinalStreamHandler struct. +// +// It accepts a variadic number of strings as keywords. If any keywords are provided, +// the DefaultKeywords variable is updated with the provided keywords. +// +// DefualtKeywords is map of the agents final out prefix keywords. +// +// The function returns a pointer to the created AgentFinalStreamHandler struct. func NewFinalStreamHandler(keywords ...string) *AgentFinalStreamHandler { if len(keywords) > 0 { DefaultKeywords = keywords @@ -32,10 +42,20 @@ func NewFinalStreamHandler(keywords ...string) *AgentFinalStreamHandler { } } +// GetEgress returns the egress channel of the AgentFinalStreamHandler. +// +// It does not take any parameters. +// It returns a channel of type []byte. func (handler *AgentFinalStreamHandler) GetEgress() chan []byte { return handler.egress } +// ReadFromEgress reads data from the egress channel and invokes the provided +// callback function with each chunk of data. +// +// The callback function receives two parameters: +// - ctx: the context.Context object for the egress operation. +// - chunk: a byte slice representing a chunk of data from the egress channel. func (handler *AgentFinalStreamHandler) ReadFromEgress(callback func(ctx context.Context, chunk []byte)) { egressCtx, cancel := context.WithCancel(handler.ctx) defer cancel() @@ -48,6 +68,13 @@ func (handler *AgentFinalStreamHandler) ReadFromEgress(callback func(ctx context }() } +// HandleStreamingFunc implements the callback interface that handles the streaming +// of data in the AgentFinalStreamHandler. The handler reads the incoming data and checks for the +// agents final output keywords, ie, `Final Answer:`, `Final:`, `AI:`. Upon detection of +// the keyword, it starst to stream the agents final output to the egress channel. +// +// It takes in the context and a chunk of bytes as parameters. +// There is no return type for this function. func (handler *AgentFinalStreamHandler) HandleStreamingFunc(ctx context.Context, chunk []byte) { handler.ctx = ctx chunkStr := string(chunk)