Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
160 changes: 160 additions & 0 deletions cmd/llm/chat2.go
Original file line number Diff line number Diff line change
@@ -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)
}
1 change: 1 addition & 0 deletions cmd/llm/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"`
Expand Down
9 changes: 8 additions & 1 deletion context.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,9 @@
package llm

import "context"
import (
"context"
"time"
)

//////////////////////////////////////////////////////////////////
// TYPES
Expand Down Expand Up @@ -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
}
12 changes: 12 additions & 0 deletions pkg/internal/impl/session.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package impl
import (
"context"
"encoding/json"
"time"

// Packages
"github.com/mutablelogic/go-llm"
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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

Expand Down
50 changes: 50 additions & 0 deletions pkg/ui/telegram/opt.go
Original file line number Diff line number Diff line change
@@ -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
}
}
Loading