Merge pull request #11 from mutablelogic/dev

Added telegram as a test

Author: djthorpe

cmd/llm/chat2.go

@@ -0,0 +1,160 @@
+package main
+
+import (
+	"context"
+	"errors"
+	"log"
+	"sync"
+	"time"
+
+	// Packages
+	llm "github.com/mutablelogic/go-llm"
+	telegram "github.com/mutablelogic/go-llm/pkg/ui/telegram"
+)
+
+////////////////////////////////////////////////////////////////////////////////
+// TYPES
+
+type Chat2Cmd struct {
+	Model string `arg:"" help:"Model name"`
+	Token string `env:"TELEGRAM_TOKEN" help:"Telegram token" required:""`
+}
+
+type Server struct {
+	sync.RWMutex
+	*telegram.Client
+
+	// Model and toolkit
+	toolkit llm.ToolKit
+	model   llm.Model
+
+	// Map of active sessions
+	sessions map[string]llm.Context
+}
+
+////////////////////////////////////////////////////////////////////////////////
+// LIFECYCLE
+
+func NewTelegramServer(token string, model llm.Model, toolkit llm.ToolKit, opts ...telegram.Opt) (*Server, error) {
+	server := new(Server)
+	server.sessions = make(map[string]llm.Context)
+	server.model = model
+	server.toolkit = toolkit
+
+	// Create a new telegram client
+	opts = append(opts, telegram.WithCallback(server.receive))
+	if telegram, err := telegram.New(token, opts...); err != nil {
+		return nil, err
+	} else {
+		server.Client = telegram
+	}
+
+	// Return success
+	return server, nil
+}
+
+////////////////////////////////////////////////////////////////////////////////
+// PUBLIC METHODS
+
+func (cmd *Chat2Cmd) Run(globals *Globals) error {
+	return run(globals, cmd.Model, func(ctx context.Context, model llm.Model) error {
+		server, err := NewTelegramServer(cmd.Token, model, globals.toolkit, telegram.WithDebug(globals.Debug))
+		if err != nil {
+			return err
+		}
+
+		log.Printf("Running Telegram bot %q\n", server.Client.Name())
+
+		var result error
+		var wg sync.WaitGroup
+		wg.Add(2)
+		go func(ctx context.Context) {
+			defer wg.Done()
+			if err := server.Run(ctx); err != nil {
+				result = errors.Join(result, err)
+			}
+		}(ctx)
+		go func(ctx context.Context) {
+			defer wg.Done()
+			ticker := time.NewTicker(time.Minute)
+			defer ticker.Stop()
+			for {
+				select {
+				case <-ctx.Done():
+					return
+				case <-ticker.C:
+					server.Purge()
+				}
+			}
+		}(ctx)
+
+		// Wait for completion
+		wg.Wait()
+
+		// Return any errors
+		return result
+	})
+}
+
+// //////////////////////////////////////////////////////////////////////////////
+// PRIVATE METHODS
+
+func (telegram *Server) Purge() {
+	telegram.Lock()
+	defer telegram.Unlock()
+	for user, session := range telegram.sessions {
+		if session.SinceLast() > 10*time.Minute {
+			log.Printf("Purging session for %q\n", user)
+			delete(telegram.sessions, user)
+		}
+	}
+}
+
+func (telegram *Server) session(user string) llm.Context {
+	telegram.Lock()
+	defer telegram.Unlock()
+	if session, exists := telegram.sessions[user]; exists {
+		return session
+	}
+	session := telegram.model.Context(
+		llm.WithToolKit(telegram.toolkit),
+		llm.WithSystemPrompt("Please reply to messages in markdown format."),
+	)
+	telegram.sessions[user] = session
+	return session
+}
+
+func (telegram *Server) receive(ctx context.Context, msg telegram.Message) error {
+	// Get an active session
+	session := telegram.session(msg.Sender())
+
+	// Process the message
+	text := msg.Text()
+	text += "\n\nPlease reply in markdown format."
+	if err := session.FromUser(ctx, text); err != nil {
+		return err
+	}
+
+	// Run tool calls
+	for {
+		calls := session.ToolCalls(0)
+		if len(calls) == 0 {
+			break
+		}
+		if text := session.Text(0); text != "" {
+			msg.Reply(ctx, text, false)
+		} else {
+			msg.Reply(ctx, "_Gathering information_", true)
+		}
+
+		results, err := telegram.toolkit.Run(ctx, calls...)
+		if err != nil {
+			return err
+		} else if err := session.FromTool(ctx, results...); err != nil {
+			return err
+		}
+	}
+
+	// Reply with the text
+	return msg.Reply(ctx, session.Text(0), true)
+}

