Skip to content

Commit

Permalink
feat: add user handle to prompt
Browse files Browse the repository at this point in the history
  • Loading branch information
davidramiro committed Feb 9, 2025
1 parent cd61272 commit c476bb6
Show file tree
Hide file tree
Showing 5 changed files with 27 additions and 6 deletions.
3 changes: 3 additions & 0 deletions config.sample.toml
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,9 @@ bot_token = "4242:telegram-bot-token"
api_key = "sk-api-key"
system_prompt = '''
You are HSBot, a helpful assistant.
There might be multiple persons in a single conversation interacting with you.
The username of the person is at the beginning of the prompt, starting with an @.
If there is no @, it's just the first name of the user. You can directly address a person if you got provided with an @ handle.
'''

[fal]
Expand Down
15 changes: 14 additions & 1 deletion internal/adapters/handler/command.go
Original file line number Diff line number Diff line change
Expand Up @@ -48,9 +48,13 @@ func (h *CommandHandler) Handle(ctx context.Context, b *bot.Bot, update *models.
}
if update.Message.ReplyToMessage.From.ID == botUser.ID {
isReplyToBot = true
quotedText = update.Message.ReplyToMessage.Text
} else {
quotedText = getUserNameFromMessage(update.Message.ReplyToMessage.From) + ": " +
update.Message.ReplyToMessage.Text
}

*replyToMessageID = update.Message.ReplyToMessage.ID
quotedText = update.Message.ReplyToMessage.Text
}

imageURL := make(chan string)
Expand All @@ -64,6 +68,7 @@ func (h *CommandHandler) Handle(ctx context.Context, b *bot.Bot, update *models.
ID: update.Message.ID,
ChatID: update.Message.Chat.ID,
Text: update.Message.Text,
Username: getUserNameFromMessage(update.Message.From),
ReplyToMessageID: replyToMessageID,
IsReplyToBot: isReplyToBot,
QuotedText: quotedText,
Expand Down Expand Up @@ -147,3 +152,11 @@ func findLargestImage(photos []models.PhotoSize) string {

return maxID
}

func getUserNameFromMessage(user *models.User) string {
if user.Username == "" {
return user.FirstName
}

return "@" + user.Username
}
4 changes: 3 additions & 1 deletion internal/core/domain/commands/chat.go
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,8 @@ func (h *ChatHandler) Respond(ctx context.Context, timeout time.Duration, messag
promptText += ": " + transcript
}

promptText = message.Username + ": " + promptText

h.mutex.Lock()
defer h.mutex.Unlock()

Expand All @@ -99,7 +101,7 @@ func (h *ChatHandler) Respond(ctx context.Context, timeout time.Duration, messag
}

conversation.messages = append(conversation.messages, domain.Prompt{Author: domain.User,
Prompt: message.Text})
Prompt: promptText})
} else {
conversation.messages = append(conversation.messages, domain.Prompt{Author: domain.User,
Prompt: promptText, ImageURL: message.ImageURL})
Expand Down
10 changes: 6 additions & 4 deletions internal/core/domain/commands/chat_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -64,21 +64,23 @@ func TestChatHandlerCache(t *testing.T) {

assert.NotNil(t, chatHandler)

err := chatHandler.Respond(context.Background(), time.Minute, &domain.Message{ChatID: 1, ID: 1, Text: "/chat prompt"})
err := chatHandler.Respond(context.Background(), time.Minute, &domain.Message{
ChatID: 1, ID: 1, Username: "@unit", Text: "/chat prompt"})

require.NoError(t, err)
assert.Equal(t, "mock response", ms.Message)
assert.Len(t, chatHandler.cache, 1)

err = chatHandler.Respond(context.Background(), time.Minute, &domain.Message{ChatID: 1, ID: 2, Text: "/chat prompt2"})
err = chatHandler.Respond(context.Background(), time.Minute, &domain.Message{
ChatID: 1, ID: 2, Username: "@unit", Text: "/chat prompt2"})
require.NoError(t, err)
assert.Len(t, chatHandler.cache, 1)

assert.Len(t, chatHandler.cache[1].messages, 4)

assert.Equal(t, "prompt", chatHandler.cache[1].messages[0].Prompt)
assert.Equal(t, "@unit: prompt", chatHandler.cache[1].messages[0].Prompt)
assert.Equal(t, "mock response", chatHandler.cache[1].messages[1].Prompt)
assert.Equal(t, "prompt2", chatHandler.cache[1].messages[2].Prompt)
assert.Equal(t, "@unit: prompt2", chatHandler.cache[1].messages[2].Prompt)
assert.Equal(t, "mock response", chatHandler.cache[1].messages[3].Prompt)
}

Expand Down
1 change: 1 addition & 0 deletions internal/core/domain/model.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ type Prompt struct {
type Message struct {
ID int
ChatID int64
Username string
ReplyToMessageID *int
IsReplyToBot bool
QuotedText string
Expand Down

0 comments on commit c476bb6

Please sign in to comment.