refactor: objectify both gobot and rustbot

gobot/jet/reader.go

@@ -0,0 +1,54 @@
+package jet
+
+import (
+	"context"
+	"fmt"
+	"log"
+
+	"github.com/nats-io/nats.go/jetstream"
+)
+
+type Reader struct {
+	cons    jetstream.Consumer
+	subject string
+}
+
+func (r Reader) Read(ctx context.Context, prompts chan string) error {
+	log.Println("launching JetStream Writer")
+	defer log.Println("done reading from JetStream")
+	iter, err := r.cons.Messages()
+	if err != nil {
+		return err
+	}
+	go func() {
+		<-ctx.Done()
+		iter.Stop()
+	}()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+			msg, err := iter.Next()
+			if err != nil {
+				// NOTE: we are handling ErrMsgIteratorClosed
+				// because the only way the iterator is closed
+				// is if we do it when the context is cancelled.
+				if err == jetstream.ErrMsgIteratorClosed {
+					return nil
+				}
+				return err
+			}
+			fmt.Printf("\n[Q]: %s\n", string(msg.Data()))
+			if err := msg.Ack(); err != nil {
+				return err
+			}
+			select {
+			case <-ctx.Done():
+				return ctx.Err()
+			case prompts <- string(msg.Data()):
+			}
+		}
+	}
+}

gobot/jet/stream.go

@@ -0,0 +1,69 @@
+package jet
+
+import (
+	"context"
+	"errors"
+	"fmt"
+
+	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
+)
+
+type Config struct {
+	StreamURL   string
+	StreamName  string
+	DurableName string
+	PubSubject  string
+	SubSubject  string
+}
+
+type Stream struct {
+	Writer Writer
+	Reader Reader
+}
+
+func NewStream(ctx context.Context, c Config) (*Stream, error) {
+	nc, err := nats.Connect(c.StreamURL)
+	if err != nil {
+		return nil, fmt.Errorf("failed connecting to NATS: %v", err)
+	}
+
+	js, err := jetstream.New(nc)
+	if err != nil {
+		return nil, fmt.Errorf("failed creating a new stream: %v", err)
+	}
+
+	stream, err := js.CreateStream(ctx, jetstream.StreamConfig{
+		Name:     c.StreamName,
+		Subjects: []string{c.SubSubject, c.PubSubject},
+	})
+	if err != nil {
+		if !errors.Is(err, jetstream.ErrStreamNameAlreadyInUse) {
+			return nil, err
+		}
+		var jsErr error
+		stream, jsErr = js.Stream(ctx, c.StreamName)
+		if jsErr != nil {
+			return nil, err
+		}
+	}
+
+	cons, err := stream.CreateOrUpdateConsumer(ctx, jetstream.ConsumerConfig{
+		Durable:       c.DurableName,
+		FilterSubject: c.SubSubject,
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	return &Stream{
+		Writer: Writer{
+			stream:  js,
+			subject: c.PubSubject,
+		},
+		Reader: Reader{
+			cons:    cons,
+			subject: c.SubSubject,
+		},
+	}, nil
+}

gobot/jet/writer.go

@@ -0,0 +1,38 @@
+package jet
+
+import (
+	"context"
+	"fmt"
+	"log"
+
+	"github.com/nats-io/nats.go/jetstream"
+)
+
+type Writer struct {
+	stream  jetstream.JetStream
+	subject string
+}
+
+func (w Writer) Write(ctx context.Context, chunks chan []byte) error {
+	log.Println("launching JetStream Reader")
+	defer log.Println("done writing to JetStream")
+	msg := []byte{}
+	for {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case chunk := <-chunks:
+			if len(chunk) == 0 {
+				fmt.Printf("\n[A]: %s\n", string(msg))
+				_, err := w.stream.Publish(ctx, w.subject, msg)
+				if err != nil {
+					return err
+				}
+				// reset the msg slice instead of reallocating
+				msg = msg[:0]
+				break
+			}
+			msg = append(msg, chunk...)
+		}
+	}
+}

gobot/llm/llm.go

@@ -33,21 +33,21 @@ func New(c Config) (*LLM, error) {
 	}, nil
 }
 
-func Stream(ctx context.Context, llm *LLM, prompts chan string, chunks chan []byte) error {
+func (l *LLM) Stream(ctx context.Context, prompts chan string, chunks chan []byte) error {
 	log.Println("launching LLM stream")
 	defer log.Println("done streaming LLM")
-	chat := NewHistory(int(llm.histSize))
-	chat.Add(llm.seedPrompt)
+	chat := NewHistory(int(l.histSize))
+	chat.Add(l.seedPrompt)
 
-	fmt.Println("Seed prompt: ", llm.seedPrompt)
+	fmt.Println("Seed prompt: ", l.seedPrompt)
 
 	for {
 		select {
 		case <-ctx.Done():
 			return ctx.Err()
 		case prompt := <-prompts:
 			chat.Add(prompt)
-			_, err := llms.GenerateFromSinglePrompt(ctx, llm.model, chat.String(),
+			_, err := llms.GenerateFromSinglePrompt(ctx, l.model, chat.String(),
 				llms.WithStreamingFunc(func(_ context.Context, chunk []byte) error {
 					select {
 					case <-ctx.Done():