cmd/llm/main.go

@@ -78,6 +78,7 @@ type CLI struct {
 	// Commands
 	Download  DownloadModelCmd `cmd:"" help:"Download a model"`
 	Chat      ChatCmd          `cmd:"" help:"Start a chat session"`
+	Chat2     Chat2Cmd         `cmd:"" help:"Start a chat session (2)"`
 	Complete  CompleteCmd      `cmd:"" help:"Complete a prompt"`
 	Embedding EmbeddingCmd     `cmd:"" help:"Generate an embedding"`
 	Version   VersionCmd       `cmd:"" help:"Print the version of this tool"`

context.go

@@ -1,6 +1,9 @@
 package llm
 
-import "context"
+import (
+	"context"
+	"time"
+)
 
 //////////////////////////////////////////////////////////////////
 // TYPES
@@ -44,4 +47,8 @@ type Context interface {
 	// Generate a response from a tool, passing the results
 	// from the tool call
 	FromTool(context.Context, ...ToolResult) error
+
+	// Return the duration since the last completion was made
+	// or zero
+	SinceLast() time.Duration
 }

pkg/internal/impl/session.go

@@ -3,6 +3,7 @@ package impl
 import (
 	"context"
 	"encoding/json"
+	"time"
 
 	// Packages
 	"github.com/mutablelogic/go-llm"
@@ -37,6 +38,7 @@ type session struct {
 	opts    []llm.Opt        // Options to apply to the session
 	seq     []llm.Completion // Sequence of messages
 	factory MessageFactory   // Factory for generating messages
+	last    time.Time        // Last completion time
 }
 
 var _ llm.Context = (*session)(nil)
@@ -129,10 +131,20 @@ func (session *session) chat(ctx context.Context, messages ...llm.Completion) er
 	// Append the first choice
 	session.Append(completion.Choice(0))
 
+	// Update the last completion time
+	session.last = time.Now()
+
 	// Success
 	return nil
 }
 
+func (session *session) SinceLast() time.Duration {
+	if len(session.seq) == 0 || session.last.IsZero() {
+		return 0
+	}
+	return time.Since(session.last)
+}
+
 ///////////////////////////////////////////////////////////////////////////////
 // PUBLIC METHODS - COMPLETION
 

pkg/ui/telegram/opt.go

@@ -0,0 +1,50 @@
+package telegram
+
+import "context"
+
+///////////////////////////////////////////////////////////////////////////////
+// TYPES
+
+// A generic option type, which can set options on an agent or session
+type Opt func(*opts) error
+
+// set of options
+type opts struct {
+	token    string
+	callback CallbackFunc
+	debug    bool
+}
+
+type CallbackFunc func(context.Context, Message) error
+
+////////////////////////////////////////////////////////////////////////////////
+// LIFECYCLE
+
+// applyOpts returns a structure of options
+func applyOpts(token string, opt ...Opt) (*opts, error) {
+	o := new(opts)
+	o.token = token
+	for _, opt := range opt {
+		if err := opt(o); err != nil {
+			return nil, err
+		}
+	}
+	return o, nil
+}
+
+////////////////////////////////////////////////////////////////////////////////
+// PUBLIC METHODS
+
+func WithCallback(fn CallbackFunc) Opt {
+	return func(o *opts) error {
+		o.callback = fn
+		return nil
+	}
+}
+
+func WithDebug(v bool) Opt {
+	return func(o *opts) error {
+		o.debug = v
+		return nil
+	}
+}