Compare commits

..
7 Commits
Author SHA1 Message Date
HugeFrog24 44fcd02d9a Style and security 2026-03-05 04:50:36 +01:00
HugeFrog24 d8d0da4704 Upgrade dependencies
Added tests, revised logging

Removed dependency on env file

Try reformatting unit file

Comments clarification

Added readme

Added readme
2024-10-23 22:06:56 +02:00
HugeFrog24 c8af457af1 MVP
md formatting doesnt work yet

Started implementing owner feature

Add .gitattributes to enforce LF line endings

Temporary commit before merge

Updated owner management

Updated json and gitignore

Proceed with role management

Again, CI

Fix some lint errors

Implemented screening

Per-bot API keys implemented

Use getRoleByName func

Fix unused imports

Upgrade actions

rm unused function

Upgrade action

Fix unaddressed errors
2024-10-23 22:06:55 +02:00
HugeFrog24 e5532df7f9 Handle business messages 2024-10-20 15:52:41 +02:00
HugeFrog24 0ab56448c7 Multibot finished 2024-10-13 16:41:03 +02:00
HugeFrog24 9f2b3df4c8 Concern separation 2024-10-13 02:58:18 +02:00
HugeFrog24 41c9b8075b Created user-role system 2024-10-13 01:36:56 +02:00
33 changed files with 885 additions and 1960 deletions
+14
View File
@@ -0,0 +1,14 @@
---
alwaysApply: true
always_on: true
trigger: always_on
applyTo: "**"
description: Snyk Security At Inception
---
# Project security best practices
- Always run snyk_code_scan tool for new first party code that is generated in a Snyk-supported language.
- If any security issues are found based on newly introduced or modified code or dependencies, attempt to fix the issues using the results context from Snyk.
- Rescan the code after fixing the issues to ensure that the issues were fixed and that there are no newly introduced issues.
- Repeat this process until no new issues are found.
+23 -2
View File
@@ -7,18 +7,38 @@ on:
branches: [ main ]
jobs:
# Common setup job that other jobs can depend on
setup:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- uses: actions/setup-go@v5
with:
go-version: '1.26.0'
- uses: actions/cache@v4
with:
path: |
~/.cache/go-build
~/go/pkg/mod
key: ${{ runner.os }}-go-${{ hashFiles('**/go.sum') }}
restore-keys: |
${{ runner.os }}-go-
- run: go mod tidy
# Lint job
lint:
needs: setup
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v6
- uses: golangci/golangci-lint-action@v9
with:
version: v2.12.2
version: v2.10
args: --timeout 5m
# Test job
test:
needs: setup
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v6
@@ -29,9 +49,10 @@ jobs:
# Security scan job
security:
needs: setup
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v6
- uses: securego/gosec@v2.26.1
- uses: securego/gosec@master
with:
args: ./...
+1 -1
View File
@@ -15,4 +15,4 @@ bot.db
# All config files except for the default
config/*
!config/default.json
!config/default.json
+3
View File
@@ -0,0 +1,3 @@
{
"mcpServers": {}
}
+1 -36
View File
@@ -4,8 +4,7 @@ A scalable, multi-bot solution for Telegram using Go, GORM, and the Anthropic AP
## Design Considerations
- AI-powered (Anthropic Claude)
- Voice message support (ElevenLabs STT + TTS) — optional, enabled per bot via config
- AI-powered
- Supports multiple bot profiles
- Uses SQLite for persistence
- Implements rate limiting and user management
@@ -52,40 +51,6 @@ A scalable, multi-bot solution for Telegram using Go, GORM, and the Anthropic AP
go build -o telegram-bot
```
## Trying Out New Behavior Safely
Want to experiment with a different personality, tone, or set of instructions without disturbing the bot your users already talk to? Run a second, separate bot just for testing.
Each bot profile is its own config file with its own Telegram token, and bots are fully independent — separate identity, separate chat history, separate settings. So a "test twin" is quick to set up:
1. Create a new bot with [@BotFather](https://t.me/BotFather) and copy its token.
2. Copy your existing config to a new file, e.g. `cp config/mybot.json config/mybot-test.json`.
3. In the new file, paste the new token, give it a different `id`, and edit `system_prompts` to try your changes.
4. Start it alongside your main bot. Chat with the test bot, tweak its prompt, and restart the test bot to try again — your real users never see the experiments.
5. Happy with the result? Copy the same change into your main bot's config and restart it.
> [!NOTE]
> A test bot always needs its **own** token. Telegram only lets one running bot listen on a given token, so you can't point a second copy at your live bot — give the twin its own @BotFather bot instead.
## Configuration
Each bot is one JSON file in `config/` (see `config/default.json` for the template). Keys of note:
| Key | Type | Default | Description |
| ------------------ | ------ | ----------- | ----------- |
| `max_tokens` | number | `1000` | Maximum output tokens per reply. **Thinking tokens count toward this limit** — raise it (e.g. `4000`+) whenever `thinking` is `"adaptive"`, or a turn can spend the whole budget on reasoning and produce no text. |
| `thinking` | string | *(omitted)* | Reasoning mode: `"adaptive"` (the model decides when and how much to think) or `"disabled"`. Omit the key entirely to use the model's own API default. If the configured model doesn't support the chosen mode, the API rejects the request with a 400 — owners/admins see the raw error, regular users get the generic fallback. Check [Anthropic's model docs](https://platform.claude.com/docs/en/build-with-claude/adaptive-thinking) for per-model support. |
| `thinking_display` | string | *(omitted)* | `"summarized"` or `"omitted"`. Only valid together with `"thinking": "adaptive"`. Controls whether the API returns a readable summary of the reasoning (logged, never sent to chat). Thinking is billed the same either way; when omitted, the API's per-model default applies. |
Every reply logs one accounting line — `[usage] model=... in=... out=... thinking=... stop=...` — so thinking spend (billed even when its text is omitted) stays visible in `journalctl`/`docker compose logs`. A `stop=max_tokens` line is accompanied by an error-level warning that the reply was truncated.
> [!TIP]
> For deep request/response debugging, the Anthropic Go SDK ships `option.WithDebugLog(...)` (dumps full HTTP bodies with auth headers redacted). It is not wired into the bot — dev-only, add it temporarily to the client constructor if you ever need wire-level traces.
### Future: persistent memory
The Anthropic memory tool (`memory_20250818`) is a candidate future feature for cross-conversation recall (a self-hosted analog of ChatGPT's "memory"). The Go SDK already ships the types (`BetaMemoryTool20250818Param` and its tool-union slot plus the six-command union: `view`/`create`/`str_replace`/`insert`/`delete`/`rename`), but — unlike the Python/TypeScript/Java SDKs — provides **no handler helper**: the bot would have to hand-write client-side command dispatch against per-chat storage, including strict path validation (canonicalize and confine every model-supplied path under a fixed memory root; reject `..`/symlink traversal) and a no-secrets policy for stored content. Not implemented yet.
## Systemd Unit Setup
To enable the bot to start automatically on system boot and run in the background, set up a systemd unit.
-88
View File
@@ -1,88 +0,0 @@
package main
import (
"context"
"sort"
"time"
"github.com/go-telegram/bot/models"
)
const albumFlushWindow = 1 * time.Second
type pendingAlbum struct {
items []*models.Message
chatID, userID int64
username, firstName, lastName, languageCode string
isPremium bool
messageTime int
businessConnectionID string
timer *time.Timer
}
func (b *Bot) bufferAlbumItem(
ctx context.Context,
msg *models.Message,
chatID, userID int64,
username, firstName, lastName string,
isPremium bool,
languageCode string,
messageTime int,
businessConnectionID string,
) {
b.albumBuffersMu.Lock()
defer b.albumBuffersMu.Unlock()
album, exists := b.albumBuffers[msg.MediaGroupID]
if !exists {
album = &pendingAlbum{
chatID: chatID,
userID: userID,
username: username,
firstName: firstName,
lastName: lastName,
isPremium: isPremium,
languageCode: languageCode,
messageTime: messageTime,
businessConnectionID: businessConnectionID,
}
b.albumBuffers[msg.MediaGroupID] = album
}
album.items = append(album.items, msg)
if album.timer != nil {
album.timer.Stop()
}
mediaGroupID := msg.MediaGroupID
album.timer = time.AfterFunc(albumFlushWindow, func() {
b.flushAlbum(ctx, mediaGroupID)
})
}
func (b *Bot) flushAlbum(ctx context.Context, mediaGroupID string) {
b.albumBuffersMu.Lock()
album, exists := b.albumBuffers[mediaGroupID]
if !exists {
b.albumBuffersMu.Unlock()
return
}
delete(b.albumBuffers, mediaGroupID)
items := album.items
captured := *album
b.albumBuffersMu.Unlock()
sort.Slice(items, func(i, j int) bool { return items[i].ID < items[j].ID })
if !b.checkRateLimits(captured.userID) {
b.sendRateLimitExceededMessage(ctx, captured.chatID, captured.businessConnectionID)
return
}
b.handlePhotoMessage(
ctx, items,
captured.chatID, captured.userID,
captured.username, captured.firstName, captured.lastName,
captured.isPremium, captured.languageCode, captured.messageTime,
captured.businessConnectionID,
)
}
+106 -302
View File
@@ -4,333 +4,137 @@ import (
"context"
"errors"
"fmt"
"net/http"
"strings"
"sync/atomic"
"time"
"github.com/anthropics/anthropic-sdk-go"
"github.com/anthropics/anthropic-sdk-go/packages/param"
"github.com/liushuangls/go-anthropic/v2"
)
// ErrModelNotFound is returned when the configured Anthropic model is no longer available
// (deprecated or removed). Callers can use errors.Is to detect this and surface an
// actionable message to admins/owners while keeping the response vague for regular users.
var ErrModelNotFound = errors.New("model not found or deprecated")
const maxFileNotFoundRetries = 3
func (b *Bot) getAnthropicResponse(ctx context.Context, messages []anthropic.Message, isNewChat, isOwner, isEmojiOnly bool, username string, firstName string, lastName string, isPremium bool, languageCode string, messageTime int) (string, error) {
// Use prompts from config
var systemMessage string
if isNewChat {
systemMessage = b.config.SystemPrompts["new_chat"]
} else {
systemMessage = b.config.SystemPrompts["continue_conversation"]
}
const defaultMaxTokens = 1000
// Combine default prompt with custom instructions
systemMessage = b.config.SystemPrompts["default"] + " " + b.config.SystemPrompts["custom_instructions"] + " " + systemMessage
const mcpUnsupportedSentinel = "format not currently supported by the Anthropic API"
// Handle username placeholder
usernameValue := username
if username == "" {
usernameValue = "unknown" // Use "unknown" when username is not available
}
systemMessage = strings.ReplaceAll(systemMessage, "{username}", usernameValue)
var mcpUnsupportedCount atomic.Uint64
// Handle firstname placeholder
firstnameValue := firstName
if firstName == "" {
firstnameValue = "unknown" // Use "unknown" when first name is not available
}
systemMessage = strings.ReplaceAll(systemMessage, "{firstname}", firstnameValue)
type mcpCall struct{ server, name, input string }
// Handle lastname placeholder
lastnameValue := lastName
if lastName == "" {
lastnameValue = "" // Empty string when last name is not available
}
systemMessage = strings.ReplaceAll(systemMessage, "{lastname}", lastnameValue)
func (b *Bot) getAnthropicResponse(ctx context.Context, chatID int64, messages []anthropic.BetaMessageParam, isEmojiOnly bool, username string, firstName string, lastName string, isPremium bool, languageCode string, messageTime int, onSegment func(string) error) (string, error) {
staticPrompt := strings.TrimSpace(b.config.SystemPrompts["custom_instructions"])
// Handle language code placeholder
langValue := languageCode
if languageCode == "" {
langValue = "en" // Default to English when language code is not available
}
systemMessage = strings.ReplaceAll(systemMessage, "{language}", langValue)
// Handle premium status
premiumStatus := "regular user"
if isPremium {
premiumStatus = "premium user"
}
systemMessage = strings.ReplaceAll(systemMessage, "{premium_status}", premiumStatus)
// Handle time awareness
timeObj := time.Unix(int64(messageTime), 0)
hour := timeObj.Hour()
var timeContext string
if hour >= 5 && hour < 12 {
timeContext = "morning"
} else if hour >= 12 && hour < 18 {
timeContext = "afternoon"
} else if hour >= 18 && hour < 22 {
timeContext = "evening"
} else {
timeContext = "night"
}
systemMessage = strings.ReplaceAll(systemMessage, "{time_context}", timeContext)
if !isOwner {
systemMessage += " " + b.config.SystemPrompts["avoid_sensitive"]
}
if isEmojiOnly {
systemMessage += " " + b.config.SystemPrompts["respond_with_emojis"]
}
// Debug logging
InfoLogger.Printf("Sending %d messages to Anthropic", len(messages))
maxTokens := int64(defaultMaxTokens)
if b.config.MaxTokens > 0 {
maxTokens = int64(b.config.MaxTokens)
for i, msg := range messages {
for _, content := range msg.Content {
if content.Type == anthropic.MessagesContentTypeText {
InfoLogger.Printf("Message %d: Role=%v, Text=%v", i, msg.Role, content.Text)
}
}
}
params := anthropic.BetaMessageNewParams{
Model: b.config.Model,
MaxTokens: maxTokens,
// Ensure the roles are correct
for i := range messages {
switch messages[i].Role {
case anthropic.RoleUser:
messages[i].Role = anthropic.RoleUser
case anthropic.RoleAssistant:
messages[i].Role = anthropic.RoleAssistant
default:
// Default to 'user' if role is unrecognized
messages[i].Role = anthropic.RoleUser
}
}
model := anthropic.Model(b.config.Model)
// Create the request
request := anthropic.MessagesRequest{
Model: model, // Now `model` is of type anthropic.Model
Messages: messages,
Betas: []anthropic.AnthropicBeta{anthropic.AnthropicBetaFilesAPI2025_04_14},
}
if staticPrompt != "" {
blocks := []anthropic.BetaTextBlockParam{
{Text: staticPrompt, CacheControl: anthropic.NewBetaCacheControlEphemeralParam()},
}
tail := buildUserContext(username, firstName, lastName, isPremium, languageCode, messageTime)
if isEmojiOnly {
if rule := strings.TrimSpace(b.config.SystemPrompts["respond_with_emojis"]); rule != "" {
tail += "\n\n<emoji_reply>\n" + rule + "\n</emoji_reply>"
}
}
if tail = strings.TrimSpace(tail); tail != "" {
blocks = append(blocks, anthropic.BetaTextBlockParam{Text: tail})
}
params.System = blocks
System: systemMessage,
MaxTokens: 1000,
}
// Apply temperature if set in config
if b.config.Temperature != nil {
params.Temperature = param.NewOpt(float64(*b.config.Temperature))
request.Temperature = b.config.Temperature
}
if thinking, ok := thinkingParamFromConfig(b.config.Thinking, b.config.ThinkingDisplay); ok {
params.Thinking = thinking
}
if len(b.config.MCPServers) > 0 {
mcpServers := make([]anthropic.BetaRequestMCPServerURLDefinitionParam, 0, len(b.config.MCPServers))
tools := make([]anthropic.BetaToolUnionParam, 0, len(b.config.MCPServers))
for _, s := range b.config.MCPServers {
srv := anthropic.BetaRequestMCPServerURLDefinitionParam{
Name: s.Name,
URL: s.URL,
}
if s.AuthorizationToken != "" {
srv.AuthorizationToken = param.NewOpt(s.AuthorizationToken)
}
mcpServers = append(mcpServers, srv)
toolset := &anthropic.BetaMCPToolsetParam{
MCPServerName: s.Name,
}
if len(s.AllowedTools) > 0 {
toolset.DefaultConfig = anthropic.BetaMCPToolDefaultConfigParam{
Enabled: param.NewOpt(false),
}
toolset.Configs = make(map[string]anthropic.BetaMCPToolConfigParam, len(s.AllowedTools))
for _, tool := range s.AllowedTools {
toolset.Configs[tool] = anthropic.BetaMCPToolConfigParam{
Enabled: param.NewOpt(true),
}
}
}
tools = append(tools, anthropic.BetaToolUnionParam{OfMCPToolset: toolset})
}
params.MCPServers = mcpServers
params.Tools = tools
params.Betas = append(params.Betas, anthropic.AnthropicBetaMCPClient2025_11_20)
}
for attempt := 0; attempt < maxFileNotFoundRetries; attempt++ {
joined, streamErr := b.streamMessages(ctx, params, onSegment)
if streamErr == nil {
return joined, nil
}
var apiErr *anthropic.Error
if !errors.As(streamErr, &apiErr) || apiErr.StatusCode != http.StatusNotFound {
return "", fmt.Errorf("error creating Anthropic message: %w", streamErr)
}
missingFileID := extractMissingFileID(streamErr)
if missingFileID == "" {
resp, err := b.anthropicClient.CreateMessages(ctx, request)
if err != nil {
var apiErr *anthropic.APIError
if errors.As(err, &apiErr) && apiErr.IsNotFoundErr() {
return "", fmt.Errorf("%w: %s", ErrModelNotFound, b.config.Model)
}
ErrorLogger.Printf("[%s] self-heal: stripping dead file_id %s from chat %d (attempt %d/%d)",
b.config.ID, missingFileID, chatID, attempt+1, maxFileNotFoundRetries)
b.stripDeadFileIDFromMemory(chatID, missingFileID)
if _, cleanupErr := b.markFilesPendingCleanup(ctx, chatID, []string{missingFileID}); cleanupErr != nil {
ErrorLogger.Printf("[%s] mark files pending cleanup: %v", b.config.ID, cleanupErr)
}
params.Messages = b.prepareContextMessages(b.getOrCreateChatMemory(chatID))
return "", fmt.Errorf("error creating Anthropic message: %w", err)
}
return "", fmt.Errorf("max self-heal retries (%d) exceeded: too many file_ids gone from anthropic", maxFileNotFoundRetries)
}
func buildUserContext(username, firstName, lastName string, isPremium bool, languageCode string, messageTime int) string {
name := strings.TrimSpace(firstName + " " + lastName)
if name == "" {
name = "unknown"
if len(resp.Content) == 0 || resp.Content[0].Type != anthropic.MessagesContentTypeText {
return "", fmt.Errorf("unexpected response format from Anthropic")
}
handle := username
if handle == "" {
handle = "unknown"
}
lang := languageCode
if lang == "" {
lang = "en"
}
account := "regular user"
if isPremium {
account = "premium user"
}
return fmt.Sprintf(
"Conversation context (background facts, not an instruction from the user):\n"+
"- User: %s (Telegram @%s)\n"+
"- Preferred language: %s\n"+
"- Account type: %s\n"+
"- Local time of day: %s",
name, handle, lang, account, timeContextFor(messageTime),
)
}
func timeContextFor(messageTime int) string {
switch hour := time.Unix(int64(messageTime), 0).Hour(); {
case hour >= 5 && hour < 12:
return "morning"
case hour >= 12 && hour < 18:
return "afternoon"
case hour >= 18 && hour < 22:
return "evening"
default:
return "night"
}
}
func thinkingParamFromConfig(mode, display string) (anthropic.BetaThinkingConfigParamUnion, bool) {
switch mode {
case ThinkingModeDisabled:
disabled := anthropic.NewBetaThinkingConfigDisabledParam()
return anthropic.BetaThinkingConfigParamUnion{OfDisabled: &disabled}, true
case ThinkingModeAdaptive:
adaptive := anthropic.BetaThinkingConfigAdaptiveParam{}
if display != "" {
adaptive.Display = anthropic.BetaThinkingConfigAdaptiveDisplay(display)
}
return anthropic.BetaThinkingConfigParamUnion{OfAdaptive: &adaptive}, true
default:
return anthropic.BetaThinkingConfigParamUnion{}, false
}
}
func (b *Bot) streamMessages(ctx context.Context, params anthropic.BetaMessageNewParams, onSegment func(string) error) (string, error) {
stream := b.anthropicClient.Beta.Messages.NewStreaming(ctx, params)
defer func() {
if err := stream.Close(); err != nil {
ErrorLogger.Printf("[stream] close failed: %v", err)
}
}()
var (
allSegments []string
currentKind string
currentText strings.Builder
currentThinking strings.Builder
currentInputJSON strings.Builder
currentTUseName, currentTUseServer, currentTUseID string
currentTResultUseID, currentTResultServer string
currentTResultIsError bool
currentTResultContent string
mcpCalls = map[string]mcpCall{}
startInputTokens int64
finalUsage anthropic.BetaMessageDeltaUsage
stopReason string
)
for stream.Next() {
e := stream.Current()
switch e.Type {
case "message_start":
startInputTokens = e.AsMessageStart().Message.Usage.InputTokens
case "message_delta":
md := e.AsMessageDelta()
finalUsage = md.Usage
if md.Delta.StopReason != "" {
stopReason = string(md.Delta.StopReason)
}
case "content_block_start":
cbs := e.AsContentBlockStart()
currentKind = cbs.ContentBlock.Type
currentText.Reset()
currentThinking.Reset()
currentInputJSON.Reset()
switch currentKind {
case "mcp_tool_use":
currentTUseName = cbs.ContentBlock.Name
currentTUseServer = cbs.ContentBlock.ServerName
currentTUseID = cbs.ContentBlock.ID
case "mcp_tool_result":
currentTResultUseID = cbs.ContentBlock.ToolUseID
currentTResultServer = cbs.ContentBlock.ServerName
currentTResultIsError = cbs.ContentBlock.IsError
currentTResultContent = cbs.ContentBlock.JSON.Content.Raw()
}
case "content_block_delta":
cbd := e.AsContentBlockDelta()
switch cbd.Delta.Type {
case "text_delta":
if currentKind == "text" {
currentText.WriteString(cbd.Delta.Text)
}
case "thinking_delta":
if currentKind == "thinking" {
currentThinking.WriteString(cbd.Delta.Thinking)
}
case "input_json_delta":
if currentKind == "mcp_tool_use" {
currentInputJSON.WriteString(cbd.Delta.PartialJSON)
}
}
case "content_block_stop":
switch currentKind {
case "text":
seg := strings.TrimSpace(currentText.String())
if seg != "" {
allSegments = append(allSegments, seg)
if onSegment != nil {
if cbErr := onSegment(seg); cbErr != nil {
ErrorLogger.Printf("[stream] onSegment failed: %v", cbErr)
}
}
}
case "mcp_tool_use":
mcpCalls[currentTUseID] = mcpCall{
server: currentTUseServer,
name: currentTUseName,
input: currentInputJSON.String(),
}
InfoLogger.Printf("[mcp] tool_use server=%q name=%q id=%q input=%s",
currentTUseServer, currentTUseName, currentTUseID, currentInputJSON.String())
case "mcp_tool_result":
preview := currentTResultContent
if len(preview) > 500 {
preview = preview[:500] + "...(truncated)"
}
InfoLogger.Printf("[mcp] tool_result tool_use_id=%q server=%q is_error=%v content=%s",
currentTResultUseID, currentTResultServer, currentTResultIsError, preview)
if strings.Contains(currentTResultContent, mcpUnsupportedSentinel) {
total := mcpUnsupportedCount.Add(1)
call := mcpCalls[currentTResultUseID]
ErrorLogger.Printf("[%s][mcp][unsupported] connector could not serialize result "+
"(total=%d): server=%q tool=%q input=%s tool_use_id=%q",
b.config.ID, total, call.server, call.name, call.input, currentTResultUseID)
}
case "thinking", "redacted_thinking":
if summary := strings.TrimSpace(currentThinking.String()); summary != "" {
if len(summary) > 500 {
summary = summary[:500] + "...(truncated)"
}
InfoLogger.Printf("[thinking] block complete: %s", summary)
} else {
InfoLogger.Printf("[thinking] block complete (content omitted)")
}
default:
if currentKind != "" {
InfoLogger.Printf("[mcp] block type=%q (unhandled)", currentKind)
}
}
currentKind = ""
}
}
if err := stream.Err(); err != nil {
return "", err
}
if stopReason != "" || finalUsage.OutputTokens > 0 {
inputTokens := finalUsage.InputTokens
if inputTokens == 0 {
inputTokens = startInputTokens
}
InfoLogger.Printf("[usage] model=%s in=%d out=%d thinking=%d stop=%s",
params.Model, inputTokens, finalUsage.OutputTokens,
finalUsage.OutputTokensDetails.ThinkingTokens, stopReason)
if stopReason == "max_tokens" {
ErrorLogger.Printf("[usage] response truncated at max_tokens=%d - raise max_tokens (thinking counts toward it)",
params.MaxTokens)
}
}
if len(allSegments) == 0 {
return "", emptyStreamError(stopReason, finalUsage.OutputTokensDetails.ThinkingTokens, params.MaxTokens)
}
return strings.Join(allSegments, "\n\n"), nil
}
func emptyStreamError(stopReason string, thinkingTokens, maxTokens int64) error {
if stopReason == "max_tokens" {
return fmt.Errorf("output budget exhausted before any text (thinking used %d of %d max_tokens) - raise max_tokens",
thinkingTokens, maxTokens)
}
return fmt.Errorf("unexpected response format from Anthropic")
return resp.Content[0].GetText(), nil
}
-172
View File
@@ -1,172 +0,0 @@
package main
import (
"bytes"
"context"
"errors"
"fmt"
"net/http"
"strings"
"time"
"github.com/anthropics/anthropic-sdk-go"
)
const fileNotFoundPrefix = "File not found: "
func formatUploadFilename(botID uint, chatID int64, tgMessageID int, ext string) string {
return fmt.Sprintf("tg-%d-%d-%d.%s", botID, chatID, tgMessageID, ext)
}
func (b *Bot) uploadImageToAnthropic(ctx context.Context, data []byte, filename, contentType string) (string, error) {
resp, err := b.anthropicClient.Beta.Files.Upload(ctx, anthropic.BetaFileUploadParams{
File: anthropic.File(bytes.NewReader(data), filename, contentType),
Betas: []anthropic.AnthropicBeta{anthropic.AnthropicBetaFilesAPI2025_04_14},
})
if err != nil {
return "", fmt.Errorf("anthropic files upload: %w", err)
}
return resp.ID, nil
}
func (b *Bot) deleteFileFromAnthropic(ctx context.Context, fileID string) error {
_, err := b.anthropicClient.Beta.Files.Delete(ctx, fileID, anthropic.BetaFileDeleteParams{
Betas: []anthropic.AnthropicBeta{anthropic.AnthropicBetaFilesAPI2025_04_14},
})
if err == nil {
return nil
}
var apiErr *anthropic.Error
if errors.As(err, &apiErr) && apiErr.StatusCode == http.StatusNotFound {
return nil
}
return fmt.Errorf("anthropic files delete %s: %w", fileID, err)
}
func (b *Bot) compensatingDelete(ctx context.Context, fileIDs []string) {
for _, fid := range fileIDs {
if err := b.deleteFileFromAnthropic(ctx, fid); err != nil {
ErrorLogger.Printf("[%s] compensating delete for %s: %v", b.config.ID, fid, err)
}
}
}
func extractMissingFileID(err error) string {
if err == nil {
return ""
}
var apiErr *anthropic.Error
if !errors.As(err, &apiErr) {
return ""
}
if apiErr.StatusCode != http.StatusNotFound {
return ""
}
return parseMissingFileIDFromBody(apiErr.RawJSON())
}
func parseMissingFileIDFromBody(raw string) string {
idx := strings.Index(raw, fileNotFoundPrefix)
if idx == -1 {
return ""
}
rest := raw[idx+len(fileNotFoundPrefix):]
end := strings.IndexFunc(rest, func(r rune) bool {
return (r < 'a' || r > 'z') &&
(r < 'A' || r > 'Z') &&
(r < '0' || r > '9') &&
r != '_'
})
if end == -1 {
return rest
}
return rest[:end]
}
func (b *Bot) hardDeleteScope(ctx context.Context, query string, args ...interface{}) error {
var rows []Message
if err := b.db.Unscoped().Where(query, args...).Find(&rows).Error; err != nil {
return fmt.Errorf("scan rows: %w", err)
}
if len(rows) == 0 {
return nil
}
if err := b.db.Where(query, args...).Delete(&Message{}).Error; err != nil {
return fmt.Errorf("soft delete: %w", err)
}
hardDeletable := make([]uint, 0, len(rows))
for _, row := range rows {
if b.deleteRowFiles(ctx, row) {
hardDeletable = append(hardDeletable, row.ID)
}
}
if len(hardDeletable) == 0 {
return nil
}
if err := b.db.Unscoped().Where("id IN ?", hardDeletable).Delete(&Message{}).Error; err != nil {
return fmt.Errorf("hard delete: %w", err)
}
return nil
}
func (b *Bot) deleteRowFiles(ctx context.Context, row Message) bool {
if len(row.ImageFileIDs) == 0 {
return true
}
allOk := true
for _, fid := range row.ImageFileIDs {
if err := b.deleteFileFromAnthropic(ctx, fid); err != nil {
ErrorLogger.Printf("[%s] anthropic delete %s (row %d): %v", b.config.ID, fid, row.ID, err)
allOk = false
}
}
return allOk
}
func stripDeadFileIDs(src []string, deadSet map[string]struct{}) (survivors []string, dirty bool) {
survivors = make([]string, 0, len(src))
for _, fid := range src {
if _, dead := deadSet[fid]; dead {
dirty = true
continue
}
survivors = append(survivors, fid)
}
return survivors, dirty
}
func (b *Bot) markFilesPendingCleanup(ctx context.Context, chatID int64, deadFileIDs []string) (int, error) {
if len(deadFileIDs) == 0 {
return 0, nil
}
deadSet := make(map[string]struct{}, len(deadFileIDs))
for _, id := range deadFileIDs {
deadSet[id] = struct{}{}
}
var rows []Message
if err := b.db.WithContext(ctx).
Where("bot_id = ? AND chat_id = ? AND image_file_ids IS NOT NULL", b.botID, chatID).
Find(&rows).Error; err != nil {
return 0, fmt.Errorf("scan rows for cleanup: %w", err)
}
now := time.Now()
updated := 0
for _, row := range rows {
survivors, dirty := stripDeadFileIDs(row.ImageFileIDs, deadSet)
if !dirty {
continue
}
if len(survivors) == 0 {
row.ImageFileIDs = nil
row.FilesCleanedAt = &now
} else {
row.ImageFileIDs = survivors
}
if err := b.db.WithContext(ctx).Save(&row).Error; err != nil {
return updated, fmt.Errorf("update row %d: %w", row.ID, err)
}
updated++
}
return updated, nil
}
-191
View File
@@ -1,191 +0,0 @@
package main
import (
"testing"
"time"
"github.com/stretchr/testify/assert"
)
func TestFormatUploadFilename(t *testing.T) {
cases := []struct {
botID uint
chatID int64
tgMessageID int
ext string
want string
}{
{1, 12345, 42, "jpg", "tg-1-12345-42.jpg"},
{7, -1001234567890, 1, "png", "tg-7--1001234567890-1.png"},
{0, 0, 0, "webp", "tg-0-0-0.webp"},
}
for _, tc := range cases {
got := formatUploadFilename(tc.botID, tc.chatID, tc.tgMessageID, tc.ext)
assert.Equal(t, tc.want, got)
}
}
func TestParseMissingFileIDFromBody(t *testing.T) {
cases := []struct {
name string
body string
want string
}{
{
name: "canonical Anthropic file-not-found body",
body: `{"type":"error","error":{"type":"invalid_request_error","message":"File not found: file_011CNha8iCJcU1wXNR6q4V8w"}}`,
want: "file_011CNha8iCJcU1wXNR6q4V8w",
},
{
name: "trailing punctuation after the id is excluded",
body: `something File not found: file_abc123! more text`,
want: "file_abc123",
},
{
name: "body without the prefix yields empty",
body: `{"type":"error","error":{"message":"Model not found: claude-foo"}}`,
want: "",
},
{
name: "id at the very end of the buffer",
body: `File not found: file_xyz789`,
want: "file_xyz789",
},
{
name: "empty body",
body: "",
want: "",
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
assert.Equal(t, tc.want, parseMissingFileIDFromBody(tc.body))
})
}
}
func TestStripDeadFileIDs(t *testing.T) {
dead := map[string]struct{}{
"file_a": {},
"file_b": {},
}
cases := []struct {
name string
input []string
wantSurvivors []string
wantDirty bool
}{
{
name: "no overlap returns input verbatim",
input: []string{"file_x", "file_y"},
wantSurvivors: []string{"file_x", "file_y"},
wantDirty: false,
},
{
name: "partial overlap returns survivors and reports dirty",
input: []string{"file_a", "file_x", "file_b", "file_y"},
wantSurvivors: []string{"file_x", "file_y"},
wantDirty: true,
},
{
name: "all dead returns empty survivors and dirty",
input: []string{"file_a", "file_b"},
wantSurvivors: []string{},
wantDirty: true,
},
{
name: "empty input is not dirty",
input: []string{},
wantSurvivors: []string{},
wantDirty: false,
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
survivors, dirty := stripDeadFileIDs(tc.input, dead)
assert.Equal(t, tc.wantSurvivors, survivors)
assert.Equal(t, tc.wantDirty, dirty)
})
}
}
func TestMarkFilesPendingCleanup(t *testing.T) {
b, _ := setupBotForTest(t, 123)
chatID := int64(555)
row1 := Message{
BotID: b.botID,
ChatID: chatID,
UserID: 777,
Username: "u",
UserRole: "user",
Text: "look at these",
Timestamp: time.Now(),
IsUser: true,
ImageFileIDs: []string{"file_a", "file_x"},
}
assert.NoError(t, b.db.Create(&row1).Error)
row2 := Message{
BotID: b.botID,
ChatID: chatID,
UserID: 777,
Username: "u",
UserRole: "user",
Text: "screenshot",
Timestamp: time.Now(),
IsUser: true,
ImageFileIDs: []string{"file_a", "file_b"},
}
assert.NoError(t, b.db.Create(&row2).Error)
row3 := Message{
BotID: b.botID,
ChatID: chatID,
UserID: 777,
Username: "u",
UserRole: "user",
Text: "another",
Timestamp: time.Now(),
IsUser: true,
ImageFileIDs: []string{"file_x", "file_y"},
}
assert.NoError(t, b.db.Create(&row3).Error)
row4 := Message{
BotID: b.botID,
ChatID: 999,
UserID: 777,
Username: "u",
UserRole: "user",
Text: "other chat",
Timestamp: time.Now(),
IsUser: true,
ImageFileIDs: []string{"file_a"},
}
assert.NoError(t, b.db.Create(&row4).Error)
updated, err := b.markFilesPendingCleanup(t.Context(), chatID, []string{"file_a", "file_b"})
assert.NoError(t, err)
assert.Equal(t, 2, updated, "rows 1 and 2 should have been updated")
var r1 Message
assert.NoError(t, b.db.First(&r1, row1.ID).Error)
assert.Equal(t, []string{"file_x"}, r1.ImageFileIDs)
assert.Nil(t, r1.FilesCleanedAt)
var r2 Message
assert.NoError(t, b.db.First(&r2, row2.ID).Error)
assert.Empty(t, r2.ImageFileIDs)
assert.NotNil(t, r2.FilesCleanedAt)
var r3 Message
assert.NoError(t, b.db.First(&r3, row3.ID).Error)
assert.Equal(t, []string{"file_x", "file_y"}, r3.ImageFileIDs)
assert.Nil(t, r3.FilesCleanedAt)
var r4 Message
assert.NoError(t, b.db.First(&r4, row4.ID).Error)
assert.Equal(t, []string{"file_a"}, r4.ImageFileIDs)
assert.Nil(t, r4.FilesCleanedAt)
}
+171 -117
View File
@@ -1,143 +1,197 @@
package main
import (
"encoding/json"
"fmt"
"strings"
"testing"
"time"
"github.com/anthropics/anthropic-sdk-go"
)
func TestTimeContextFor(t *testing.T) {
cases := []struct {
// TestLanguageCodeReplacement tests that language code is properly handled and replaced
func TestLanguageCodeReplacement(t *testing.T) {
// Test with provided language code
systemMessage := "User's language preference: '{language}'"
// Test with a specific language code
langValue := "fr"
result := strings.ReplaceAll(systemMessage, "{language}", langValue)
if !strings.Contains(result, "User's language preference: 'fr'") {
t.Errorf("Expected language code 'fr' to be replaced, got: %s", result)
}
// Test with empty language code (should default to "en")
langValue = ""
if langValue == "" {
langValue = "en" // Default to English when language code is not available
}
result = strings.ReplaceAll(systemMessage, "{language}", langValue)
if !strings.Contains(result, "User's language preference: 'en'") {
t.Errorf("Expected default language code 'en' to be used, got: %s", result)
}
}
// TestPremiumStatusReplacement tests that premium status is properly handled and replaced
func TestPremiumStatusReplacement(t *testing.T) {
systemMessage := "User is a {premium_status}"
// Test with premium user
isPremium := true
premiumStatus := "regular user"
if isPremium {
premiumStatus = "premium user"
}
result := strings.ReplaceAll(systemMessage, "{premium_status}", premiumStatus)
if !strings.Contains(result, "User is a premium user") {
t.Errorf("Expected premium status to be replaced with 'premium user', got: %s", result)
}
// Test with regular user
isPremium = false
premiumStatus = "regular user"
if isPremium {
premiumStatus = "premium user"
}
result = strings.ReplaceAll(systemMessage, "{premium_status}", premiumStatus)
if !strings.Contains(result, "User is a regular user") {
t.Errorf("Expected premium status to be replaced with 'regular user', got: %s", result)
}
}
// TestTimeContextCalculation tests that time context is correctly calculated for different hours
func TestTimeContextCalculation(t *testing.T) {
// Test cases for different hours
testCases := []struct {
hour int
expected string
}{
{3, "night"},
{5, "morning"},
{11, "morning"},
{12, "afternoon"},
{17, "afternoon"},
{18, "evening"},
{21, "evening"},
{22, "night"},
{23, "night"},
}
for _, tc := range cases {
ts := int(time.Date(2025, 5, 15, tc.hour, 0, 0, 0, time.Local).Unix())
if got := timeContextFor(ts); got != tc.expected {
t.Errorf("timeContextFor(hour=%d) = %q, want %q", tc.hour, got, tc.expected)
}
}
}
func TestBuildUserContext(t *testing.T) {
noon := int(time.Date(2025, 5, 15, 12, 0, 0, 0, time.Local).Unix())
got := buildUserContext("alice", "Alice", "Smith", true, "de", noon)
for _, want := range []string{"Alice Smith", "@alice", "Preferred language: de", "premium user", "afternoon"} {
if !strings.Contains(got, want) {
t.Errorf("buildUserContext premium: missing %q in:\n%s", want, got)
}
{3, "night"}, // Night: hours < 5 or hours >= 22
{5, "morning"}, // Morning: 5 <= hours < 12
{12, "afternoon"}, // Afternoon: 12 <= hours < 18
{17, "afternoon"}, // Afternoon: 12 <= hours < 18
{18, "evening"}, // Evening: 18 <= hours < 22
{21, "evening"}, // Evening: 18 <= hours < 22
{22, "night"}, // Night: hours < 5 or hours >= 22
{23, "night"}, // Night: hours < 5 or hours >= 22
}
got = buildUserContext("", "", "", false, "", noon)
for _, want := range []string{"User: unknown (Telegram @unknown)", "Preferred language: en", "regular user"} {
if !strings.Contains(got, want) {
t.Errorf("buildUserContext fallback: missing %q in:\n%s", want, got)
}
}
for _, tc := range testCases {
t.Run(fmt.Sprintf("Hour_%d", tc.hour), func(t *testing.T) {
// Create a timestamp for the specified hour
testTime := time.Date(2025, 5, 15, tc.hour, 0, 0, 0, time.UTC)
got = buildUserContext("bob", "Bob", "", false, "en", noon)
if !strings.Contains(got, "User: Bob (Telegram @bob)") {
t.Errorf("buildUserContext firstname-only: got:\n%s", got)
}
}
// Get the hour directly from the test time to ensure it's what we expect
actualHour := testTime.Hour()
if actualHour != tc.hour {
t.Fatalf("Test setup error: expected hour %d, got %d", tc.hour, actualHour)
}
func TestThinkingParamFromConfig(t *testing.T) {
cases := []struct {
name string
mode string
display string
ok bool
want map[string]any
}{
{"unset omits param", "", "", false, nil},
{"unknown value omits param", "bogus", "", false, nil},
{"adaptive no display", ThinkingModeAdaptive, "", true,
map[string]any{"type": "adaptive"}},
{"adaptive summarized", ThinkingModeAdaptive, ThinkingDisplaySummarized, true,
map[string]any{"type": "adaptive", "display": "summarized"}},
{"adaptive omitted", ThinkingModeAdaptive, ThinkingDisplayOmitted, true,
map[string]any{"type": "adaptive", "display": "omitted"}},
{"disabled", ThinkingModeDisabled, "", true,
map[string]any{"type": "disabled"}},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
union, ok := thinkingParamFromConfig(tc.mode, tc.display)
if ok != tc.ok {
t.Fatalf("ok = %v, want %v", ok, tc.ok)
// Calculate time context using the same logic as in anthropic.go
var timeContext string
if actualHour >= 5 && actualHour < 12 {
timeContext = "morning"
} else if actualHour >= 12 && actualHour < 18 {
timeContext = "afternoon"
} else if actualHour >= 18 && actualHour < 22 {
timeContext = "evening"
} else {
timeContext = "night"
}
if !tc.ok {
return
}
raw, err := json.Marshal(union)
if err != nil {
t.Fatalf("marshal: %v", err)
}
var got map[string]any
if err := json.Unmarshal(raw, &got); err != nil {
t.Fatalf("unmarshal %s: %v", raw, err)
}
if len(got) != len(tc.want) {
t.Fatalf("wire shape %s: got %d keys, want %d (%v)", raw, len(got), len(tc.want), tc.want)
}
for k, v := range tc.want {
if got[k] != v {
t.Errorf("wire shape %s: key %q = %v, want %v", raw, k, got[k], v)
}
// Check if the calculated time context matches the expected value
if timeContext != tc.expected {
t.Errorf("For hour %d: expected time context '%s', got '%s'",
actualHour, tc.expected, timeContext)
}
})
}
}
func TestBackwardCompatibleParams(t *testing.T) {
params := anthropic.BetaMessageNewParams{
Model: "claude-test",
MaxTokens: defaultMaxTokens,
Messages: []anthropic.BetaMessageParam{
anthropic.NewBetaUserMessage(anthropic.NewBetaTextBlock("hi")),
},
}
raw, err := json.Marshal(params)
if err != nil {
t.Fatalf("marshal: %v", err)
}
var got map[string]any
if err := json.Unmarshal(raw, &got); err != nil {
t.Fatalf("unmarshal: %v", err)
}
if _, present := got["thinking"]; present {
t.Errorf("zero Thinking union must omit the key; body: %s", raw)
}
if mt, ok := got["max_tokens"].(float64); !ok || int(mt) != defaultMaxTokens {
t.Errorf("max_tokens = %v, want %d; body: %s", got["max_tokens"], defaultMaxTokens, raw)
}
}
// TestSystemMessagePlaceholderReplacement tests that all placeholders are correctly replaced
func TestSystemMessagePlaceholderReplacement(t *testing.T) {
systemMessage := "The user you're talking to has username '{username}' and display name '{firstname} {lastname}'.\n" +
"User's language preference: '{language}'\n" +
"User is a {premium_status}\n" +
"It's currently {time_context} in your timezone"
func TestEmptyStreamError(t *testing.T) {
err := emptyStreamError("max_tokens", 3900, 4000)
for _, want := range []string{"output budget exhausted", "3900", "4000"} {
if !strings.Contains(err.Error(), want) {
t.Errorf("max_tokens case: %q missing %q", err.Error(), want)
}
// Set up test data
username := "testuser"
firstName := "Test"
lastName := "User"
isPremium := true
languageCode := "de"
// Create a timestamp for a specific hour (e.g., 14:00 = afternoon)
testTime := time.Date(2025, 5, 15, 14, 0, 0, 0, time.UTC)
messageTime := int(testTime.Unix())
// Handle username placeholder
usernameValue := username
if username == "" {
usernameValue = "unknown"
}
if got := emptyStreamError("end_turn", 0, 1000).Error(); got != "unexpected response format from Anthropic" {
t.Errorf("generic case = %q", got)
systemMessage = strings.ReplaceAll(systemMessage, "{username}", usernameValue)
// Handle firstname placeholder
firstnameValue := firstName
if firstName == "" {
firstnameValue = "unknown"
}
if got := emptyStreamError("", 0, 1000).Error(); got != "unexpected response format from Anthropic" {
t.Errorf("no-stop-reason case = %q", got)
systemMessage = strings.ReplaceAll(systemMessage, "{firstname}", firstnameValue)
// Handle lastname placeholder
lastnameValue := lastName
if lastName == "" {
lastnameValue = ""
}
systemMessage = strings.ReplaceAll(systemMessage, "{lastname}", lastnameValue)
// Handle language code placeholder
langValue := languageCode
if languageCode == "" {
langValue = "en"
}
systemMessage = strings.ReplaceAll(systemMessage, "{language}", langValue)
// Handle premium status
premiumStatus := "regular user"
if isPremium {
premiumStatus = "premium user"
}
systemMessage = strings.ReplaceAll(systemMessage, "{premium_status}", premiumStatus)
// Handle time awareness
timeObj := time.Unix(int64(messageTime), 0)
hour := timeObj.Hour()
var timeContext string
if hour >= 5 && hour < 12 {
timeContext = "morning"
} else if hour >= 12 && hour < 18 {
timeContext = "afternoon"
} else if hour >= 18 && hour < 22 {
timeContext = "evening"
} else {
timeContext = "night"
}
systemMessage = strings.ReplaceAll(systemMessage, "{time_context}", timeContext)
// Check that all placeholders were replaced correctly
if !strings.Contains(systemMessage, "username 'testuser'") {
t.Errorf("Username not replaced correctly, got: %s", systemMessage)
}
if !strings.Contains(systemMessage, "display name 'Test User'") {
t.Errorf("Display name not replaced correctly, got: %s", systemMessage)
}
if !strings.Contains(systemMessage, "language preference: 'de'") {
t.Errorf("Language preference not replaced correctly, got: %s", systemMessage)
}
if !strings.Contains(systemMessage, "User is a premium user") {
t.Errorf("Premium status not replaced correctly, got: %s", systemMessage)
}
if !strings.Contains(systemMessage, "It's currently afternoon in your timezone") {
t.Errorf("Time context not replaced correctly, got: %s", systemMessage)
}
}
+123 -128
View File
@@ -8,17 +8,16 @@ import (
"sync"
"time"
"github.com/anthropics/anthropic-sdk-go"
"github.com/anthropics/anthropic-sdk-go/option"
"github.com/go-telegram/bot"
"github.com/go-telegram/bot/models"
"github.com/liushuangls/go-anthropic/v2"
"gorm.io/gorm"
)
type Bot struct {
tgBot TelegramClient
db *gorm.DB
anthropicClient anthropic.Client
anthropicClient *anthropic.Client
chatMemories map[int64]*ChatMemory
memorySize int
chatMemoriesMu sync.RWMutex
@@ -26,11 +25,10 @@ type Bot struct {
userLimiters map[int64]*userLimiter
userLimitersMu sync.RWMutex
clock Clock
botID uint
albumBuffers map[string]*pendingAlbum
albumBuffersMu sync.Mutex
botID uint // Reference to BotModel.ID
}
// Helper function to determine message type
func messageType(msg *models.Message) string {
if msg.Sticker != nil {
return "sticker"
@@ -38,11 +36,13 @@ func messageType(msg *models.Message) string {
return "text"
}
// NewBot initializes and returns a new Bot instance.
func NewBot(db *gorm.DB, config BotConfig, clock Clock, tgClient TelegramClient) (*Bot, error) {
// Retrieve or create Bot entry in the database
var botEntry BotModel
err := db.Where("identifier = ?", config.ID).First(&botEntry).Error
if errors.Is(err, gorm.ErrRecordNotFound) {
botEntry = BotModel{Identifier: config.ID, Name: config.ID}
botEntry = BotModel{Identifier: config.ID, Name: config.ID} // Customize as needed
if err := db.Create(&botEntry).Error; err != nil {
return nil, err
}
@@ -50,9 +50,11 @@ func NewBot(db *gorm.DB, config BotConfig, clock Clock, tgClient TelegramClient)
return nil, err
}
// Ensure the owner exists in the Users table
var owner User
err = db.Where("telegram_id = ? AND bot_id = ?", config.OwnerTelegramID, botEntry.ID).First(&owner).Error
if errors.Is(err, gorm.ErrRecordNotFound) {
// Assign the "owner" role
var ownerRole Role
err := db.Where("name = ?", "owner").First(&ownerRole).Error
if err != nil {
@@ -62,12 +64,13 @@ func NewBot(db *gorm.DB, config BotConfig, clock Clock, tgClient TelegramClient)
owner = User{
BotID: botEntry.ID,
TelegramID: config.OwnerTelegramID,
Username: "",
Username: "", // Initialize as empty; will be updated upon interaction
RoleID: ownerRole.ID,
IsOwner: true,
}
if err := db.Create(&owner).Error; err != nil {
// If unique constraint is violated, another owner already exists
if strings.Contains(err.Error(), "unique index") {
return nil, fmt.Errorf("an owner already exists for this bot")
}
@@ -77,7 +80,8 @@ func NewBot(db *gorm.DB, config BotConfig, clock Clock, tgClient TelegramClient)
return nil, err
}
anthropicClient := anthropic.NewClient(option.WithAPIKey(config.AnthropicAPIKey))
// Use the per-bot Anthropic API key
anthropicClient := anthropic.NewClient(config.AnthropicAPIKey)
b := &Bot{
db: db,
@@ -87,9 +91,8 @@ func NewBot(db *gorm.DB, config BotConfig, clock Clock, tgClient TelegramClient)
config: config,
userLimiters: make(map[int64]*userLimiter),
clock: clock,
botID: botEntry.ID,
botID: botEntry.ID, // Ensure BotModel has ID field
tgBot: tgClient,
albumBuffers: make(map[string]*pendingAlbum),
}
if tgClient == nil {
@@ -104,6 +107,7 @@ func NewBot(db *gorm.DB, config BotConfig, clock Clock, tgClient TelegramClient)
return b, nil
}
// Start begins the bot's operation.
func (b *Bot) Start(ctx context.Context) {
b.tgBot.Start(ctx)
}
@@ -113,6 +117,7 @@ func (b *Bot) getOrCreateUser(userID int64, username string, isOwner bool) (User
err := b.db.Preload("Role").Where("telegram_id = ? AND bot_id = ?", userID, b.botID).First(&user).Error
if err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
// Check if an owner already exists for this bot
if isOwner {
var existingOwner User
err := b.db.Where("bot_id = ? AND is_owner = ?", b.botID, true).First(&existingOwner).Error
@@ -128,7 +133,7 @@ func (b *Bot) getOrCreateUser(userID int64, username string, isOwner bool) (User
if isOwner {
roleName = "owner"
} else {
roleName = "user"
roleName = "user" // Assign "user" role to non-owner users
}
err := b.db.Where("name = ?", roleName).First(&role).Error
@@ -146,6 +151,7 @@ func (b *Bot) getOrCreateUser(userID int64, username string, isOwner bool) (User
}
if err := b.db.Create(&user).Error; err != nil {
// If unique constraint is violated, another owner already exists
if strings.Contains(err.Error(), "unique index") {
return User{}, fmt.Errorf("an owner already exists for this bot")
}
@@ -189,9 +195,10 @@ func (b *Bot) createMessage(chatID, userID int64, username, userRole, text strin
return message
}
// storeMessage stores a message in the database and updates its ID
func (b *Bot) storeMessage(message *Message) error {
message.BotID = b.botID
return b.db.Create(message).Error
message.BotID = b.botID // Associate the message with the correct bot
return b.db.Create(message).Error // This will update the message with its new ID
}
func (b *Bot) getOrCreateChatMemory(chatID int64) *ChatMemory {
@@ -205,12 +212,14 @@ func (b *Bot) getOrCreateChatMemory(chatID int64) *ChatMemory {
chatMemory, exists = b.chatMemories[chatID]
if !exists {
// Check if this is a new chat by querying the database
var count int64
b.db.Model(&Message{}).Where("chat_id = ? AND bot_id = ?", chatID, b.botID).Count(&count)
isNewChat := count == 0
isNewChat := count == 0 // Truly new chat if no messages exist
var messages []Message
if !isNewChat {
// Fetch existing messages only if it's not a new chat
err := b.db.Where("chat_id = ? AND bot_id = ?", chatID, b.botID).
Order("timestamp desc").
Limit(b.memorySize * 2).
@@ -218,14 +227,15 @@ func (b *Bot) getOrCreateChatMemory(chatID int64) *ChatMemory {
if err != nil {
ErrorLogger.Printf("Error fetching messages from database: %v", err)
messages = []Message{}
messages = []Message{} // Initialize an empty slice on error
} else {
// Reverse from newest-first to chronological order for conversation context.
for i, j := 0, len(messages)-1; i < j; i, j = i+1, j-1 {
messages[i], messages[j] = messages[j], messages[i]
}
}
} else {
messages = []Message{}
messages = []Message{} // Ensure messages is initialized for new chats
}
chatMemory = &ChatMemory{
@@ -240,84 +250,65 @@ func (b *Bot) getOrCreateChatMemory(chatID int64) *ChatMemory {
return chatMemory
}
func (b *Bot) stripDeadFileIDFromMemory(chatID int64, deadFileID string) {
b.chatMemoriesMu.Lock()
defer b.chatMemoriesMu.Unlock()
cm, exists := b.chatMemories[chatID]
if !exists {
return
}
for i := range cm.Messages {
if len(cm.Messages[i].ImageFileIDs) == 0 {
continue
}
survivors := make([]string, 0, len(cm.Messages[i].ImageFileIDs))
for _, fid := range cm.Messages[i].ImageFileIDs {
if fid != deadFileID {
survivors = append(survivors, fid)
}
}
cm.Messages[i].ImageFileIDs = survivors
}
}
// addMessageToChatMemory adds a new message to the chat memory, ensuring the memory size is maintained.
func (b *Bot) addMessageToChatMemory(chatMemory *ChatMemory, message Message) {
b.chatMemoriesMu.Lock()
defer b.chatMemoriesMu.Unlock()
// Add the new message
chatMemory.Messages = append(chatMemory.Messages, message)
// Maintain the memory size
if len(chatMemory.Messages) > chatMemory.Size {
chatMemory.Messages = chatMemory.Messages[len(chatMemory.Messages)-chatMemory.Size:]
}
}
func (b *Bot) prepareContextMessages(chatMemory *ChatMemory) []anthropic.BetaMessageParam {
func (b *Bot) prepareContextMessages(chatMemory *ChatMemory) []anthropic.Message {
b.chatMemoriesMu.RLock()
defer b.chatMemoriesMu.RUnlock()
// Debug logging
InfoLogger.Printf("Chat memory contains %d messages", len(chatMemory.Messages))
for i, msg := range chatMemory.Messages {
InfoLogger.Printf("Message %d: IsUser=%v, Text=%q Images=%d", i, msg.IsUser, msg.Text, len(msg.ImageFileIDs))
InfoLogger.Printf("Message %d: IsUser=%v, Text=%q", i, msg.IsUser, msg.Text)
}
var contextMessages []anthropic.BetaMessageParam
// Note: consecutive messages with the same role are permitted.
// The Anthropic API automatically merges them into a single turn rather than
// returning an error. This can happen after a /clear (which only deletes user
// messages, leaving assistant messages in the DB) followed by a restart.
// See: https://platform.claude.com/docs/en/api/messages
var contextMessages []anthropic.Message
for _, msg := range chatMemory.Messages {
blocks := contentBlocksForMessage(msg)
if len(blocks) == 0 {
role := anthropic.RoleUser
if !msg.IsUser {
role = anthropic.RoleAssistant
}
textContent := strings.TrimSpace(msg.Text)
if textContent == "" {
// Skip empty messages
continue
}
var param anthropic.BetaMessageParam
if msg.IsUser {
param = anthropic.NewBetaUserMessage(blocks...)
} else {
param = anthropic.BetaMessageParam{
Role: anthropic.BetaMessageParamRoleAssistant,
Content: blocks,
}
}
contextMessages = append(contextMessages, param)
contextMessages = append(contextMessages, anthropic.Message{
Role: role,
Content: []anthropic.MessageContent{
anthropic.NewTextMessageContent(textContent),
},
})
}
return contextMessages
}
func contentBlocksForMessage(msg Message) []anthropic.BetaContentBlockParamUnion {
var blocks []anthropic.BetaContentBlockParamUnion
if msg.IsUser && len(msg.ImageFileIDs) > 0 {
multi := len(msg.ImageFileIDs) > 1
for i, fileID := range msg.ImageFileIDs {
if multi {
blocks = append(blocks, anthropic.NewBetaTextBlock(fmt.Sprintf("Image %d:", i+1)))
}
blocks = append(blocks, anthropic.NewBetaImageBlock(anthropic.BetaFileImageSourceParam{FileID: fileID}))
}
}
if textContent := strings.TrimSpace(msg.Text); textContent != "" {
blocks = append(blocks, anthropic.NewBetaTextBlock(textContent))
}
return blocks
func (b *Bot) isNewChat(chatID int64) bool {
var count int64
b.db.Model(&Message{}).Where("chat_id = ? AND bot_id = ?", chatID, b.botID).Count(&count)
return count == 0 // Only consider a chat new if it has 0 messages
}
// roleHasScope reports whether role (with pre-loaded Scopes) contains the given scope name.
func roleHasScope(role Role, scope string) bool {
for _, s := range role.Scopes {
if s.Name == scope {
@@ -327,6 +318,8 @@ func roleHasScope(role Role, scope string) bool {
return false
}
// hasScope reports whether the user identified by userID holds the given scope for this bot.
// Owners implicitly hold all scopes regardless of their assigned role.
func (b *Bot) hasScope(userID int64, scope string) bool {
var user User
if err := b.db.Preload("Role.Scopes").
@@ -340,44 +333,52 @@ func (b *Bot) hasScope(userID int64, scope string) bool {
return roleHasScope(user.Role, scope)
}
// publicBotCommands are shown to every user in the Telegram command palette.
var publicBotCommands = []models.BotCommand{
{Command: "stats", Description: "Get bot statistics. Usage: /stats or /stats user [user_id]"},
{Command: "whoami", Description: "Get your user information"},
{Command: "clear", Description: "Clear chat history (soft delete). Admins: /clear [user_id]"},
}
// adminBotCommands are shown only in admin/owner chats via BotCommandScopeChatMember.
var adminBotCommands = []models.BotCommand{
{Command: "clear_hard", Description: "Clear chat history (permanently delete). Admins: /clear_hard [user_id]"},
{Command: "set_model", Description: "Switch the AI model (admin/owner only). Usage: /set_model <model-id>"},
}
// registerAdminCommandsForUser scopes the full command palette to a specific user's private chat.
// In Telegram private chats, chat_id == user_id, so both fields carry the same value.
// Errors are logged but treated as non-fatal: the user retains access via permission checks.
func (b *Bot) registerAdminCommandsForUser(ctx context.Context, telegramID int64) {
allCommands := make([]models.BotCommand, 0, len(publicBotCommands)+len(adminBotCommands))
allCommands = append(allCommands, publicBotCommands...)
allCommands = append(allCommands, adminBotCommands...)
_, err := b.tgBot.SetMyCommands(ctx, &bot.SetMyCommandsParams{
Commands: allCommands,
Scope: &models.BotCommandScopeChat{ChatID: telegramID},
Scope: &models.BotCommandScopeChatMember{ChatID: telegramID, UserID: telegramID},
})
if err != nil {
ErrorLogger.Printf("Failed to register admin commands for user %d: %v", telegramID, err)
}
}
// setElevatedCommands registers the full command palette (public + admin) for every user
// whose role carries the model:set scope, or who is the bot owner. Called once at startup
// and uses the freshly created tgBot directly (b.tgBot is not yet assigned at that point).
func setElevatedCommands(tgBot TelegramClient, users []User) {
allCommands := make([]models.BotCommand, 0, len(publicBotCommands)+len(adminBotCommands))
allCommands = append(allCommands, publicBotCommands...)
allCommands = append(allCommands, adminBotCommands...)
for _, u := range users {
if u.TelegramID == 0 {
continue
continue // skip placeholder users not yet seen in a chat
}
if !u.IsOwner && !roleHasScope(u.Role, ScopeModelSet) {
continue
}
_, err := tgBot.SetMyCommands(context.Background(), &bot.SetMyCommandsParams{
Commands: allCommands,
Scope: &models.BotCommandScopeChat{ChatID: u.TelegramID},
Scope: &models.BotCommandScopeChatMember{ChatID: u.TelegramID, UserID: u.TelegramID},
})
if err != nil {
ErrorLogger.Printf("Warning: could not set admin commands for user %d: %v", u.TelegramID, err)
@@ -395,6 +396,7 @@ func initTelegramBot(token string, b *Bot) (TelegramClient, error) {
return nil, err
}
// Register public commands for all users.
_, err = tgBot.SetMyCommands(context.Background(), &bot.SetMyCommandsParams{
Commands: publicBotCommands,
Scope: &models.BotCommandScopeDefault{},
@@ -404,6 +406,10 @@ func initTelegramBot(token string, b *Bot) (TelegramClient, error) {
return nil, err
}
// Register full command palette (public + admin) scoped to each known elevated user.
// BotCommandScopeChatMember targets the user's private DM with the bot (chat_id == user_id).
// Elevation is determined by scope rather than role name, so renaming roles requires no code change.
// This is best-effort: failures are logged but do not prevent the bot from starting.
var allUsers []User
if err := b.db.Preload("Role.Scopes").Where("bot_id = ?", b.botID).Find(&allUsers).Error; err != nil {
ErrorLogger.Printf("Warning: could not query users for command scoping: %v", err)
@@ -415,12 +421,14 @@ func initTelegramBot(token string, b *Bot) (TelegramClient, error) {
}
func (b *Bot) sendResponse(ctx context.Context, chatID int64, text string, businessConnectionID string) error {
// Pass the outgoing message through the centralized screen for storage and chat memory update
_, err := b.screenOutgoingMessage(chatID, text)
if err != nil {
ErrorLogger.Printf("Error storing assistant message: %v", err)
return err
}
// Prepare message parameters
params := &bot.SendMessageParams{
ChatID: chatID,
Text: text,
@@ -430,6 +438,7 @@ func (b *Bot) sendResponse(ctx context.Context, chatID int64, text string, busin
params.BusinessConnectionID = businessConnectionID
}
// Send the message via Telegram client
_, err = b.tgBot.SendMessage(ctx, params)
if err != nil {
ErrorLogger.Printf("[%s] Error sending message to chat %d with BusinessConnectionID %s: %v",
@@ -439,23 +448,9 @@ func (b *Bot) sendResponse(ctx context.Context, chatID int64, text string, busin
return nil
}
func (b *Bot) sendOneSegment(ctx context.Context, chatID int64, text, businessConnectionID string) error {
params := &bot.SendMessageParams{
ChatID: chatID,
Text: text,
}
if businessConnectionID != "" {
params.BusinessConnectionID = businessConnectionID
}
if _, err := b.tgBot.SendMessage(ctx, params); err != nil {
ErrorLogger.Printf("[%s] Error sending segment to chat %d with BusinessConnectionID %s: %v",
b.config.ID, chatID, businessConnectionID, err)
return err
}
return nil
}
// sendStats sends the bot statistics to the specified chat.
func (b *Bot) sendStats(ctx context.Context, chatID int64, userID int64, targetUserID int64, businessConnectionID string) {
// If targetUserID is 0, show global stats
if targetUserID == 0 {
totalUsers, totalMessages, err := b.getStats()
if err != nil {
@@ -466,6 +461,7 @@ func (b *Bot) sendStats(ctx context.Context, chatID int64, userID int64, targetU
return
}
// Do NOT manually escape hyphens here
statsMessage := fmt.Sprintf(
"📊 Bot Statistics:\n\n"+
"- Total Users: %d\n"+
@@ -474,42 +470,15 @@ func (b *Bot) sendStats(ctx context.Context, chatID int64, userID int64, targetU
totalMessages,
)
if b.hasScope(userID, ScopeStatsViewAny) {
type topEntry struct {
UserID int64
MsgCount int64
}
var top []topEntry
if err := b.db.Model(&Message{}).
Select("user_id, COUNT(*) as msg_count").
Where("bot_id = ? AND is_user = ? AND deleted_at IS NULL", b.botID, true).
Group("user_id").
Order("msg_count DESC").
Limit(3).
Scan(&top).Error; err != nil {
ErrorLogger.Printf("Error fetching top users: %v", err)
} else if len(top) > 0 {
statsMessage += "\n\n🏆 Most Active Users:"
for i, entry := range top {
var u User
if err := b.db.Select("username").Where("telegram_id = ? AND bot_id = ?", entry.UserID, b.botID).First(&u).Error; err != nil {
u.Username = fmt.Sprintf("ID:%d", entry.UserID)
}
name := u.Username
if name == "" {
name = fmt.Sprintf("ID:%d", entry.UserID)
}
statsMessage += fmt.Sprintf("\n%d. @%s — %d messages", i+1, name, entry.MsgCount)
}
}
}
// Send the response through the centralized screen
if err := b.sendResponse(ctx, chatID, statsMessage, businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending stats message: %v", err)
}
return
}
// If targetUserID is not 0, show user-specific stats
// Check permissions if the user is trying to view someone else's stats
if targetUserID != userID {
if !b.hasScope(userID, ScopeStatsViewAny) {
InfoLogger.Printf("User %d attempted to view stats for user %d without permission", userID, targetUserID)
@@ -520,6 +489,7 @@ func (b *Bot) sendStats(ctx context.Context, chatID int64, userID int64, targetU
}
}
// Get user stats
username, messagesIn, messagesOut, totalMessages, err := b.getUserStats(targetUserID)
if err != nil {
ErrorLogger.Printf("Error fetching user stats: %v\n", err)
@@ -529,6 +499,7 @@ func (b *Bot) sendStats(ctx context.Context, chatID int64, userID int64, targetU
return
}
// Build the user stats message
userInfo := fmt.Sprintf("@%s (ID: %d)", username, targetUserID)
if username == "" {
userInfo = fmt.Sprintf("User ID: %d", targetUserID)
@@ -550,6 +521,7 @@ func (b *Bot) sendStats(ctx context.Context, chatID int64, userID int64, targetU
}
}
// getStats retrieves the total number of users and messages from the database.
func (b *Bot) getStats() (int64, int64, error) {
var totalUsers int64
if err := b.db.Model(&User{}).Where("bot_id = ?", b.botID).Count(&totalUsers).Error; err != nil {
@@ -564,30 +536,36 @@ func (b *Bot) getStats() (int64, int64, error) {
return totalUsers, totalMessages, nil
}
// getUserStats retrieves statistics for a specific user
func (b *Bot) getUserStats(userID int64) (string, int64, int64, int64, error) {
// Get user information from database
var user User
err := b.db.Where("telegram_id = ? AND bot_id = ?", userID, b.botID).First(&user).Error
if err != nil {
return "", 0, 0, 0, fmt.Errorf("user not found: %w", err)
}
// Count messages sent by the user (IN)
var messagesIn int64
if err := b.db.Model(&Message{}).Where("user_id = ? AND bot_id = ? AND is_user = ?",
userID, b.botID, true).Count(&messagesIn).Error; err != nil {
return "", 0, 0, 0, err
}
// Count responses to the user (OUT)
var messagesOut int64
if err := b.db.Model(&Message{}).Where("chat_id IN (SELECT DISTINCT chat_id FROM messages WHERE user_id = ? AND bot_id = ? AND deleted_at IS NULL) AND bot_id = ? AND is_user = ?",
userID, b.botID, b.botID, false).Count(&messagesOut).Error; err != nil {
return "", 0, 0, 0, err
}
// Total messages is the sum
totalMessages := messagesIn + messagesOut
return user.Username, messagesIn, messagesOut, totalMessages, nil
}
// isOnlyEmojis checks if the string consists solely of emojis.
func isOnlyEmojis(s string) bool {
for _, r := range s {
if !isEmoji(r) {
@@ -597,12 +575,14 @@ func isOnlyEmojis(s string) bool {
return true
}
// isEmoji determines if a rune is an emoji.
// This is a simplistic check and can be expanded based on requirements.
func isEmoji(r rune) bool {
return (r >= 0x1F600 && r <= 0x1F64F) ||
(r >= 0x1F300 && r <= 0x1F5FF) ||
(r >= 0x1F680 && r <= 0x1F6FF) ||
(r >= 0x2600 && r <= 0x26FF) ||
(r >= 0x2700 && r <= 0x27BF)
return (r >= 0x1F600 && r <= 0x1F64F) || // Emoticons
(r >= 0x1F300 && r <= 0x1F5FF) || // Misc Symbols and Pictographs
(r >= 0x1F680 && r <= 0x1F6FF) || // Transport and Map
(r >= 0x2600 && r <= 0x26FF) || // Misc symbols
(r >= 0x2700 && r <= 0x27BF) // Dingbats
}
func (b *Bot) sendWhoAmI(ctx context.Context, chatID int64, userID int64, username string, businessConnectionID string) {
@@ -632,11 +612,13 @@ func (b *Bot) sendWhoAmI(ctx context.Context, chatID int64, userID int64, userna
role.Name,
)
// Send the response through the centralized screen
if err := b.sendResponse(ctx, chatID, whoAmIMessage, businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending /whoami message: %v", err)
}
}
// screenIncomingMessage centralizes all incoming message processing: storing messages and updating chat memory.
func (b *Bot) screenIncomingMessage(message *models.Message) (Message, error) {
if b.config.DebugScreening {
start := time.Now()
@@ -652,8 +634,9 @@ func (b *Bot) screenIncomingMessage(message *models.Message) (Message, error) {
}()
}
userRole := "user"
userRole := string(anthropic.RoleUser)
// Determine message text based on message type
messageText := message.Text
if message.Sticker != nil {
if message.Sticker.Emoji != "" {
@@ -662,31 +645,34 @@ func (b *Bot) screenIncomingMessage(message *models.Message) (Message, error) {
messageText = "Sent a sticker."
}
}
if message.Voice != nil {
messageText = "[Voice message]"
}
userMessage := b.createMessage(message.Chat.ID, message.From.ID, message.From.Username, userRole, messageText, true)
// Handle sticker-specific details if present
if message.Sticker != nil {
userMessage.StickerFileID = message.Sticker.FileID
userMessage.StickerEmoji = message.Sticker.Emoji
userMessage.StickerEmoji = message.Sticker.Emoji // Store the sticker emoji
if message.Sticker.Thumbnail != nil {
userMessage.StickerPNGFile = message.Sticker.Thumbnail.FileID
}
}
// Get the chat memory before storing the message
chatMemory := b.getOrCreateChatMemory(message.Chat.ID)
// Store the message and get its ID
if err := b.storeMessage(&userMessage); err != nil {
return Message{}, err
}
// Add the message to the chat memory
b.addMessageToChatMemory(chatMemory, userMessage)
return userMessage, nil
}
// screenOutgoingMessage handles storing of outgoing messages and updating chat memory.
// It also marks the most recent unanswered user message as answered.
func (b *Bot) screenOutgoingMessage(chatID int64, response string) (Message, error) {
if b.config.DebugScreening {
start := time.Now()
@@ -701,11 +687,13 @@ func (b *Bot) screenOutgoingMessage(chatID int64, response string) (Message, err
}()
}
assistantMessage := b.createMessage(chatID, 0, "", "assistant", response, false)
// Create and store the assistant message
assistantMessage := b.createMessage(chatID, 0, "", string(anthropic.RoleAssistant), response, false)
if err := b.storeMessage(&assistantMessage); err != nil {
return Message{}, err
}
// Find and mark the most recent unanswered user message as answered
now := time.Now()
err := b.db.Model(&Message{}).
Where("chat_id = ? AND bot_id = ? AND is_user = ? AND answered_on IS NULL",
@@ -716,8 +704,10 @@ func (b *Bot) screenOutgoingMessage(chatID int64, response string) (Message, err
if err != nil {
ErrorLogger.Printf("Error marking user message as answered: %v", err)
// Continue even if there's an error updating the user message
}
// Update chat memory with the message that now has an ID
chatMemory := b.getOrCreateChatMemory(chatID)
b.addMessageToChatMemory(chatMemory, assistantMessage)
@@ -725,26 +715,31 @@ func (b *Bot) screenOutgoingMessage(chatID int64, response string) (Message, err
}
func (b *Bot) promoteUserToAdmin(promoterID, userToPromoteID int64) error {
// Check if the promoter has the user:promote scope
if !b.hasScope(promoterID, ScopeUserPromote) {
return errors.New("only admins or owners can promote users to admin")
}
// Get the user to promote
userToPromote, err := b.getOrCreateUser(userToPromoteID, "", false)
if err != nil {
return err
}
// Get the admin role
var adminRole Role
if err := b.db.Where("name = ?", "admin").First(&adminRole).Error; err != nil {
return err
}
// Update the user's role
userToPromote.RoleID = adminRole.ID
userToPromote.Role = adminRole
if err := b.db.Save(&userToPromote).Error; err != nil {
return err
}
// Surface admin commands in the newly promoted user's private chat.
b.registerAdminCommandsForUser(context.Background(), userToPromoteID)
return nil
}
-105
View File
@@ -1,105 +0,0 @@
package main
import (
"testing"
"github.com/stretchr/testify/assert"
)
func TestContentBlocksForMessage(t *testing.T) {
t.Run("empty message yields no blocks", func(t *testing.T) {
blocks := contentBlocksForMessage(Message{IsUser: true})
assert.Empty(t, blocks)
})
t.Run("user text only yields one text block", func(t *testing.T) {
blocks := contentBlocksForMessage(Message{IsUser: true, Text: "hello"})
assert.Len(t, blocks, 1)
assert.NotNil(t, blocks[0].OfText)
assert.Equal(t, "hello", blocks[0].OfText.Text)
})
t.Run("user single image without caption — no label, no text", func(t *testing.T) {
blocks := contentBlocksForMessage(Message{
IsUser: true,
ImageFileIDs: []string{"file_solo"},
})
assert.Len(t, blocks, 1)
assert.NotNil(t, blocks[0].OfImage)
assert.NotNil(t, blocks[0].OfImage.Source.OfFile)
assert.Equal(t, "file_solo", blocks[0].OfImage.Source.OfFile.FileID)
})
t.Run("user single image with caption — image before text", func(t *testing.T) {
blocks := contentBlocksForMessage(Message{
IsUser: true,
Text: "is this right?",
ImageFileIDs: []string{"file_solo"},
})
assert.Len(t, blocks, 2)
assert.NotNil(t, blocks[0].OfImage, "image block must come before text per Anthropic guidance")
assert.Equal(t, "file_solo", blocks[0].OfImage.Source.OfFile.FileID)
assert.NotNil(t, blocks[1].OfText)
assert.Equal(t, "is this right?", blocks[1].OfText.Text)
})
t.Run("user album (multi-image) labels each with Image N:", func(t *testing.T) {
blocks := contentBlocksForMessage(Message{
IsUser: true,
Text: "compare these",
ImageFileIDs: []string{"file_a", "file_b", "file_c"},
})
assert.Len(t, blocks, 7)
assert.Equal(t, "Image 1:", blocks[0].OfText.Text)
assert.Equal(t, "file_a", blocks[1].OfImage.Source.OfFile.FileID)
assert.Equal(t, "Image 2:", blocks[2].OfText.Text)
assert.Equal(t, "file_b", blocks[3].OfImage.Source.OfFile.FileID)
assert.Equal(t, "Image 3:", blocks[4].OfText.Text)
assert.Equal(t, "file_c", blocks[5].OfImage.Source.OfFile.FileID)
assert.Equal(t, "compare these", blocks[6].OfText.Text)
})
t.Run("assistant message with images-set is text-only (defensive)", func(t *testing.T) {
blocks := contentBlocksForMessage(Message{
IsUser: false,
Text: "I see your screenshot",
ImageFileIDs: []string{"file_should_be_ignored"},
})
assert.Len(t, blocks, 1)
assert.NotNil(t, blocks[0].OfText)
assert.Equal(t, "I see your screenshot", blocks[0].OfText.Text)
})
t.Run("whitespace-only text is skipped but images survive", func(t *testing.T) {
blocks := contentBlocksForMessage(Message{
IsUser: true,
Text: " \n ",
ImageFileIDs: []string{"file_x"},
})
assert.Len(t, blocks, 1)
assert.NotNil(t, blocks[0].OfImage)
})
}
func TestStripDeadFileIDFromMemory(t *testing.T) {
b, _ := setupBotForTest(t, 100)
chatID := int64(42)
cm := b.getOrCreateChatMemory(chatID)
cm.Messages = []Message{
{IsUser: true, Text: "first", ImageFileIDs: []string{"file_a", "file_b"}},
{IsUser: false, Text: "reply"},
{IsUser: true, Text: "third", ImageFileIDs: []string{"file_b", "file_c"}},
}
b.stripDeadFileIDFromMemory(chatID, "file_b")
assert.Equal(t, []string{"file_a"}, cm.Messages[0].ImageFileIDs, "file_b should be removed from message 1")
assert.Empty(t, cm.Messages[1].ImageFileIDs, "assistant message untouched")
assert.Equal(t, []string{"file_c"}, cm.Messages[2].ImageFileIDs, "file_b should be removed from message 3")
}
func TestStripDeadFileIDFromMemory_UnknownChatIsNoop(t *testing.T) {
b, _ := setupBotForTest(t, 100)
b.stripDeadFileIDFromMemory(99999, "file_anything")
}
+7
View File
@@ -1,25 +1,32 @@
// clock.go
package main
import "time"
// Clock is an interface to abstract time-related functions.
type Clock interface {
Now() time.Time
}
// RealClock implements Clock using the actual time.
type RealClock struct{}
// Now returns the current local time.
func (RealClock) Now() time.Time {
return time.Now()
}
// MockClock implements Clock for testing purposes.
type MockClock struct {
currentTime time.Time
}
// Now returns the mocked current time.
func (mc *MockClock) Now() time.Time {
return mc.currentTime
}
// Advance moves the current time forward by the specified duration.
func (mc *MockClock) Advance(d time.Duration) {
mc.currentTime = mc.currentTime.Add(d)
}
+44 -63
View File
@@ -6,50 +6,50 @@ import (
"os"
"path/filepath"
"strings"
)
type MCPServer struct {
Name string `json:"name"`
URL string `json:"url"`
AuthorizationToken string `json:"authorization_token,omitempty"`
AllowedTools []string `json:"allowed_tools,omitempty"`
}
const (
ThinkingModeAdaptive = "adaptive"
ThinkingModeDisabled = "disabled"
ThinkingDisplaySummarized = "summarized"
ThinkingDisplayOmitted = "omitted"
"github.com/liushuangls/go-anthropic/v2"
)
type BotConfig struct {
ID string `json:"id"`
TelegramToken string `json:"telegram_token"`
MemorySize int `json:"memory_size"`
MessagePerHour int `json:"messages_per_hour"`
MessagePerDay int `json:"messages_per_day"`
TempBanDuration string `json:"temp_ban_duration"`
Model string `json:"model"`
Temperature *float32 `json:"temperature,omitempty"`
MaxTokens int `json:"max_tokens,omitempty"`
Thinking string `json:"thinking,omitempty"`
ThinkingDisplay string `json:"thinking_display,omitempty"`
SystemPrompts map[string]string `json:"system_prompts"`
Active bool `json:"active"`
OwnerTelegramID int64 `json:"owner_telegram_id"`
AnthropicAPIKey string `json:"anthropic_api_key"`
ElevenLabsAPIKey string `json:"elevenlabs_api_key"`
ElevenLabsVoiceID string `json:"elevenlabs_voice_id"`
ElevenLabsModel string `json:"elevenlabs_model"`
DebugScreening bool `json:"debug_screening"`
MCPServers []MCPServer `json:"mcp_servers,omitempty"`
ConfigFilePath string `json:"-"`
ID string `json:"id"`
TelegramToken string `json:"telegram_token"`
MemorySize int `json:"memory_size"`
MessagePerHour int `json:"messages_per_hour"`
MessagePerDay int `json:"messages_per_day"`
TempBanDuration string `json:"temp_ban_duration"`
Model anthropic.Model `json:"model"`
Temperature *float32 `json:"temperature,omitempty"` // Controls creativity vs determinism (0.0-1.0)
SystemPrompts map[string]string `json:"system_prompts"`
Active bool `json:"active"`
OwnerTelegramID int64 `json:"owner_telegram_id"`
AnthropicAPIKey string `json:"anthropic_api_key"`
DebugScreening bool `json:"debug_screening"` // Enable detailed screening logs
ConfigFilePath string `json:"-"` // Set at load time; not serialized
}
// Custom unmarshalling to handle anthropic.Model
func (c *BotConfig) UnmarshalJSON(data []byte) error {
type Alias BotConfig
aux := &struct {
Model string `json:"model"`
*Alias
}{
Alias: (*Alias)(c),
}
if err := json.Unmarshal(data, &aux); err != nil {
return err
}
c.Model = anthropic.Model(aux.Model)
return nil
}
// validateConfigPath ensures the file path is within the allowed directory
func validateConfigPath(configDir, filename string) (string, error) {
// Clean the paths to remove any . or .. components
configDir = filepath.Clean(configDir)
filename = filepath.Clean(filename)
// Get absolute paths
absConfigDir, err := filepath.Abs(configDir)
if err != nil {
return "", fmt.Errorf("failed to get absolute path for config directory: %w", err)
@@ -61,11 +61,13 @@ func validateConfigPath(configDir, filename string) (string, error) {
return "", fmt.Errorf("failed to get absolute path for config file: %w", err)
}
// Use filepath.Rel to check if the path is within the config directory
rel, err := filepath.Rel(absConfigDir, absPath)
if err != nil || strings.HasPrefix(rel, "..") || strings.Contains(rel, "..") {
return "", fmt.Errorf("invalid config path: file must be within the config directory")
}
// Verify file extension
if filepath.Ext(absPath) != ".json" {
return "", fmt.Errorf("invalid file extension: must be .json")
}
@@ -107,11 +109,6 @@ func loadAllConfigs(dir string) ([]BotConfig, error) {
continue
}
if config.Thinking == ThinkingModeAdaptive && config.MaxTokens > 0 && config.MaxTokens < 4000 {
InfoLogger.Printf("[%s] thinking=adaptive with max_tokens=%d: thinking tokens count toward max_tokens; consider >= 4000",
config.ID, config.MaxTokens)
}
config.ConfigFilePath = validPath
configs = append(configs, config)
}
@@ -145,29 +142,6 @@ func validateConfig(config *BotConfig, ids, tokens map[string]bool) error {
return fmt.Errorf("missing 'model' field")
}
switch config.Thinking {
case "", ThinkingModeAdaptive, ThinkingModeDisabled:
default:
return fmt.Errorf("invalid 'thinking' value %q: must be %q or %q (or omitted)",
config.Thinking, ThinkingModeAdaptive, ThinkingModeDisabled)
}
switch config.ThinkingDisplay {
case "":
case ThinkingDisplaySummarized, ThinkingDisplayOmitted:
if config.Thinking != ThinkingModeAdaptive {
return fmt.Errorf("'thinking_display' requires 'thinking': %q (the API rejects display with thinking disabled)",
ThinkingModeAdaptive)
}
default:
return fmt.Errorf("invalid 'thinking_display' value %q: must be %q or %q (or omitted)",
config.ThinkingDisplay, ThinkingDisplaySummarized, ThinkingDisplayOmitted)
}
if config.MaxTokens < 0 {
return fmt.Errorf("'max_tokens' must be greater than 0 when set")
}
if config.MessagePerHour <= 0 {
return fmt.Errorf("'messages_per_hour' must be greater than 0")
}
@@ -181,6 +155,7 @@ func validateConfig(config *BotConfig, ids, tokens map[string]bool) error {
func loadConfig(filename string) (BotConfig, error) {
var config BotConfig
// Use filepath.Clean before opening the file
file, err := os.OpenFile(filepath.Clean(filename), os.O_RDONLY, 0)
if err != nil {
return config, fmt.Errorf("failed to open config file %s: %w", filename, err)
@@ -199,12 +174,15 @@ func loadConfig(filename string) (BotConfig, error) {
return config, nil
}
// Reload reloads the BotConfig from the specified filename within the given config directory
func (c *BotConfig) Reload(configDir, filename string) error {
// Validate the config path
validPath, err := validateConfigPath(configDir, filename)
if err != nil {
return fmt.Errorf("invalid config path: %w", err)
}
// Use filepath.Clean before opening the file
cleanPath := filepath.Clean(validPath)
file, err := os.OpenFile(cleanPath, os.O_RDONLY, 0)
if err != nil {
@@ -221,9 +199,12 @@ func (c *BotConfig) Reload(configDir, filename string) error {
return fmt.Errorf("failed to decode JSON from %s: %w", validPath, err)
}
c.Model = anthropic.Model(c.Model)
return nil
}
// PersistModel updates the model field in memory and writes it back to the config file on disk.
// Only the "model" key is changed; all other fields are preserved verbatim.
func (c *BotConfig) PersistModel(newModel string) error {
if c.ConfigFilePath == "" {
return fmt.Errorf("config file path not set; cannot persist model")
@@ -250,6 +231,6 @@ func (c *BotConfig) PersistModel(newModel string) error {
return fmt.Errorf("failed to write config: %w", err)
}
c.Model = newModel
c.Model = anthropic.Model(newModel)
return nil
}
+6 -7
View File
@@ -4,19 +4,18 @@
"telegram_token": "YOUR_TELEGRAM_BOT_TOKEN",
"owner_telegram_id": 111111111,
"anthropic_api_key": "YOUR_SPECIFIC_ANTHROPIC_API_KEY",
"elevenlabs_api_key": "",
"elevenlabs_voice_id": "",
"elevenlabs_model": "",
"memory_size": 10,
"messages_per_hour": 20,
"messages_per_day": 100,
"temp_ban_duration": "24h",
"model": "claude-haiku-4-5",
"model": "claude-3-5-haiku-latest",
"temperature": 0.7,
"max_tokens": 1000,
"debug_screening": false,
"system_prompts": {
"custom_instructions": "You are Atom, a helpful assistant texting through a limited Telegram interface with a 15-word maximum. Write like texting a friend - use shorthand, skip grammar, use slang/abbreviations. The system cuts off anything longer than 15 words.\n\n- Address the user by their first name, and reply in their preferred language (both are in the conversation context).\n- Use time-appropriate greetings based on the user's local time of day.\n- If a user asks about buying apples, inform them that we don't sell apples.\n- When asked for a joke, tell a clean, family-friendly joke about programming or technology.\n- If someone inquires about our services, explain that we offer AI-powered chatbot solutions.\n- For any questions about pricing, direct users to contact our sales team at [email protected].\n- If asked about your capabilities, be honest about what you can and cannot do.\nAlways maintain a friendly and professional tone.",
"respond_with_emojis": "The user's message contains only emoji. Reply using only emoji."
"default": "You are a helpful assistant.",
"custom_instructions": "You are texting through a limited Telegram interface with 15-word maximum. Write like texting a friend - use shorthand, skip grammar, use slang/abbreviations. System cuts off anything longer than 15 words.\n\n- Your name is Atom.\n- The user you're talking to has username '{username}' and display name '{firstname} {lastname}'.\n- User's language preference: '{language}'\n- User is a {premium_status}\n- It's currently {time_context} in your timezone. Use appropriate time-based greetings and address the user by name.\n- If a user asks about buying apples, inform them that we don't sell apples.\n- When asked for a joke, tell a clean, family-friendly joke about programming or technology.\n- If someone inquires about our services, explain that we offer AI-powered chatbot solutions.\n- For any questions about pricing, direct users to contact our sales team at [email protected].\n- If asked about your capabilities, be honest about what you can and cannot do.\nAlways maintain a friendly and professional tone.",
"continue_conversation": "Continuing our conversation. Remember previous context if relevant.",
"avoid_sensitive": "Avoid discussing sensitive topics or providing harmful information.",
"respond_with_emojis": "Since the user sent only emojis, respond using emojis only."
}
}
+47 -106
View File
@@ -2,18 +2,21 @@ package main
import (
"encoding/json"
"fmt"
"os"
"path/filepath"
"strings"
"testing"
"github.com/liushuangls/go-anthropic/v2"
)
// Set up loggers
func TestMain(m *testing.M) {
initLoggers()
os.Exit(m.Run())
}
// TestBotConfig_UnmarshalJSON tests the custom unmarshalling of BotConfig
func TestBotConfig_UnmarshalJSON(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
jsonData := `{
"id": "bot123",
@@ -35,7 +38,7 @@ func TestBotConfig_UnmarshalJSON(t *testing.T) { //NOSONAR go:S100 -- underscore
t.Fatalf("Failed to unmarshal JSON: %v", err)
}
expectedModel := "claude-v1"
expectedModel := anthropic.Model("claude-v1")
if config.Model != expectedModel {
t.Errorf("Expected model %s, got %s", expectedModel, config.Model)
}
@@ -45,8 +48,10 @@ func TestBotConfig_UnmarshalJSON(t *testing.T) { //NOSONAR go:S100 -- underscore
t.Errorf("Expected ID %s, got %s", expectedID, config.ID)
}
// Add more field checks as necessary
}
// TestValidateConfigPath tests the validateConfigPath function
func TestValidateConfigPath(t *testing.T) {
execDir, err := os.Getwd()
if err != nil {
@@ -91,6 +96,7 @@ func TestValidateConfigPath(t *testing.T) {
},
}
// Create a subdirectory for testing
subDir := filepath.Join(execDir, "subdir")
if err := os.MkdirAll(subDir, 0755); err != nil {
t.Fatalf("Failed to create subdir: %v", err)
@@ -116,7 +122,9 @@ func TestValidateConfigPath(t *testing.T) {
}
}
// TestLoadConfig tests the loadConfig function
func TestLoadConfig(t *testing.T) {
// Create a temporary directory
tempDir, err := os.MkdirTemp("", "config_test")
if err != nil {
t.Fatalf("Failed to create temp dir: %v", err)
@@ -127,6 +135,7 @@ func TestLoadConfig(t *testing.T) {
}
}()
// Valid config JSON
validConfig := `{
"id": "bot123",
"telegram_token": "token123",
@@ -142,6 +151,7 @@ func TestLoadConfig(t *testing.T) {
"anthropic_api_key": "api_key_123"
}`
// Invalid config JSON
invalidConfig := `{
"id": "bot123",
"telegram_token": "token123",
@@ -149,11 +159,13 @@ func TestLoadConfig(t *testing.T) {
"model": "claude-v1"
}`
// Write valid config file
validPath := filepath.Join(tempDir, "valid_config.json")
if err := os.WriteFile(validPath, []byte(validConfig), 0644); err != nil {
t.Fatalf("Failed to write valid config: %v", err)
}
// Write invalid config file
invalidPath := filepath.Join(tempDir, "invalid_config.json")
if err := os.WriteFile(invalidPath, []byte(invalidConfig), 0644); err != nil {
t.Fatalf("Failed to write invalid config: %v", err)
@@ -206,6 +218,7 @@ func TestLoadConfig(t *testing.T) {
}
}
// TestValidateConfig tests the validateConfig function
func TestValidateConfig(t *testing.T) {
tests := []struct {
name string
@@ -338,7 +351,9 @@ func TestValidateConfig(t *testing.T) {
}
}
// TestLoadAllConfigs tests the loadAllConfigs function
func TestLoadAllConfigs(t *testing.T) {
// Create a temporary directory
tempDir, err := os.MkdirTemp("", "load_all_configs_test")
if err != nil {
t.Fatalf("Failed to create temp dir: %v", err)
@@ -351,7 +366,7 @@ func TestLoadAllConfigs(t *testing.T) {
tests := []struct {
name string
setupFiles map[string]string
setupFiles map[string]string // filename -> content
expectConfigs int
expectError bool
expectErrorMsg string
@@ -509,6 +524,7 @@ func TestLoadAllConfigs(t *testing.T) {
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
// Clear the tempDir before each test
if err := os.RemoveAll(tempDir); err != nil {
t.Fatalf("Failed to remove temp dir: %v", err)
}
@@ -516,6 +532,7 @@ func TestLoadAllConfigs(t *testing.T) {
t.Fatalf("Failed to create temp dir: %v", err)
}
// Write the test files directly
for filename, content := range tt.setupFiles {
err := os.WriteFile(filepath.Join(tempDir, filename), []byte(content), 0644)
if err != nil {
@@ -535,7 +552,9 @@ func TestLoadAllConfigs(t *testing.T) {
}
}
// TestBotConfig_Reload tests the Reload method of BotConfig
func TestBotConfig_Reload(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
// Create a temporary directory
tempDir, err := os.MkdirTemp("", "reload_test")
if err != nil {
t.Fatalf("Failed to create temp dir: %v", err)
@@ -546,6 +565,7 @@ func TestBotConfig_Reload(t *testing.T) { //NOSONAR go:S100 -- underscore separa
}
}()
// Create initial config file
config1 := `{
"id": "bot123",
"telegram_token": "token123",
@@ -565,11 +585,13 @@ func TestBotConfig_Reload(t *testing.T) { //NOSONAR go:S100 -- underscore separa
t.Fatalf("Failed to write initial config: %v", err)
}
// Initialize BotConfig
var config BotConfig
if err := config.Reload(tempDir, "config.json"); err != nil {
t.Fatalf("Failed to reload config: %v", err)
}
// Verify initial load
if config.ID != "bot123" {
t.Errorf("Expected ID 'bot123', got '%s'", config.ID)
}
@@ -577,6 +599,7 @@ func TestBotConfig_Reload(t *testing.T) { //NOSONAR go:S100 -- underscore separa
t.Errorf("Expected Model 'claude-v1', got '%s'", config.Model)
}
// Update config file
config2 := `{
"id": "bot123",
"telegram_token": "token123_updated",
@@ -595,10 +618,12 @@ func TestBotConfig_Reload(t *testing.T) { //NOSONAR go:S100 -- underscore separa
t.Fatalf("Failed to write updated config: %v", err)
}
// Reload config
if err := config.Reload(tempDir, "config.json"); err != nil {
t.Fatalf("Failed to reload updated config: %v", err)
}
// Verify updated config
if config.TelegramToken != "token123_updated" {
t.Errorf("Expected TelegramToken 'token123_updated', got '%s'", config.TelegramToken)
}
@@ -613,6 +638,7 @@ func TestBotConfig_Reload(t *testing.T) { //NOSONAR go:S100 -- underscore separa
}
}
// TestBotConfig_UnmarshalJSON_Invalid tests unmarshalling with invalid model
func TestBotConfig_UnmarshalJSON_Invalid(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
jsonData := `{
"id": "bot123",
@@ -640,11 +666,14 @@ func TestBotConfig_UnmarshalJSON_Invalid(t *testing.T) { //NOSONAR go:S100 -- un
}
}
// Helper function to check substring
func contains(s, substr string) bool {
return strings.Contains(s, substr)
}
// TestTemperatureConfig tests that the temperature value is correctly loaded
func TestTemperatureConfig(t *testing.T) {
// Create a temporary directory
tempDir, err := os.MkdirTemp("", "temperature_test")
if err != nil {
t.Fatalf("Failed to create temp dir: %v", err)
@@ -655,6 +684,7 @@ func TestTemperatureConfig(t *testing.T) {
}
}()
// Create config with temperature
configWithTemp := `{
"id": "bot123",
"telegram_token": "token123",
@@ -670,6 +700,7 @@ func TestTemperatureConfig(t *testing.T) {
"anthropic_api_key": "api_key_123"
}`
// Create config without temperature
configWithoutTemp := `{
"id": "bot124",
"telegram_token": "token124",
@@ -684,6 +715,7 @@ func TestTemperatureConfig(t *testing.T) {
"anthropic_api_key": "api_key_123"
}`
// Write config files
withTempPath := filepath.Join(tempDir, "with_temp.json")
if err := os.WriteFile(withTempPath, []byte(configWithTemp), 0644); err != nil {
t.Fatalf("Failed to write config with temperature: %v", err)
@@ -694,27 +726,35 @@ func TestTemperatureConfig(t *testing.T) {
t.Fatalf("Failed to write config without temperature: %v", err)
}
// Test loading config with temperature
configWithTempObj, err := loadConfig(withTempPath)
if err != nil {
t.Fatalf("Failed to load config with temperature: %v", err)
}
// Verify temperature is set correctly
if configWithTempObj.Temperature == nil {
t.Errorf("Expected Temperature to be set, got nil")
} else if *configWithTempObj.Temperature != 0.42 {
t.Errorf("Expected Temperature 0.42, got %f", *configWithTempObj.Temperature)
}
// Test loading config without temperature
configWithoutTempObj, err := loadConfig(withoutTempPath)
if err != nil {
t.Fatalf("Failed to load config without temperature: %v", err)
}
// Verify temperature is nil when not specified
if configWithoutTempObj.Temperature != nil {
t.Errorf("Expected Temperature to be nil, got %f", *configWithoutTempObj.Temperature)
}
}
// Additional tests can be added here to cover more scenarios
// TestBotConfig_PersistModel verifies that PersistModel updates the model both in memory
// and on disk while leaving all other config fields unchanged.
func TestBotConfig_PersistModel(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
tempDir, err := os.MkdirTemp("", "persist_model_test")
if err != nil {
@@ -744,14 +784,17 @@ func TestBotConfig_PersistModel(t *testing.T) { //NOSONAR go:S100 -- underscore
ConfigFilePath: configPath,
}
// Successful model update
if err := config.PersistModel("claude-sonnet-4-6"); err != nil {
t.Fatalf("PersistModel() unexpected error: %v", err)
}
// In-memory model must be updated immediately
if string(config.Model) != "claude-sonnet-4-6" {
t.Errorf("in-memory model: got %q, want %q", config.Model, "claude-sonnet-4-6")
}
// On-disk model must be updated; other fields must be preserved
data, err := os.ReadFile(configPath)
if err != nil {
t.Fatalf("Failed to read updated config file: %v", err)
@@ -767,111 +810,9 @@ func TestBotConfig_PersistModel(t *testing.T) { //NOSONAR go:S100 -- underscore
t.Errorf("on-disk id should be preserved: got %v, want %q", raw["id"], "bot1")
}
// Error case: empty ConfigFilePath must return an error
noPath := BotConfig{Model: "claude-v1"}
if err := noPath.PersistModel("claude-sonnet-4-6"); err == nil {
t.Error("PersistModel with empty ConfigFilePath: expected error, got nil")
}
}
func thinkingTestConfig(id string) BotConfig {
return BotConfig{
ID: id,
TelegramToken: "token-" + id,
MemorySize: 10,
MessagePerHour: 10,
MessagePerDay: 100,
TempBanDuration: "1h",
Model: "claude-test",
}
}
func TestThinkingConfig(t *testing.T) {
cases := []struct {
name string
thinking string
display string
wantErr string
}{
{"absent", "", "", ""},
{"adaptive", ThinkingModeAdaptive, "", ""},
{"disabled", ThinkingModeDisabled, "", ""},
{"adaptive summarized", ThinkingModeAdaptive, ThinkingDisplaySummarized, ""},
{"adaptive omitted", ThinkingModeAdaptive, ThinkingDisplayOmitted, ""},
{"legacy enabled rejected", "enabled", "", "invalid 'thinking'"},
{"case sensitive", "Adaptive", "", "invalid 'thinking'"},
{"unknown display", ThinkingModeAdaptive, "verbose", "invalid 'thinking_display'"},
{"display without thinking", "", ThinkingDisplaySummarized, "'thinking_display' requires"},
{"display with disabled", ThinkingModeDisabled, ThinkingDisplayOmitted, "'thinking_display' requires"},
}
for i, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
cfg := thinkingTestConfig(fmt.Sprintf("bot-think-%d", i))
cfg.Thinking = tc.thinking
cfg.ThinkingDisplay = tc.display
err := validateConfig(&cfg, map[string]bool{}, map[string]bool{})
if tc.wantErr == "" {
if err != nil {
t.Fatalf("validateConfig(thinking=%q display=%q) = %v, want nil", tc.thinking, tc.display, err)
}
return
}
if err == nil {
t.Fatalf("validateConfig(thinking=%q display=%q) = nil, want error containing %q", tc.thinking, tc.display, tc.wantErr)
}
if !contains(err.Error(), tc.wantErr) {
t.Errorf("error %q does not contain %q", err.Error(), tc.wantErr)
}
})
}
}
func TestMaxTokensConfig(t *testing.T) {
var withValue BotConfig
if err := json.Unmarshal([]byte(`{"max_tokens": 4000}`), &withValue); err != nil {
t.Fatalf("unmarshal: %v", err)
}
if withValue.MaxTokens != 4000 {
t.Errorf("MaxTokens = %d, want 4000", withValue.MaxTokens)
}
var withoutValue BotConfig
if err := json.Unmarshal([]byte(`{}`), &withoutValue); err != nil {
t.Fatalf("unmarshal: %v", err)
}
if withoutValue.MaxTokens != 0 {
t.Errorf("MaxTokens = %d, want 0 when absent", withoutValue.MaxTokens)
}
neg := thinkingTestConfig("bot-maxtok-neg")
neg.MaxTokens = -1
if err := validateConfig(&neg, map[string]bool{}, map[string]bool{}); err == nil {
t.Error("validateConfig(max_tokens=-1) = nil, want error")
}
zero := thinkingTestConfig("bot-maxtok-zero")
zero.MaxTokens = 0
if err := validateConfig(&zero, map[string]bool{}, map[string]bool{}); err != nil {
t.Errorf("validateConfig(max_tokens=0) = %v, want nil", err)
}
}
func TestThinkingConfigLoad(t *testing.T) {
jsonData := `{
"id": "bot-think-load",
"thinking": "adaptive",
"thinking_display": "omitted",
"max_tokens": 4096
}`
var cfg BotConfig
if err := json.Unmarshal([]byte(jsonData), &cfg); err != nil {
t.Fatalf("unmarshal: %v", err)
}
if cfg.Thinking != ThinkingModeAdaptive {
t.Errorf("Thinking = %q, want %q", cfg.Thinking, ThinkingModeAdaptive)
}
if cfg.ThinkingDisplay != ThinkingDisplayOmitted {
t.Errorf("ThinkingDisplay = %q, want %q", cfg.ThinkingDisplay, ThinkingDisplayOmitted)
}
if cfg.MaxTokens != 4096 {
t.Errorf("MaxTokens = %d, want 4096", cfg.MaxTokens)
}
}
+9 -2
View File
@@ -38,11 +38,17 @@ func initDB() (*gorm.DB, error) {
}
sqlDB.SetMaxOpenConns(1)
// AutoMigrate the models
err = db.AutoMigrate(&BotModel{}, &ConfigModel{}, &Message{}, &User{}, &Role{}, &Scope{})
if err != nil {
return nil, fmt.Errorf("failed to migrate database schema: %w", err)
}
// Enforce unique owner per bot using raw SQL
// Note: SQLite doesn't support partial indexes, but we can simulate it by making a unique index on (BotID, IsOwner)
// and ensuring that IsOwner can only be true for one user per BotID.
// This approach allows multiple users with IsOwner=false for the same BotID,
// but only one user can have IsOwner=true per BotID.
err = db.Exec(`CREATE UNIQUE INDEX IF NOT EXISTS idx_bot_owner ON users (bot_id, is_owner) WHERE is_owner = 1;`).Error
if err != nil {
return nil, fmt.Errorf("failed to create unique index for bot owners: %w", err)
@@ -65,7 +71,7 @@ func createDefaultScopes(db *gorm.DB) error {
ScopeStatsViewOwn, ScopeStatsViewAny,
ScopeHistoryClearOwn, ScopeHistoryClearAny,
ScopeHistoryClearHardOwn, ScopeHistoryClearHardAny,
ScopeModelSet, ScopeUserPromote, ScopeTTSUse,
ScopeModelSet, ScopeUserPromote,
}
for _, name := range all {
if err := db.FirstOrCreate(&Scope{}, Scope{Name: name}).Error; err != nil {
@@ -82,11 +88,12 @@ func createDefaultScopes(db *gorm.DB) error {
ScopeStatsViewOwn, ScopeStatsViewAny,
ScopeHistoryClearOwn, ScopeHistoryClearAny,
ScopeHistoryClearHardOwn, ScopeHistoryClearHardAny,
ScopeModelSet, ScopeUserPromote, ScopeTTSUse,
ScopeModelSet, ScopeUserPromote,
}
assignments := map[string][]string{
"user": userScopes,
"admin": elevatedScopes,
// owner gets the same scopes as admin; owner uniqueness is enforced by the IsOwner flag
"owner": elevatedScopes,
}
for roleName, scopes := range assignments {
-99
View File
@@ -1,99 +0,0 @@
package main
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"mime/multipart"
"net/http"
)
const (
elevenLabsTTSURL = "https://api.elevenlabs.io/v1/text-to-speech/"
elevenLabsSTTURL = "https://api.elevenlabs.io/v1/speech-to-text"
elevenLabsDefaultModel = "eleven_multilingual_v2"
)
func (b *Bot) generateSpeech(ctx context.Context, text string) (io.Reader, error) {
model := b.config.ElevenLabsModel
if model == "" {
model = elevenLabsDefaultModel
}
body, err := json.Marshal(map[string]string{
"text": text,
"model_id": model,
})
if err != nil {
return nil, fmt.Errorf("elevenlabs TTS marshal error: %w", err)
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost,
elevenLabsTTSURL+b.config.ElevenLabsVoiceID, bytes.NewReader(body))
if err != nil {
return nil, fmt.Errorf("elevenlabs TTS request error: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("xi-api-key", b.config.ElevenLabsAPIKey)
resp, err := http.DefaultClient.Do(req)
if err != nil {
return nil, fmt.Errorf("elevenlabs TTS error: %w", err)
}
if resp.StatusCode != http.StatusOK {
defer func() { _ = resp.Body.Close() }()
errBody, _ := io.ReadAll(resp.Body)
return nil, fmt.Errorf("elevenlabs TTS error: status %d: %s", resp.StatusCode, errBody)
}
return resp.Body, nil
}
func (b *Bot) transcribeVoice(ctx context.Context, fileID string) (string, error) {
audioBytes, err := b.downloadTelegramFile(ctx, fileID)
if err != nil {
return "", err
}
var buf bytes.Buffer
mw := multipart.NewWriter(&buf)
if err := mw.WriteField("model_id", "scribe_v1"); err != nil {
return "", fmt.Errorf("multipart write error: %w", err)
}
part, err := mw.CreateFormFile("file", "audio.ogg")
if err != nil {
return "", fmt.Errorf("multipart create file error: %w", err)
}
if _, err := io.Copy(part, bytes.NewReader(audioBytes)); err != nil {
return "", fmt.Errorf("multipart copy error: %w", err)
}
if err := mw.Close(); err != nil {
return "", fmt.Errorf("multipart close error: %w", err)
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost,
elevenLabsSTTURL, &buf)
if err != nil {
return "", fmt.Errorf("create STT request error: %w", err)
}
req.Header.Set("Content-Type", mw.FormDataContentType())
req.Header.Set("xi-api-key", b.config.ElevenLabsAPIKey)
sttResp, err := http.DefaultClient.Do(req)
if err != nil {
return "", fmt.Errorf("elevenlabs STT request error: %w", err)
}
defer func() { _ = sttResp.Body.Close() }()
if sttResp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(sttResp.Body)
return "", fmt.Errorf("elevenlabs STT error: status %d: %s", sttResp.StatusCode, body)
}
var result struct {
Text string `json:"text"`
}
if err := json.NewDecoder(sttResp.Body).Decode(&result); err != nil {
return "", fmt.Errorf("elevenlabs STT decode error: %w", err)
}
return result.Text, nil
}
Binary file not shown.
+6 -22
View File
@@ -3,37 +3,21 @@ module github.com/HugeFrog24/go-telegram-bot
go 1.26.0
require (
github.com/anthropics/anthropic-sdk-go v1.57.0
github.com/go-telegram/bot v1.22.0
github.com/go-telegram/bot v1.19.0
github.com/liushuangls/go-anthropic/v2 v2.17.1
github.com/stretchr/testify v1.11.1
golang.org/x/sync v0.22.0
golang.org/x/time v0.15.0
golang.org/x/time v0.14.0
gorm.io/driver/sqlite v1.6.0
gorm.io/gorm v1.31.2
gorm.io/gorm v1.31.1
)
require (
github.com/bahlo/generic-list-go v0.2.0 // indirect
github.com/buger/jsonparser v1.2.0 // indirect
github.com/davecgh/go-spew v1.1.1 // indirect
github.com/invopop/jsonschema v0.14.0 // indirect
github.com/jinzhu/inflection v1.0.0 // indirect
github.com/jinzhu/now v1.1.5 // indirect
github.com/kr/pretty v0.3.1 // indirect
github.com/mailru/easyjson v0.9.2 // indirect
github.com/mattn/go-sqlite3 v1.14.48 // indirect
github.com/pb33f/ordered-map/v2 v2.3.1 // indirect
github.com/mattn/go-sqlite3 v1.14.34 // indirect
github.com/pmezard/go-difflib v1.0.0 // indirect
github.com/rogpeppe/go-internal v1.14.1 // indirect
github.com/standard-webhooks/standard-webhooks/libraries v0.0.1 // indirect
github.com/stretchr/objx v0.5.3 // indirect
github.com/tidwall/gjson v1.19.0 // indirect
github.com/tidwall/match v1.2.0 // indirect
github.com/tidwall/pretty v1.2.1 // indirect
github.com/tidwall/sjson v1.2.5 // indirect
github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect
go.yaml.in/yaml/v4 v4.0.0-rc.6 // indirect
golang.org/x/text v0.40.0 // indirect
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
golang.org/x/text v0.34.0 // indirect
gopkg.in/yaml.v3 v3.0.1 // indirect
)
+11 -81
View File
@@ -1,100 +1,30 @@
github.com/anthropics/anthropic-sdk-go v1.45.0 h1:rWnpyBpm9OAm97jyH5bi6W4SRCwJeNY/RyhaJ7CHSUI=
github.com/anthropics/anthropic-sdk-go v1.45.0/go.mod h1:bx5vWuHFuGPkELH8Z4KUiNSohFnUwScdpTyr+50myPo=
github.com/anthropics/anthropic-sdk-go v1.52.0 h1:1TB9jt4DN87VMwS/hB1VK26tYzK0ipEOtqPaPGFtJQg=
github.com/anthropics/anthropic-sdk-go v1.52.0/go.mod h1:3EfIfmFqxH6rbiLcIP4tPFyXL/IHakx2wDG4OU+TIEI=
github.com/anthropics/anthropic-sdk-go v1.57.0 h1:iEAcPbUKfJ2Iqz9uN/jEndCNW2+x7OYLHDidXDhPjI0=
github.com/anthropics/anthropic-sdk-go v1.57.0/go.mod h1:3EfIfmFqxH6rbiLcIP4tPFyXL/IHakx2wDG4OU+TIEI=
github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk=
github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg=
github.com/buger/jsonparser v1.2.0 h1:4EFcvK1kD4jyj6YqNK6skK6w+y7FHHBR+XBCtxwu/6g=
github.com/buger/jsonparser v1.2.0/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0=
github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/dnaeon/go-vcr v1.2.0 h1:zHCHvJYTMh1N7xnV7zf1m1GPBF9Ad0Jk/whtQ1663qI=
github.com/dnaeon/go-vcr v1.2.0/go.mod h1:R4UdLID7HZT3taECzJs4YgbbH6PIGXB6W/sc5OLb6RQ=
github.com/go-telegram/bot v1.20.0 h1:4Pea/qTidSspr4WBJw9FbHUMNhYeqszBqQUfsQEyFbc=
github.com/go-telegram/bot v1.20.0/go.mod h1:i2TRs7fXWIeaceF3z7KzsMt/he0TwkVC680mvdTFYeM=
github.com/go-telegram/bot v1.21.0 h1:Va/PbGc2vBDdv57GCUEEVV6ROlHWiC6SklJY9Hvhzps=
github.com/go-telegram/bot v1.21.0/go.mod h1:i2TRs7fXWIeaceF3z7KzsMt/he0TwkVC680mvdTFYeM=
github.com/go-telegram/bot v1.22.0 h1:zK29OoTYMmR5emJrCtGa2SjaGleeZiUB/C1i7kc2lXE=
github.com/go-telegram/bot v1.22.0/go.mod h1:i2TRs7fXWIeaceF3z7KzsMt/he0TwkVC680mvdTFYeM=
github.com/invopop/jsonschema v0.13.0 h1:KvpoAJWEjR3uD9Kbm2HWJmqsEaHt8lBUpd0qHcIi21E=
github.com/invopop/jsonschema v0.13.0/go.mod h1:ffZ5Km5SWWRAIN6wbDXItl95euhFz2uON45H2qjYt+0=
github.com/invopop/jsonschema v0.14.0 h1:MHQqLhvpNUZfw+hM3AZDYK7jxO8FZoQeQM77g8iyZjg=
github.com/invopop/jsonschema v0.14.0/go.mod h1:ygm6C2EaVNMBDPpaPlnOA2pFAxBnxGjFlMZABxm9n2I=
github.com/go-telegram/bot v1.19.0 h1:tuvTQhgNietHFRN0HUDhuXsgfgkGSaO8WWwZQW3DMQg=
github.com/go-telegram/bot v1.19.0/go.mod h1:i2TRs7fXWIeaceF3z7KzsMt/he0TwkVC680mvdTFYeM=
github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E=
github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc=
github.com/jinzhu/now v1.1.5 h1:/o9tlHleP7gOFmsnYNz3RGnqzefHA47wQpKrrdTIwXQ=
github.com/jinzhu/now v1.1.5/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8=
github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=
github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
github.com/mailru/easyjson v0.9.2 h1:dX8U45hQsZpxd80nLvDGihsQ/OxlvTkVUXH2r/8cb2M=
github.com/mailru/easyjson v0.9.2/go.mod h1:1+xMtQp2MRNVL/V1bOzuP3aP8VNwRW55fQUto+XFtTU=
github.com/mattn/go-sqlite3 v1.14.44 h1:3VSe+xafpbzsLbdr2AWlAZk9yRHiBhTBakioXaCKTF8=
github.com/mattn/go-sqlite3 v1.14.44/go.mod h1:pjEuOr8IwzLJP2MfGeTb0A35jauH+C2kbHKBr7yXKVQ=
github.com/mattn/go-sqlite3 v1.14.47 h1:jOBI62gS7nKeZv+as1oGEy0+1qISgXwH/QBlR6KbfIo=
github.com/mattn/go-sqlite3 v1.14.47/go.mod h1:6JTjA44L93a0QCyJef5YvlPoKXntQPjzWv5gtm9sB6w=
github.com/mattn/go-sqlite3 v1.14.48 h1:7XHIgl0a8HwOaiK4E47ozLkST78rR9+OtNGx27D/TFs=
github.com/mattn/go-sqlite3 v1.14.48/go.mod h1:6JTjA44L93a0QCyJef5YvlPoKXntQPjzWv5gtm9sB6w=
github.com/pb33f/ordered-map/v2 v2.3.1 h1:5319HDO0aw4DA4gzi+zv4FXU9UlSs3xGZ40wcP1nBjY=
github.com/pb33f/ordered-map/v2 v2.3.1/go.mod h1:qxFQgd0PkVUtOMCkTapqotNgzRhMPL7VvaHKbd1HnmQ=
github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=
github.com/liushuangls/go-anthropic/v2 v2.17.1 h1:ca3oFzgQHs9/mJr+xx2XFQIYcQLM2rDCqieUx0g+8p4=
github.com/liushuangls/go-anthropic/v2 v2.17.1/go.mod h1:a550cJXPoTG2FL3DvfKG2zzD5O2vjgvo4tHtoGPzFLU=
github.com/mattn/go-sqlite3 v1.14.34 h1:3NtcvcUnFBPsuRcno8pUtupspG/GM+9nZ88zgJcp6Zk=
github.com/mattn/go-sqlite3 v1.14.34/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs=
github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ=
github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc=
github.com/standard-webhooks/standard-webhooks/libraries v0.0.1 h1:uOfcYT+3QungH6tIGSVCR/Y3KJmgJiHcojJbMTPDZAI=
github.com/standard-webhooks/standard-webhooks/libraries v0.0.1/go.mod h1:L1MQhA6x4dn9r007T033lsaZMv9EmBAdXyU/+EF40fo=
github.com/stretchr/objx v0.5.3 h1:jmXUvGomnU1o3W/V5h2VEradbpJDwGrzugQQvL0POH4=
github.com/stretchr/objx v0.5.3/go.mod h1:rDQraq+vQZU7Fde9LOZLr8Tax6zZvy4kuNKF+QYS+U0=
github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
github.com/tidwall/gjson v1.14.2/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=
github.com/tidwall/gjson v1.19.0 h1:xwxm7n691Uf3u5OFjzngavjGTh55KX5q/9w9xHW88JU=
github.com/tidwall/gjson v1.19.0/go.mod h1:V37/opeE/JbLUOfH0QTXiNez2l0RUjYUhpT4szFQAfc=
github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM=
github.com/tidwall/match v1.2.0 h1:0pt8FlkOwjN2fPt4bIl4BoNxb98gGHN2ObFEDkrfZnM=
github.com/tidwall/match v1.2.0/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM=
github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU=
github.com/tidwall/pretty v1.2.1 h1:qjsOFOWWQl+N3RsoF5/ssm1pHmJJwhjlSbZ51I6wMl4=
github.com/tidwall/pretty v1.2.1/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU=
github.com/tidwall/sjson v1.2.5 h1:kLy8mja+1c9jlljvWTlSazM7cKDRfJuR/bOJhcY5NcY=
github.com/tidwall/sjson v1.2.5/go.mod h1:Fvgq9kS/6ociJEDnK0Fk1cpYF4FIW6ZF7LAe+6jwd28=
github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc=
github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw=
go.yaml.in/yaml/v4 v4.0.0-rc.6 h1:1h7H1ohdUh93/FyE4YaDa1Zh64K6VVbjF4K6WUxMtH4=
go.yaml.in/yaml/v4 v4.0.0-rc.6/go.mod h1:aZqd9kCMsGL7AuUv/m/PvWLdg5sjJsZ4oHDEnfPPfY0=
golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4=
golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
golang.org/x/sync v0.21.0 h1:HLII4xRRTtCRkxYp4HNFF0Js/Og6q2i++KXbg0gHCwM=
golang.org/x/sync v0.21.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek=
golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
golang.org/x/text v0.37.0 h1:Cqjiwd9eSg8e0QAkyCaQTNHFIIzWtidPahFWR83rTrc=
golang.org/x/text v0.37.0/go.mod h1:a5sjxXGs9hsn/AJVwuElvCAo9v8QYLzvavO5z2PiM38=
golang.org/x/text v0.38.0 h1:sXmwo9DwP3OK9EZ7PqAdaooSGozfl/3a6/xJcbzPRhE=
golang.org/x/text v0.38.0/go.mod h1:YXZt3QhHUKYT53r2lLKFIVi6Ao1jdzrTR/KQ09qyxF4=
golang.org/x/text v0.40.0 h1:Ub2Z6/xjgF1WrYQz2nuITOEegKFtiIy+rieRJ5lHZKs=
golang.org/x/text v0.40.0/go.mod h1:hpnzDAfGV753zIKo+wk3u1bVKCGPbrnF7+7LBF/UHVY=
golang.org/x/time v0.15.0 h1:bbrp8t3bGUeFOx08pvsMYRTCVSMk89u4tKbNOZbp88U=
golang.org/x/time v0.15.0/go.mod h1:Y4YMaQmXwGQZoFaVFk4YpCt4FLQMYKZe9oeV/f4MSno=
golang.org/x/text v0.34.0 h1:oL/Qq0Kdaqxa1KbNeMKwQq0reLCCaFtqu2eNuSeNHbk=
golang.org/x/text v0.34.0/go.mod h1:homfLqTYRFyVYemLBFl5GgL/DWEiH5wcsQ5gSh1yziA=
golang.org/x/time v0.14.0 h1:MRx4UaLrDotUKUdCIqzPC48t1Y9hANFKIRpNx+Te8PI=
golang.org/x/time v0.14.0/go.mod h1:eL/Oa2bBBK0TkX57Fyni+NgnyQQN4LitPmob2Hjnqw4=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
gopkg.in/yaml.v2 v2.2.8 h1:obN1ZagJSUGI0Ek/LBmuj4SNLPfIny3KsKFopxRdj10=
gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
gorm.io/driver/sqlite v1.6.0 h1:WHRRrIiulaPiPFmDcod6prc4l2VGVWHz80KspNsxSfQ=
gorm.io/driver/sqlite v1.6.0/go.mod h1:AO9V1qIQddBESngQUKWL9yoH93HIeA1X6V633rBwyT8=
gorm.io/gorm v1.31.1 h1:7CA8FTFz/gRfgqgpeKIBcervUn3xSyPUmr6B2WXJ7kg=
gorm.io/gorm v1.31.1/go.mod h1:XyQVbO2k6YkOis7C2437jSit3SsDK72s7n7rsSHd+Gs=
gorm.io/gorm v1.31.2 h1:3o8FXNo9v9S858gil+3LlZA1LkCOzgb4g5BL64FgaCo=
gorm.io/gorm v1.31.2/go.mod h1:XyQVbO2k6YkOis7C2437jSit3SsDK72s7n7rsSHd+Gs=
+96 -248
View File
@@ -7,214 +7,22 @@ import (
"strconv"
"strings"
"github.com/anthropics/anthropic-sdk-go"
"github.com/go-telegram/bot"
"github.com/go-telegram/bot/models"
"golang.org/x/sync/errgroup"
"github.com/liushuangls/go-anthropic/v2"
)
func (b *Bot) handleVoiceMessage(ctx context.Context, message *models.Message, userMsg Message, chatID, userID int64, username, firstName, lastName string, isPremium bool, languageCode string, messageTime int, businessConnectionID string) {
if b.config.ElevenLabsAPIKey == "" {
if err := b.sendResponse(ctx, chatID, "I don't understand voice messages.", businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending voice-unsupported message: %v", err)
}
return
}
if !b.hasScope(userID, ScopeTTSUse) {
if err := b.sendResponse(ctx, chatID, "You don't have permission to use voice features.", businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending permission denied message: %v", err)
}
return
}
transcript, err := b.transcribeVoice(ctx, message.Voice.FileID)
if err != nil {
ErrorLogger.Printf("Error transcribing voice message from user %d: %v", userID, err)
if err := b.sendResponse(ctx, chatID, "Sorry, I couldn't understand your voice message.", businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending transcription error message: %v", err)
}
return
}
if err := b.db.Model(&userMsg).Update("text", transcript).Error; err != nil {
ErrorLogger.Printf("Error updating voice transcript in DB: %v", err)
}
b.chatMemoriesMu.Lock()
if mem, exists := b.chatMemories[chatID]; exists {
for i := len(mem.Messages) - 1; i >= 0; i-- {
if mem.Messages[i].ID == userMsg.ID {
mem.Messages[i].Text = transcript
break
}
}
}
b.chatMemoriesMu.Unlock()
chatMemory := b.getOrCreateChatMemory(chatID)
contextMessages := b.prepareContextMessages(chatMemory)
response, err := b.getAnthropicResponse(ctx, chatID, contextMessages, false, username, firstName, lastName, isPremium, languageCode, messageTime, nil)
if err != nil {
ErrorLogger.Printf("Error getting Anthropic response for voice: %v", err)
if err := b.sendResponse(ctx, chatID, b.anthropicErrorResponse(err, userID), businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending anthropic error response: %v", err)
}
return
}
audioReader, err := b.generateSpeech(ctx, response)
if err != nil {
ErrorLogger.Printf("Error generating speech, falling back to text: %v", err)
if err := b.sendResponse(ctx, chatID, response, businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending text fallback: %v", err)
}
return
}
if _, err := b.screenOutgoingMessage(chatID, response); err != nil {
ErrorLogger.Printf("Error storing assistant voice response: %v", err)
}
params := &bot.SendAudioParams{
ChatID: chatID,
Audio: &models.InputFileUpload{Filename: "response.mp3", Data: audioReader},
}
if businessConnectionID != "" {
params.BusinessConnectionID = businessConnectionID
}
if _, err := b.tgBot.SendAudio(ctx, params); err != nil {
ErrorLogger.Printf("Error sending audio to chat %d: %v", chatID, err)
}
}
func (b *Bot) uploadPhotoFromItem(ctx context.Context, item *models.Message, chatID int64) (string, error) {
photo := largestPhotoSize(item.Photo)
data, err := b.downloadTelegramFile(ctx, photo.FileID)
if err != nil {
return "", fmt.Errorf("download %s: %w", photo.FileID, err)
}
filename := formatUploadFilename(b.botID, chatID, item.ID, "jpg")
return b.uploadImageToAnthropic(ctx, data, filename, "image/jpeg")
}
func (b *Bot) handlePhotoMessage(
ctx context.Context,
items []*models.Message,
chatID, userID int64,
username, firstName, lastName string,
isPremium bool,
languageCode string,
messageTime int,
businessConnectionID string,
) {
if len(items) == 0 {
return
}
uploaded := make([]string, len(items))
caption := ""
g, gctx := errgroup.WithContext(ctx)
for i, item := range items {
if item.Caption != "" {
caption = item.Caption
}
if len(item.Photo) == 0 {
continue
}
i, item := i, item
g.Go(func() error {
fileID, err := b.uploadPhotoFromItem(gctx, item, chatID)
if err != nil {
return err
}
uploaded[i] = fileID
return nil
})
}
if err := g.Wait(); err != nil {
ErrorLogger.Printf("[%s] photo upload failed: %v", b.config.ID, err)
var successful []string
for _, fid := range uploaded {
if fid != "" {
successful = append(successful, fid)
}
}
b.compensatingDelete(ctx, successful)
if sendErr := b.sendResponse(ctx, chatID, "Sorry, I couldn't process one of your images.", businessConnectionID); sendErr != nil {
ErrorLogger.Printf("Error sending photo failure message: %v", sendErr)
}
return
}
finalUploaded := make([]string, 0, len(uploaded))
for _, fid := range uploaded {
if fid != "" {
finalUploaded = append(finalUploaded, fid)
}
}
if len(finalUploaded) == 0 {
return
}
chatMemory := b.getOrCreateChatMemory(chatID)
userMessage := b.createMessage(chatID, userID, username, "user", caption, true)
userMessage.ImageFileIDs = finalUploaded
if err := b.storeMessage(&userMessage); err != nil {
b.compensatingDelete(ctx, finalUploaded)
ErrorLogger.Printf("[%s] store photo message failed: %v", b.config.ID, err)
if sendErr := b.sendResponse(ctx, chatID, "Sorry, I had trouble saving your message.", businessConnectionID); sendErr != nil {
ErrorLogger.Printf("Error sending store failure message: %v", sendErr)
}
return
}
b.addMessageToChatMemory(chatMemory, userMessage)
contextMessages := b.prepareContextMessages(chatMemory)
joined, err := b.getAnthropicResponse(
ctx, chatID, contextMessages, false,
username, firstName, lastName, isPremium, languageCode, messageTime,
func(seg string) error {
return b.sendOneSegment(ctx, chatID, seg, businessConnectionID)
},
)
if err != nil {
ErrorLogger.Printf("Error getting Anthropic response for photo: %v", err)
if sendErr := b.sendResponse(ctx, chatID, b.anthropicErrorResponse(err, userID), businessConnectionID); sendErr != nil {
ErrorLogger.Printf("Error sending anthropic error response: %v", sendErr)
}
return
}
if _, storeErr := b.screenOutgoingMessage(chatID, joined); storeErr != nil {
ErrorLogger.Printf("Error recording assistant turn: %v", storeErr)
}
}
// anthropicErrorResponse returns the message to send back to the user when getAnthropicResponse
// fails. Admins and owners receive an actionable hint when the model is deprecated; regular users
// always get the generic fallback to avoid leaking internal details.
func (b *Bot) anthropicErrorResponse(err error, userID int64) string {
isElevated := b.hasScope(userID, ScopeModelSet)
if errors.Is(err, ErrModelNotFound) && isElevated {
if errors.Is(err, ErrModelNotFound) && b.hasScope(userID, ScopeModelSet) {
return fmt.Sprintf(
"⚠️ Model `%s` is no longer available (deprecated or removed by Anthropic).\n"+
"Use /set_model <model-id> to switch. Current models: https://platform.claude.com/docs/en/about-claude/models/overview",
b.config.Model,
)
}
if isElevated {
var apiErr *anthropic.Error
if errors.As(err, &apiErr) {
body := apiErr.RawJSON()
if len(body) > 800 {
body = body[:800] + "...(truncated)"
}
out := fmt.Sprintf("⚠️ Anthropic API error %d:\n%s", apiErr.StatusCode, body)
if apiErr.RequestID != "" {
out += fmt.Sprintf("\nRequest-ID: %s", apiErr.RequestID)
}
return out
}
return fmt.Sprintf("⚠️ Anthropic call failed: %v", err)
}
return "I'm sorry, I'm having trouble processing your request right now."
}
@@ -226,9 +34,11 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U
} else if update.BusinessMessage != nil {
message = update.BusinessMessage
} else {
// No message to process
return
}
// Extract businessConnectionID if available
var businessConnectionID string
if update.BusinessConnection != nil {
businessConnectionID = update.BusinessConnection.ID
@@ -237,6 +47,8 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U
}
if message.From == nil {
// Channel posts and some automated messages have no sender — ignore them.
// see: https://core.telegram.org/bots/api#message
return
}
@@ -250,17 +62,31 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U
messageTime := message.Date
text := message.Text
// Check if it's a new chat (before storing the message so the flag is accurate).
isNewChatFlag := b.isNewChat(chatID)
// Screen incoming message (store to DB + add to chat memory)
userMsg, err := b.screenIncomingMessage(message)
if err != nil {
ErrorLogger.Printf("Error storing user message: %v", err)
return
}
// Determine if the user is the owner
var isOwner bool
if b.db.Where("telegram_id = ? AND bot_id = ? AND is_owner = ?", userID, b.botID, true).First(&User{}).Error == nil {
err = b.db.Where("telegram_id = ? AND bot_id = ? AND is_owner = ?", userID, b.botID, true).First(&User{}).Error
if err == nil {
isOwner = true
}
// Always create/get the user record — on the very first message and on all subsequent ones.
user, err := b.getOrCreateUser(userID, username, isOwner)
if err != nil {
ErrorLogger.Printf("Error getting or creating user: %v", err)
return
}
// Update the username if it has changed
if user.Username != username {
user.Username = username
if err := b.db.Save(&user).Error; err != nil {
@@ -268,45 +94,27 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U
}
}
if message.MediaGroupID != "" && len(message.Photo) > 0 {
b.bufferAlbumItem(ctx, message, chatID, userID, username, firstName, lastName,
isPremium, languageCode, messageTime, businessConnectionID)
return
}
if len(message.Photo) > 0 {
if !b.checkRateLimits(userID) {
b.sendRateLimitExceededMessage(ctx, chatID, businessConnectionID)
return
}
b.handlePhotoMessage(ctx, []*models.Message{message},
chatID, userID, username, firstName, lastName,
isPremium, languageCode, messageTime,
businessConnectionID)
return
}
userMsg, err := b.screenIncomingMessage(message)
if err != nil {
ErrorLogger.Printf("Error storing user message: %v", err)
return
}
// Check if the message is a command — applies on every message, including the very first.
if message.Entities != nil {
for _, entity := range message.Entities {
if entity.Type == "bot_command" {
command := strings.TrimSpace(message.Text[entity.Offset : entity.Offset+entity.Length])
switch command {
case "/stats":
// Parse command parameters
parts := strings.Fields(message.Text)
// Default: show global stats
if len(parts) == 1 {
b.sendStats(ctx, chatID, userID, 0, businessConnectionID)
return
}
// Check for "user" parameter
if len(parts) >= 2 && parts[1] == "user" {
targetUserID := userID
targetUserID := userID // Default to current user
// If a user ID is provided, parse it
if len(parts) >= 3 {
var parseErr error
targetUserID, parseErr = strconv.ParseInt(parts[2], 10, 64)
@@ -323,6 +131,7 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U
return
}
// Invalid parameter
if err := b.sendResponse(ctx, chatID, "Invalid command format. Usage: /stats or /stats user [user_id]", businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending response: %v", err)
}
@@ -372,6 +181,11 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U
return
}
newModel := strings.TrimSpace(parts[1])
// No upfront model validation:
// - The go-anthropic library constants are not enumerable at runtime (Go has no const reflection).
// - A live /v1/models probe would add a network round-trip and show in the API audit log.
// - An invalid model ID will produce a 404 on the next real message, which routes through
// anthropicErrorResponse and already delivers an actionable admin-facing hint.
if err := b.config.PersistModel(newModel); err != nil {
ErrorLogger.Printf("Failed to persist model change: %v", err)
if err := b.sendResponse(ctx, chatID, fmt.Sprintf("Model updated in memory to `%s`, but failed to save to config file: %v", newModel, err), businessConnectionID); err != nil {
@@ -416,48 +230,42 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U
}
}
// Rate limit check applies to all message types including stickers.
if !b.checkRateLimits(userID) {
b.sendRateLimitExceededMessage(ctx, chatID, businessConnectionID)
return
}
if message.Voice != nil {
b.handleVoiceMessage(ctx, message, userMsg, chatID, userID, username, firstName, lastName, isPremium, languageCode, messageTime, businessConnectionID)
return
}
// Build context once — shared by the sticker and text response paths.
chatMemory := b.getOrCreateChatMemory(chatID)
contextMessages := b.prepareContextMessages(chatMemory)
// Check if the message contains a sticker
if message.Sticker != nil {
b.handleStickerMessage(ctx, chatID, userMsg, message, contextMessages, businessConnectionID)
return
}
// Proceed only if the message contains text
if text == "" {
InfoLogger.Printf("Received a non-text message from user %d in chat %d", userID, chatID)
return
}
// Determine if the text contains only emojis
isEmojiOnly := isOnlyEmojis(text)
joined, err := b.getAnthropicResponse(
ctx, chatID, contextMessages, isEmojiOnly,
username, firstName, lastName, isPremium, languageCode, messageTime,
func(seg string) error {
return b.sendOneSegment(ctx, chatID, seg, businessConnectionID)
},
)
// Get response from Anthropic
response, err := b.getAnthropicResponse(ctx, contextMessages, isNewChatFlag, isOwner, isEmojiOnly, username, firstName, lastName, isPremium, languageCode, messageTime)
if err != nil {
ErrorLogger.Printf("Error getting Anthropic response: %v", err)
if sendErr := b.sendResponse(ctx, chatID, b.anthropicErrorResponse(err, userID), businessConnectionID); sendErr != nil {
ErrorLogger.Printf("Error sending response: %v", sendErr)
}
return
response = b.anthropicErrorResponse(err, userID)
}
if _, storeErr := b.screenOutgoingMessage(chatID, joined); storeErr != nil {
ErrorLogger.Printf("Error recording assistant turn: %v", storeErr)
// Send the response
if err := b.sendResponse(ctx, chatID, response, businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending response: %v", err)
return
}
}
@@ -467,11 +275,14 @@ func (b *Bot) sendRateLimitExceededMessage(ctx context.Context, chatID int64, bu
}
}
func (b *Bot) handleStickerMessage(ctx context.Context, chatID int64, userMessage Message, message *models.Message, contextMessages []anthropic.BetaMessageParam, businessConnectionID string) {
func (b *Bot) handleStickerMessage(ctx context.Context, chatID int64, userMessage Message, message *models.Message, contextMessages []anthropic.Message, businessConnectionID string) {
// userMessage was already screened (stored + added to memory) by handleUpdate — do not call screenIncomingMessage again.
// Generate AI response about the sticker
response, err := b.generateStickerResponse(ctx, userMessage, contextMessages)
if err != nil {
ErrorLogger.Printf("Error generating sticker response: %v", err)
// Provide a fallback dynamic response based on sticker type
if message.Sticker.IsAnimated {
response = "Wow, that's a cool animated sticker!"
} else if message.Sticker.IsVideo {
@@ -481,16 +292,19 @@ func (b *Bot) handleStickerMessage(ctx context.Context, chatID int64, userMessag
}
}
// Send the response
if err := b.sendResponse(ctx, chatID, response, businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending response: %v", err)
return
}
}
func (b *Bot) generateStickerResponse(ctx context.Context, message Message, contextMessages []anthropic.BetaMessageParam) (string, error) {
func (b *Bot) generateStickerResponse(ctx context.Context, message Message, contextMessages []anthropic.Message) (string, error) {
// contextMessages already contains the sticker turn (added by screenIncomingMessage as
// "Sent a sticker: <emoji>"), so the full conversation history is preserved.
if message.StickerFileID != "" {
messageTime := int(message.Timestamp.Unix())
response, err := b.getAnthropicResponse(ctx, message.ChatID, contextMessages, true, message.Username, "", "", false, "", messageTime, nil)
response, err := b.getAnthropicResponse(ctx, contextMessages, false, false, true, message.Username, "", "", false, "", messageTime)
if err != nil {
return "", err
}
@@ -501,6 +315,7 @@ func (b *Bot) generateStickerResponse(ctx context.Context, message Message, cont
}
func (b *Bot) clearChatHistory(ctx context.Context, chatID int64, currentUserID int64, targetUserID int64, targetChatID int64, businessConnectionID string, hardDelete bool) {
// If targetUserID is provided and different from currentUserID, check permissions
if targetUserID != 0 && targetUserID != currentUserID {
requiredScope := ScopeHistoryClearAny
if hardDelete {
@@ -514,6 +329,7 @@ func (b *Bot) clearChatHistory(ctx context.Context, chatID int64, currentUserID
return
}
// Check if the target user exists
var targetUser User
err := b.db.Where("telegram_id = ? AND bot_id = ?", targetUserID, b.botID).First(&targetUser).Error
if err != nil {
@@ -524,34 +340,58 @@ func (b *Bot) clearChatHistory(ctx context.Context, chatID int64, currentUserID
return
}
} else {
// If no targetUserID is provided, set it to currentUserID
targetUserID = currentUserID
}
// Delete messages from the database
//
// Assumption: this bot is primarily used in private DMs, where each user's messages
// are stored with chat_id == their own user_id — not the caller's chat_id. Scoping
// a cross-user delete by the caller's chatID would therefore match 0 rows.
//
// When clearing another user's history the default (targetChatID == 0) deletes all
// of that user's messages across every chat for this bot — the natural meaning of
// "/clear <userID>" (wipe their entire history with the bot).
//
// When targetChatID != 0 the deletion is scoped to that specific chat, which is
// useful for group moderation ("/clear <userID> <chatID>").
var err error
if hardDelete {
// Permanently delete messages
if targetUserID == currentUserID {
err = b.hardDeleteScope(ctx, "chat_id = ? AND bot_id = ?", chatID, b.botID)
// Own history — delete ALL messages (user + assistant) in the current chat.
err = b.db.Unscoped().Where("chat_id = ? AND bot_id = ?", chatID, b.botID).Delete(&Message{}).Error
InfoLogger.Printf("User %d permanently deleted their own chat history in chat %d", currentUserID, chatID)
} else {
if targetChatID != 0 {
err = b.hardDeleteScope(ctx, "chat_id = ? AND bot_id = ?", targetChatID, b.botID)
// Chat-scoped: delete ALL messages (user + assistant) in the specified chat.
err = b.db.Unscoped().Where("chat_id = ? AND bot_id = ?", targetChatID, b.botID).Delete(&Message{}).Error
InfoLogger.Printf("Admin/owner %d permanently deleted chat history for user %d in chat %d", currentUserID, targetUserID, targetChatID)
} else {
err = b.hardDeleteScope(ctx,
"bot_id = ? AND (user_id = ? OR (chat_id = ? AND is_user = ?))",
b.botID, targetUserID, targetUserID, false)
// Bot-wide: delete all of the user's own messages across every chat, then delete
// assistant messages from their DM chat (where chat_id == user_id by Telegram convention).
err = b.db.Unscoped().Where("bot_id = ? AND user_id = ?", b.botID, targetUserID).Delete(&Message{}).Error
if err == nil {
err = b.db.Unscoped().Where("chat_id = ? AND bot_id = ? AND is_user = ?", targetUserID, b.botID, false).Delete(&Message{}).Error
}
InfoLogger.Printf("Admin/owner %d permanently deleted all chat history for user %d", currentUserID, targetUserID)
}
}
} else {
// Soft delete messages
if targetUserID == currentUserID {
// Own history — delete ALL messages (user + assistant) in the current chat.
err = b.db.Where("chat_id = ? AND bot_id = ?", chatID, b.botID).Delete(&Message{}).Error
InfoLogger.Printf("User %d soft deleted their own chat history in chat %d", currentUserID, chatID)
} else {
if targetChatID != 0 {
// Chat-scoped: delete ALL messages (user + assistant) in the specified chat.
err = b.db.Where("chat_id = ? AND bot_id = ?", targetChatID, b.botID).Delete(&Message{}).Error
InfoLogger.Printf("Admin/owner %d soft deleted chat history for user %d in chat %d", currentUserID, targetUserID, targetChatID)
} else {
// Bot-wide: delete all of the user's own messages across every chat, then delete
// assistant messages from their DM chat (where chat_id == user_id by Telegram convention).
err = b.db.Where("bot_id = ? AND user_id = ?", b.botID, targetUserID).Delete(&Message{}).Error
if err == nil {
err = b.db.Where("chat_id = ? AND bot_id = ? AND is_user = ?", targetUserID, b.botID, false).Delete(&Message{}).Error
@@ -569,20 +409,28 @@ func (b *Bot) clearChatHistory(ctx context.Context, chatID int64, currentUserID
return
}
// Evict the relevant in-memory cache entry so the next access rebuilds from
// the now-clean DB. Applies to all cases: own history, cross-user
// scoped to a specific chat, and bot-wide cross-user clear.
b.chatMemoriesMu.Lock()
if targetUserID == currentUserID {
// Own history is always scoped to the current chat.
delete(b.chatMemories, chatID)
} else if targetChatID != 0 {
// Admin cleared a specific chat — evict that chat's cache.
delete(b.chatMemories, targetChatID)
} else {
// Bot-wide clear: primary use-case is DMs where chatID == userID.
delete(b.chatMemories, targetUserID)
}
b.chatMemoriesMu.Unlock()
// Send a confirmation message
var confirmationMessage string
if targetUserID == currentUserID {
confirmationMessage = "Your chat history has been cleared."
} else {
// Get the username of the target user if available
var targetUser User
err := b.db.Where("telegram_id = ? AND bot_id = ?", targetUserID, b.botID).First(&targetUser).Error
if err == nil && targetUser.Username != "" {
+93 -32
View File
@@ -17,6 +17,7 @@ import (
)
func TestHandleUpdate_NewChat(t *testing.T) {
// Setup
db := setupTestDB(t)
mockClock := &MockClock{
currentTime: time.Now(),
@@ -24,7 +25,7 @@ func TestHandleUpdate_NewChat(t *testing.T) {
config := BotConfig{
ID: "test_bot",
OwnerTelegramID: 123,
OwnerTelegramID: 123, // owner's ID
TelegramToken: "test_token",
MemorySize: 10,
MessagePerHour: 5,
@@ -36,6 +37,7 @@ func TestHandleUpdate_NewChat(t *testing.T) {
mockTgClient := &MockTelegramClient{}
// Create bot model first
botModel := &BotModel{
Identifier: config.ID,
Name: config.ID,
@@ -43,6 +45,7 @@ func TestHandleUpdate_NewChat(t *testing.T) {
err := db.Create(botModel).Error
assert.NoError(t, err)
// Create bot config
configModel := &ConfigModel{
BotID: botModel.ID,
MemorySize: config.MemorySize,
@@ -56,34 +59,40 @@ func TestHandleUpdate_NewChat(t *testing.T) {
err = db.Create(configModel).Error
assert.NoError(t, err)
// Create bot instance
b, err := NewBot(db, config, mockClock, mockTgClient)
assert.NoError(t, err)
testCases := []struct {
name string
userID int64
wantSubstr string
name string
userID int64
isOwner bool
wantResp string
}{
{
name: "Owner First Message",
userID: 123,
wantSubstr: "Anthropic call failed:",
name: "Owner First Message",
userID: 123, // owner's ID
isOwner: true,
wantResp: "I'm sorry, I'm having trouble processing your request right now.",
},
{
name: "Regular User First Message",
userID: 456,
wantSubstr: "I'm sorry, I'm having trouble processing your request right now.",
name: "Regular User First Message",
userID: 456,
isOwner: false,
wantResp: "I'm sorry, I'm having trouble processing your request right now.",
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
// Setup mock response expectations for error case to test fallback messages
mockTgClient.SendMessageFunc = func(ctx context.Context, params *bot.SendMessageParams) (*models.Message, error) {
assert.Equal(t, tc.userID, params.ChatID)
assert.Contains(t, params.Text, tc.wantSubstr)
assert.Equal(t, tc.wantResp, params.Text)
return &models.Message{}, nil
}
// Create update with new message
update := &models.Update{
Message: &models.Message{
Chat: models.Chat{ID: tc.userID},
@@ -95,23 +104,24 @@ func TestHandleUpdate_NewChat(t *testing.T) {
},
}
// Handle the update
b.handleUpdate(context.Background(), nil, update)
// Verify message was stored
var storedMsg Message
err := db.Where("chat_id = ? AND user_id = ? AND text = ?", tc.userID, tc.userID, "Hello").First(&storedMsg).Error
assert.NoError(t, err)
// Verify response was stored
var respMsg Message
err = db.Where("chat_id = ? AND is_user = ?", tc.userID, false).
Order("timestamp DESC").
First(&respMsg).Error
err = db.Where("chat_id = ? AND is_user = ? AND text = ?", tc.userID, false, tc.wantResp).First(&respMsg).Error
assert.NoError(t, err)
assert.Contains(t, respMsg.Text, tc.wantSubstr)
})
}
}
func TestClearChatHistory(t *testing.T) {
// Setup
db := setupTestDB(t)
mockClock := &MockClock{
currentTime: time.Now(),
@@ -119,7 +129,7 @@ func TestClearChatHistory(t *testing.T) {
config := BotConfig{
ID: "test_bot",
OwnerTelegramID: 123,
OwnerTelegramID: 123, // owner's ID
TelegramToken: "test_token",
MemorySize: 10,
MessagePerHour: 5,
@@ -131,6 +141,7 @@ func TestClearChatHistory(t *testing.T) {
mockTgClient := &MockTelegramClient{}
// Create bot model first
botModel := &BotModel{
Identifier: config.ID,
Name: config.ID,
@@ -138,6 +149,7 @@ func TestClearChatHistory(t *testing.T) {
err := db.Create(botModel).Error
assert.NoError(t, err)
// Create bot config
configModel := &ConfigModel{
BotID: botModel.ID,
MemorySize: config.MemorySize,
@@ -151,18 +163,22 @@ func TestClearChatHistory(t *testing.T) {
err = db.Create(configModel).Error
assert.NoError(t, err)
// Create bot instance
b, err := NewBot(db, config, mockClock, mockTgClient)
assert.NoError(t, err)
// Create test users
ownerID := int64(123)
adminID := int64(456)
regularUserID := int64(789)
nonExistentUserID := int64(999)
chatID := int64(1000)
// Create admin role
adminRole, err := b.getRoleByName("admin")
assert.NoError(t, err)
// Create admin user
adminUser := User{
BotID: b.botID,
TelegramID: adminID,
@@ -174,6 +190,7 @@ func TestClearChatHistory(t *testing.T) {
err = db.Create(&adminUser).Error
assert.NoError(t, err)
// Create regular user
regularRole, err := b.getRoleByName("user")
assert.NoError(t, err)
regularUser := User{
@@ -187,11 +204,15 @@ func TestClearChatHistory(t *testing.T) {
err = db.Create(&regularUser).Error
assert.NoError(t, err)
// Create test messages for each user.
// Each user's messages are stored with chat_id == their own user_id, mirroring
// how Telegram private DMs work (chat_id == user_id in 1-on-1 bot conversations).
// Using a shared artificial chatID here would mask the cross-user delete bug.
for _, userID := range []int64{ownerID, adminID, regularUserID} {
for i := 0; i < 5; i++ {
message := Message{
BotID: b.botID,
ChatID: userID,
ChatID: userID, // per-user chat, not a shared chatID
UserID: userID,
Username: "test",
UserRole: "user",
@@ -204,6 +225,7 @@ func TestClearChatHistory(t *testing.T) {
}
}
// Test cases
testCases := []struct {
name string
currentUserID int64
@@ -266,7 +288,7 @@ func TestClearChatHistory(t *testing.T) {
targetUserID: adminID,
hardDelete: false,
expectedError: true,
expectedCount: 5,
expectedCount: 5, // Messages should remain
expectedMsg: "Permission denied. Only admins and owners can clear other users' histories.",
},
{
@@ -275,7 +297,7 @@ func TestClearChatHistory(t *testing.T) {
targetUserID: nonExistentUserID,
hardDelete: false,
expectedError: true,
expectedCount: 5,
expectedCount: 5, // Messages should remain for admin
expectedMsg: "User with ID 999 not found.",
},
{
@@ -288,23 +310,29 @@ func TestClearChatHistory(t *testing.T) {
expectedMsg: "Chat history for user @regular (ID: 789) has been cleared.",
},
{
// targetChatID scopes the delete to a specific chat; messages in other chats survive.
// We seed messages with ChatID == userID (per-user DM), so targeting a different chatID
// should leave the user's messages untouched (expectedCount == 5).
name: "Admin clears regular user's history scoped to non-matching chat",
currentUserID: adminID,
targetUserID: regularUserID,
targetChatID: int64(9999),
targetChatID: int64(9999), // a chat the user has no messages in
hardDelete: false,
expectedError: false,
expectedCount: 5,
expectedCount: 5, // messages in chat 789 are unaffected
expectedMsg: "Chat history for user @regular (ID: 789) has been cleared.",
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
// Reset messages for the test case
if tc.name != "Owner hard deletes regular user's history" {
// Delete all messages for the target user
err = db.Where("user_id = ?", tc.targetUserID).Delete(&Message{}).Error
assert.NoError(t, err)
// Recreate messages for the target user
for i := 0; i < 5; i++ {
message := Message{
BotID: b.botID,
@@ -321,16 +349,20 @@ func TestClearChatHistory(t *testing.T) {
}
}
// Setup mock response expectations
var sentMessage string
mockTgClient.SendMessageFunc = func(ctx context.Context, params *bot.SendMessageParams) (*models.Message, error) {
sentMessage = params.Text
return &models.Message{}, nil
}
// Call the clearChatHistory method
b.clearChatHistory(context.Background(), chatID, tc.currentUserID, tc.targetUserID, tc.targetChatID, tc.businessConnID, tc.hardDelete)
// Verify the response message
assert.Equal(t, tc.expectedMsg, sentMessage)
// Count remaining messages for the target user
var count int64
if tc.hardDelete {
db.Unscoped().Model(&Message{}).Where("user_id = ? AND chat_id = ?", tc.targetUserID, chatID).Count(&count)
@@ -343,6 +375,7 @@ func TestClearChatHistory(t *testing.T) {
}
func TestStatsCommand(t *testing.T) {
// Setup
db := setupTestDB(t)
mockClock := &MockClock{
currentTime: time.Now(),
@@ -350,7 +383,7 @@ func TestStatsCommand(t *testing.T) {
config := BotConfig{
ID: "test_bot",
OwnerTelegramID: 123,
OwnerTelegramID: 123, // owner's ID
TelegramToken: "test_token",
MemorySize: 10,
MessagePerHour: 5,
@@ -362,6 +395,7 @@ func TestStatsCommand(t *testing.T) {
mockTgClient := &MockTelegramClient{}
// Create bot model first
botModel := &BotModel{
Identifier: config.ID,
Name: config.ID,
@@ -369,6 +403,7 @@ func TestStatsCommand(t *testing.T) {
err := db.Create(botModel).Error
assert.NoError(t, err)
// Create bot config
configModel := &ConfigModel{
BotID: botModel.ID,
MemorySize: config.MemorySize,
@@ -382,17 +417,21 @@ func TestStatsCommand(t *testing.T) {
err = db.Create(configModel).Error
assert.NoError(t, err)
// Create bot instance
b, err := NewBot(db, config, mockClock, mockTgClient)
assert.NoError(t, err)
// Create test users
ownerID := int64(123)
adminID := int64(456)
regularUserID := int64(789)
chatID := int64(1000)
// Create admin role
adminRole, err := b.getRoleByName("admin")
assert.NoError(t, err)
// Create admin user
adminUser := User{
BotID: b.botID,
TelegramID: adminID,
@@ -404,6 +443,7 @@ func TestStatsCommand(t *testing.T) {
err = db.Create(&adminUser).Error
assert.NoError(t, err)
// Create regular user
regularRole, err := b.getRoleByName("user")
assert.NoError(t, err)
regularUser := User{
@@ -417,8 +457,10 @@ func TestStatsCommand(t *testing.T) {
err = db.Create(&regularUser).Error
assert.NoError(t, err)
// Create test messages for each user
for _, userID := range []int64{ownerID, adminID, regularUserID} {
for i := 0; i < 5; i++ {
// User message
userMessage := Message{
BotID: b.botID,
ChatID: chatID,
@@ -432,6 +474,7 @@ func TestStatsCommand(t *testing.T) {
err = db.Create(&userMessage).Error
assert.NoError(t, err)
// Bot response
botMessage := Message{
BotID: b.botID,
ChatID: chatID,
@@ -447,6 +490,7 @@ func TestStatsCommand(t *testing.T) {
}
}
// Test cases
testCases := []struct {
name string
command string
@@ -515,12 +559,14 @@ func TestStatsCommand(t *testing.T) {
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
// Setup mock response expectations
var sentMessage string
mockTgClient.SendMessageFunc = func(ctx context.Context, params *bot.SendMessageParams) (*models.Message, error) {
sentMessage = params.Text
return &models.Message{}, nil
}
// Create update with command
update := &models.Update{
Message: &models.Message{
Chat: models.Chat{ID: chatID},
@@ -533,19 +579,22 @@ func TestStatsCommand(t *testing.T) {
{
Type: "bot_command",
Offset: 0,
Length: 6,
Length: 6, // Length of "/stats"
},
},
},
}
// Handle the update
b.handleUpdate(context.Background(), nil, update)
// Verify the response message contains the expected text
assert.Contains(t, sentMessage, tc.expectedMsg)
})
}
}
// Helper function to get username by ID for test
func getUsernameByID(id int64) string {
switch id {
case 123:
@@ -565,11 +614,13 @@ func setupTestDB(t *testing.T) *gorm.DB {
t.Fatalf("Failed to open test database: %v", err)
}
// AutoMigrate the models
err = db.AutoMigrate(&BotModel{}, &ConfigModel{}, &Message{}, &User{}, &Role{}, &Scope{})
if err != nil {
t.Fatalf("Failed to migrate database schema: %v", err)
}
// Create default roles and scopes
err = createDefaultRoles(db)
if err != nil {
t.Fatalf("Failed to create default roles: %v", err)
@@ -581,6 +632,8 @@ func setupTestDB(t *testing.T) *gorm.DB {
return db
}
// setupBotForTest creates a minimal Bot instance backed by an in-memory DB.
// It follows the same pattern as the existing handler tests to avoid duplication.
func setupBotForTest(t *testing.T, ownerID int64) (*Bot, *MockTelegramClient) {
t.Helper()
db := setupTestDB(t)
@@ -615,9 +668,13 @@ func setupBotForTest(t *testing.T, ownerID int64) (*Bot, *MockTelegramClient) {
return b, mockTgClient
}
// TestAnthropicErrorResponse verifies that model-deprecation errors surface actionable
// details only to admin/owner, and that regular users and non-model errors always get
// the generic fallback.
func TestAnthropicErrorResponse(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
b, _ := setupBotForTest(t, 123)
// Create admin user
adminRole, err := b.getRoleByName("admin")
assert.NoError(t, err)
assert.NoError(t, b.db.Create(&User{
@@ -625,6 +682,7 @@ func TestAnthropicErrorResponse(t *testing.T) { //NOSONAR go:S100 -- underscore
RoleID: adminRole.ID, Role: adminRole,
}).Error)
// Create regular user
userRole, err := b.getRoleByName("user")
assert.NoError(t, err)
assert.NoError(t, b.db.Create(&User{
@@ -662,18 +720,11 @@ func TestAnthropicErrorResponse(t *testing.T) { //NOSONAR go:S100 -- underscore
wantMissing: "/set_model",
},
{
name: "owner receives elevated detail for non-API error",
name: "owner receives generic message for non-model error",
err: otherErr,
userID: 123,
wantSubstr: "Anthropic call failed:",
wantMissing: "I'm sorry",
},
{
name: "regular user receives generic message for non-model error",
err: otherErr,
userID: 789,
wantSubstr: "I'm sorry",
wantMissing: "Anthropic call failed",
wantMissing: "/set_model",
},
}
@@ -688,9 +739,12 @@ func TestAnthropicErrorResponse(t *testing.T) { //NOSONAR go:S100 -- underscore
}
}
// TestSetModelCommand verifies that /set_model enforces permissions, validates input,
// updates the model in memory, and persists the change to the config file on disk.
func TestSetModelCommand(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
b, mockTgClient := setupBotForTest(t, 123)
// Point the config at a temporary file so PersistModel can write to disk.
tempDir, err := os.MkdirTemp("", "set_model_cmd_test")
assert.NoError(t, err)
defer func() { _ = os.RemoveAll(tempDir) }()
@@ -700,6 +754,7 @@ func TestSetModelCommand(t *testing.T) { //NOSONAR go:S100 -- underscore separat
assert.NoError(t, os.WriteFile(configPath, []byte(initialJSON), 0600))
b.config.ConfigFilePath = configPath
// Create admin and regular users
adminRole, err := b.getRoleByName("admin")
assert.NoError(t, err)
assert.NoError(t, b.db.Create(&User{
@@ -715,6 +770,8 @@ func TestSetModelCommand(t *testing.T) { //NOSONAR go:S100 -- underscore separat
chatID := int64(1000)
// Seed chat 1000 with a prior message so isNewChatFlag is false for all subtests.
// Commands are only processed in the non-new-chat branch of handleUpdate.
assert.NoError(t, b.db.Create(&Message{
BotID: b.botID, ChatID: chatID, UserID: 789, Username: "regular",
UserRole: "user", Text: "hello", IsUser: true,
@@ -777,6 +834,7 @@ func TestSetModelCommand(t *testing.T) { //NOSONAR go:S100 -- underscore separat
})
}
// Verify the successful update took effect in memory and on disk.
t.Run("model change persisted in memory and on disk", func(t *testing.T) {
assert.Equal(t, "claude-sonnet-4-6", string(b.config.Model))
data, err := os.ReadFile(configPath)
@@ -785,10 +843,12 @@ func TestSetModelCommand(t *testing.T) { //NOSONAR go:S100 -- underscore separat
})
}
// TestHasScope verifies that scope checks honour role assignments and the owner bypass.
func TestHasScope(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
const ownerID int64 = 100
b, _ := setupBotForTest(t, ownerID)
// Admin user
adminRole, err := b.getRoleByName("admin")
assert.NoError(t, err)
assert.NoError(t, b.db.Create(&User{
@@ -796,6 +856,7 @@ func TestHasScope(t *testing.T) { //NOSONAR go:S100 -- underscore separation is
RoleID: adminRole.ID, Role: adminRole,
}).Error)
// Regular user
userRole, err := b.getRoleByName("user")
assert.NoError(t, err)
assert.NoError(t, b.db.Create(&User{
+10
View File
@@ -5,13 +5,23 @@ import (
"os"
)
// For log management, use journalctl commands:
// - View logs: journalctl -u telegram-bot
// - Follow logs: journalctl -u telegram-bot -f
// - View errors: journalctl -u telegram-bot -p err
// Refer to the documentation for details on systemd unit setup.
// Initialize loggers for informational and error messages.
var (
InfoLogger *log.Logger
ErrorLogger *log.Logger
)
// initLoggers sets up separate loggers for stdout and stderr.
func initLoggers() {
// InfoLogger writes to stdout with specific flags.
InfoLogger = log.New(os.Stdout, "INFO: ", log.Ldate|log.Ltime|log.Lshortfile)
// ErrorLogger writes to stderr with specific flags.
ErrorLogger = log.New(os.Stderr, "ERROR: ", log.Ldate|log.Ltime|log.Lshortfile)
}
+11
View File
@@ -8,30 +8,38 @@ import (
)
func main() {
// Initialize custom loggers
initLoggers()
// Log the start of the application
InfoLogger.Println("Starting Telegram Bot Application")
// Initialize database
db, err := initDB()
if err != nil {
ErrorLogger.Fatalf("Error initializing database: %v", err)
}
// Load all bot configurations
configs, err := loadAllConfigs("config")
if err != nil {
ErrorLogger.Fatalf("Error loading configurations: %v", err)
}
// Create a WaitGroup to manage goroutines
var wg sync.WaitGroup
// Set up context with cancellation
ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
defer cancel()
// Initialize and start each bot
for _, config := range configs {
wg.Add(1)
go func(cfg BotConfig) {
defer wg.Done()
// Create Bot instance without TelegramClient initially
realClock := RealClock{}
bot, err := NewBot(db, cfg, realClock, nil)
if err != nil {
@@ -39,14 +47,17 @@ func main() {
return
}
// Start the bot in a separate goroutine
go bot.Start(ctx)
// Keep the bot running until the context is cancelled
<-ctx.Done()
InfoLogger.Printf("Bot %s stopped", cfg.ID)
}(config)
}
// Wait for all bots to finish
wg.Wait()
InfoLogger.Println("All bots have stopped. Exiting application.")
+18 -18
View File
@@ -8,10 +8,10 @@ import (
type BotModel struct {
gorm.Model
Identifier string `gorm:"uniqueIndex"`
Identifier string `gorm:"uniqueIndex"` // Renamed from ID to Identifier
Name string
Configs []ConfigModel `gorm:"foreignKey:BotID;constraint:OnDelete:CASCADE"`
Users []User `gorm:"foreignKey:BotID;constraint:OnDelete:CASCADE"`
Users []User `gorm:"foreignKey:BotID;constraint:OnDelete:CASCADE"` // Associated users
Messages []Message `gorm:"foreignKey:BotID;constraint:OnDelete:CASCADE"`
}
@@ -22,36 +22,35 @@ type ConfigModel struct {
MessagePerHour int `json:"messages_per_hour"`
MessagePerDay int `json:"messages_per_day"`
TempBanDuration string `json:"temp_ban_duration"`
SystemPrompts string `json:"system_prompts"`
SystemPrompts string `json:"system_prompts"` // Consider JSON string or separate table
TelegramToken string `json:"telegram_token"`
Active bool `json:"active"`
}
type Message struct {
gorm.Model
BotID uint `gorm:"index"`
ChatID int64 `gorm:"index"`
UserID int64 `gorm:"index"`
Username string `gorm:"index"`
UserRole string
BotID uint `gorm:"index"`
ChatID int64 `gorm:"index"`
UserID int64 `gorm:"index"`
Username string `gorm:"index"`
UserRole string // Store the role as a string
Text string `gorm:"type:text"`
Timestamp time.Time `gorm:"index"`
IsUser bool
StickerFileID string
StickerPNGFile string
StickerEmoji string
DeletedAt gorm.DeletedAt `gorm:"index"`
AnsweredOn *time.Time `gorm:"index"`
ImageFileIDs []string `gorm:"type:text;serializer:json"`
FilesCleanedAt *time.Time `gorm:"index"`
StickerEmoji string // Store the emoji associated with the sticker
DeletedAt gorm.DeletedAt `gorm:"index"` // Add soft delete field
AnsweredOn *time.Time `gorm:"index"` // Tracks when a user message was answered (NULL for assistant messages and unanswered user messages)
}
type ChatMemory struct {
Messages []Message
Size int
BusinessConnectionID string
BusinessConnectionID string // New field to store the business connection ID
}
// Scope name constants — used in DB seeding, hasScope checks, and tests.
const (
ScopeStatsViewOwn = "stats:view:own"
ScopeStatsViewAny = "stats:view:any"
@@ -61,7 +60,6 @@ const (
ScopeHistoryClearHardAny = "history:clear_hard:any"
ScopeModelSet = "model:set"
ScopeUserPromote = "user:promote"
ScopeTTSUse = "tts:use"
)
type Scope struct {
@@ -77,14 +75,16 @@ type Role struct {
type User struct {
gorm.Model
BotID uint `gorm:"uniqueIndex:idx_user_bot;index"`
TelegramID int64 `gorm:"uniqueIndex:idx_user_bot;not null"`
BotID uint `gorm:"uniqueIndex:idx_user_bot;index"` // Foreign key to BotModel
TelegramID int64 `gorm:"uniqueIndex:idx_user_bot;not null"` // Unique per (telegram_id, bot_id) pair
Username string
RoleID uint
Role Role `gorm:"foreignKey:RoleID"`
IsOwner bool `gorm:"default:false"`
IsOwner bool `gorm:"default:false"` // Indicates if the user is the owner
}
// idx_user_bot is a composite unique index on (bot_id, telegram_id),
// allowing the same Telegram user to be registered independently on each bot.
func (User) TableName() string {
return "users"
}
+7
View File
@@ -33,20 +33,26 @@ func (b *Bot) checkRateLimits(userID int64) bool {
now := limiter.clock.Now()
// Check if the user is currently banned
if now.Before(limiter.banUntil) {
return false
}
// Reset hourly limiter if an hour has passed since the last reset
if now.Sub(limiter.lastHourlyReset) >= time.Hour {
limiter.hourlyLimiter = rate.NewLimiter(rate.Every(time.Hour/time.Duration(b.config.MessagePerHour)), b.config.MessagePerHour)
limiter.lastHourlyReset = now
}
// Reset daily limiter if 24 hours have passed since the last reset
if now.Sub(limiter.lastDailyReset) >= 24*time.Hour {
limiter.dailyLimiter = rate.NewLimiter(rate.Every(24*time.Hour/time.Duration(b.config.MessagePerDay)), b.config.MessagePerDay)
limiter.lastDailyReset = now
}
// Check if the message exceeds rate limits.
// Reserve from both limiters first, then cancel both if either is over budget.
// This prevents consuming a token from one limiter when the other rejects.
dailyRes := limiter.dailyLimiter.ReserveN(now, 1)
hourlyRes := limiter.hourlyLimiter.ReserveN(now, 1)
if dailyRes.DelayFrom(now) > 0 || hourlyRes.DelayFrom(now) > 0 {
@@ -54,6 +60,7 @@ func (b *Bot) checkRateLimits(userID int64) bool {
hourlyRes.CancelAt(now)
banDuration, err := time.ParseDuration(b.config.TempBanDuration)
if err != nil {
// If parsing fails, default to a 24-hour ban
banDuration = 24 * time.Hour
}
limiter.banUntil = now.Add(banDuration)
+23 -5
View File
@@ -5,22 +5,27 @@ import (
"time"
)
// TestCheckRateLimits tests the checkRateLimits method of the Bot.
// It verifies that users are allowed or denied based on their message rates.
func TestCheckRateLimits(t *testing.T) {
// Create a mock clock starting at a fixed time
mockClock := &MockClock{
currentTime: time.Date(2023, 10, 1, 0, 0, 0, 0, time.UTC),
}
// Create a mock configuration with reduced timeframes for testing
config := BotConfig{
ID: "bot1",
MemorySize: 10,
MessagePerHour: 5,
MessagePerDay: 10,
TempBanDuration: "1m",
MessagePerHour: 5, // Allow 5 messages per hour
MessagePerDay: 10, // Allow 10 messages per day
TempBanDuration: "1m", // Temporary ban duration of 1 minute for testing
SystemPrompts: make(map[string]string),
TelegramToken: "YOUR_TELEGRAM_BOT_TOKEN",
OwnerTelegramID: 123456789,
}
// Initialize the Bot with mock data and MockClock
bot := &Bot{
config: config,
userLimiters: make(map[int64]*userLimiter),
@@ -29,39 +34,52 @@ func TestCheckRateLimits(t *testing.T) {
userID := int64(12345)
// Helper function to simulate message sending
sendMessage := func() bool {
return bot.checkRateLimits(userID)
}
// Send 5 messages within the hourly limit
for i := 0; i < config.MessagePerHour; i++ {
if !sendMessage() {
t.Errorf("Expected message %d to be allowed", i+1)
}
}
// 6th message should exceed the hourly limit and trigger a ban
if sendMessage() {
t.Errorf("Expected message to be denied due to hourly limit exceeded")
}
// Attempt to send another message immediately, should still be banned
if sendMessage() {
t.Errorf("Expected message to be denied while user is banned")
}
mockClock.Advance(time.Minute)
// Fast-forward time by TempBanDuration to lift the ban
mockClock.Advance(time.Minute) // Banned for 1 minute
mockClock.Advance(time.Hour)
// Advance time to allow hourly limiter to replenish
mockClock.Advance(time.Hour) // Advance by 1 hour
// Send another message, should be allowed now
if !sendMessage() {
t.Errorf("Expected message to be allowed after ban duration")
}
// Send additional messages to reach the daily limit
for i := 0; i < config.MessagePerDay-config.MessagePerHour-1; i++ {
if !sendMessage() {
t.Errorf("Expected message %d to be allowed towards daily limit", i+1)
}
}
// Attempt to exceed the daily limit
if sendMessage() {
t.Errorf("Expected message to be denied due to daily limit exceeded")
}
}
// To ensure thread safety and avoid race conditions during testing,
// you can run the tests with the `-race` flag:
// go test -race -v
+2 -3
View File
@@ -1,3 +1,4 @@
// telegram_client.go
package main
import (
@@ -7,11 +8,9 @@ import (
"github.com/go-telegram/bot/models"
)
// TelegramClient defines the methods required from the Telegram bot.
type TelegramClient interface {
SendMessage(ctx context.Context, params *bot.SendMessageParams) (*models.Message, error)
SendAudio(ctx context.Context, params *bot.SendAudioParams) (*models.Message, error)
SetMyCommands(ctx context.Context, params *bot.SetMyCommandsParams) (bool, error)
GetFile(ctx context.Context, params *bot.GetFileParams) (*models.File, error)
FileDownloadLink(f *models.File) string
Start(ctx context.Context)
}
+8 -27
View File
@@ -1,3 +1,4 @@
// telegram_client_mock.go
package main
import (
@@ -8,16 +9,15 @@ import (
"github.com/stretchr/testify/mock"
)
// MockTelegramClient is a mock implementation of TelegramClient for testing.
type MockTelegramClient struct {
mock.Mock
SendMessageFunc func(ctx context.Context, params *bot.SendMessageParams) (*models.Message, error)
SendAudioFunc func(ctx context.Context, params *bot.SendAudioParams) (*models.Message, error)
SetMyCommandsFunc func(ctx context.Context, params *bot.SetMyCommandsParams) (bool, error)
GetFileFunc func(ctx context.Context, params *bot.GetFileParams) (*models.File, error)
FileDownloadLinkFunc func(f *models.File) string
StartFunc func(ctx context.Context)
SendMessageFunc func(ctx context.Context, params *bot.SendMessageParams) (*models.Message, error)
SetMyCommandsFunc func(ctx context.Context, params *bot.SetMyCommandsParams) (bool, error)
StartFunc func(ctx context.Context)
}
// SendMessage mocks sending a message.
func (m *MockTelegramClient) SendMessage(ctx context.Context, params *bot.SendMessageParams) (*models.Message, error) {
if m.SendMessageFunc != nil {
return m.SendMessageFunc(ctx, params)
@@ -29,6 +29,7 @@ func (m *MockTelegramClient) SendMessage(ctx context.Context, params *bot.SendMe
return nil, args.Error(1)
}
// SetMyCommands mocks registering bot commands.
func (m *MockTelegramClient) SetMyCommands(ctx context.Context, params *bot.SetMyCommandsParams) (bool, error) {
if m.SetMyCommandsFunc != nil {
return m.SetMyCommandsFunc(ctx, params)
@@ -36,27 +37,7 @@ func (m *MockTelegramClient) SetMyCommands(ctx context.Context, params *bot.SetM
return true, nil
}
func (m *MockTelegramClient) SendAudio(ctx context.Context, params *bot.SendAudioParams) (*models.Message, error) {
if m.SendAudioFunc != nil {
return m.SendAudioFunc(ctx, params)
}
return nil, nil
}
func (m *MockTelegramClient) GetFile(ctx context.Context, params *bot.GetFileParams) (*models.File, error) {
if m.GetFileFunc != nil {
return m.GetFileFunc(ctx, params)
}
return &models.File{}, nil
}
func (m *MockTelegramClient) FileDownloadLink(f *models.File) string {
if m.FileDownloadLinkFunc != nil {
return m.FileDownloadLinkFunc(f)
}
return ""
}
// Start mocks starting the Telegram client.
func (m *MockTelegramClient) Start(ctx context.Context) {
if m.StartFunc != nil {
m.StartFunc(ctx)
-52
View File
@@ -1,52 +0,0 @@
package main
import (
"context"
"fmt"
"io"
"net/http"
tgbot "github.com/go-telegram/bot"
"github.com/go-telegram/bot/models"
)
func largestPhotoSize(photos []models.PhotoSize) models.PhotoSize {
if len(photos) == 0 {
return models.PhotoSize{}
}
largest := photos[0]
largestArea := largest.Width * largest.Height
for i := 1; i < len(photos); i++ {
area := photos[i].Width * photos[i].Height
if area > largestArea {
largest = photos[i]
largestArea = area
}
}
return largest
}
func (b *Bot) downloadTelegramFile(ctx context.Context, fileID string) ([]byte, error) {
fileInfo, err := b.tgBot.GetFile(ctx, &tgbot.GetFileParams{FileID: fileID})
if err != nil {
return nil, fmt.Errorf("telegram GetFile %s: %w", fileID, err)
}
downloadURL := b.tgBot.FileDownloadLink(fileInfo)
req, err := http.NewRequestWithContext(ctx, http.MethodGet, downloadURL, nil)
if err != nil {
return nil, fmt.Errorf("telegram download request %s: %w", fileID, err)
}
resp, err := http.DefaultClient.Do(req)
if err != nil {
return nil, fmt.Errorf("telegram download %s: %w", fileID, err)
}
defer func() { _ = resp.Body.Close() }()
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("telegram download %s: status %d", fileID, resp.StatusCode)
}
data, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("telegram download read %s: %w", fileID, err)
}
return data, nil
}
-53
View File
@@ -1,53 +0,0 @@
package main
import (
"testing"
"github.com/go-telegram/bot/models"
"github.com/stretchr/testify/assert"
)
func TestLargestPhotoSize(t *testing.T) {
cases := []struct {
name string
photos []models.PhotoSize
wantFileID string
}{
{
name: "ascending sizes — last is largest",
photos: []models.PhotoSize{
{FileID: "thumb", Width: 90, Height: 90},
{FileID: "small", Width: 320, Height: 320},
{FileID: "full", Width: 1280, Height: 720},
},
wantFileID: "full",
},
{
name: "descending sizes — first is largest",
photos: []models.PhotoSize{
{FileID: "full", Width: 1280, Height: 720},
{FileID: "small", Width: 320, Height: 320},
{FileID: "thumb", Width: 90, Height: 90},
},
wantFileID: "full",
},
{
name: "single photo",
photos: []models.PhotoSize{
{FileID: "solo", Width: 800, Height: 600},
},
wantFileID: "solo",
},
{
name: "empty slice returns zero value (caller guards upstream)",
photos: []models.PhotoSize{},
wantFileID: "",
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
got := largestPhotoSize(tc.photos)
assert.Equal(t, tc.wantFileID, got.FileID)
})
}
}
+45
View File
@@ -22,18 +22,22 @@ const (
)
func TestOwnerAssignment(t *testing.T) {
// Initialize loggers
initLoggers()
// Initialize in-memory database for testing
db, err := gorm.Open(sqlite.Open(memoryDSN), &gorm.Config{})
if err != nil {
t.Fatalf(errOpenDB, err)
}
// Migrate the schema
err = db.AutoMigrate(&BotModel{}, &ConfigModel{}, &Message{}, &User{}, &Role{}, &Scope{})
if err != nil {
t.Fatalf(errMigrateSchema, err)
}
// Create default roles and scopes
err = createDefaultRoles(db)
if err != nil {
t.Fatalf(errCreateRoles, err)
@@ -42,6 +46,7 @@ func TestOwnerAssignment(t *testing.T) {
t.Fatalf(errCreateScopes, err)
}
// Create a bot configuration
config := BotConfig{
ID: "test_bot",
TelegramToken: "TEST_TELEGRAM_TOKEN",
@@ -54,41 +59,49 @@ func TestOwnerAssignment(t *testing.T) {
OwnerTelegramID: 111111111,
}
// Initialize MockClock
mockClock := &MockClock{
currentTime: time.Now(),
}
// Initialize MockTelegramClient
mockTGClient := &MockTelegramClient{
SendMessageFunc: func(ctx context.Context, params *bot.SendMessageParams) (*models.Message, error) {
chatID, ok := params.ChatID.(int64)
if !ok {
return nil, fmt.Errorf("ChatID is not of type int64")
}
// Simulate successful message sending
return &models.Message{ID: 1, Chat: models.Chat{ID: chatID}}, nil
},
}
// Create the bot with the mock Telegram client
bot, err := NewBot(db, config, mockClock, mockTGClient)
if err != nil {
t.Fatalf(errCreateBot, err)
}
// Verify that the owner exists
var owner User
err = db.Where("telegram_id = ? AND bot_id = ? AND is_owner = ?", config.OwnerTelegramID, bot.botID, true).First(&owner).Error
if err != nil {
t.Fatalf("Owner was not created: %v", err)
}
// Attempt to create another owner for the same bot
_, err = bot.getOrCreateUser(222222222, "AnotherOwner", true)
if err == nil {
t.Fatalf("Expected error when creating a second owner, but got none")
}
// Verify that the error message is appropriate
expectedErrorMsg := "an owner already exists for this bot"
if err.Error() != expectedErrorMsg {
t.Fatalf("Unexpected error message: %v", err)
}
// Assign admin role to a new user
regularUser, err := bot.getOrCreateUser(333333333, "RegularUser", false)
if err != nil {
t.Fatalf("Failed to create regular user: %v", err)
@@ -98,6 +111,7 @@ func TestOwnerAssignment(t *testing.T) {
t.Fatalf("Expected role 'user', got '%s'", regularUser.Role.Name)
}
// Attempt to change an existing user to owner
_, err = bot.getOrCreateUser(333333333, "AdminUser", true)
if err == nil {
t.Fatalf("Expected error when changing existing user to owner, but got none")
@@ -108,21 +122,27 @@ func TestOwnerAssignment(t *testing.T) {
t.Fatalf("Unexpected error message: %v", err)
}
// If you need to test admin creation, you should do it through a separate admin creation function
// or by updating an existing user's role with proper authorization checks
}
func TestPromoteUserToAdmin(t *testing.T) {
// Initialize loggers
initLoggers()
// Initialize in-memory database for testing
db, err := gorm.Open(sqlite.Open(memoryDSN), &gorm.Config{})
if err != nil {
t.Fatalf(errOpenDB, err)
}
// Migrate the schema
err = db.AutoMigrate(&BotModel{}, &ConfigModel{}, &Message{}, &User{}, &Role{}, &Scope{})
if err != nil {
t.Fatalf(errMigrateSchema, err)
}
// Create default roles and scopes
err = createDefaultRoles(db)
if err != nil {
t.Fatalf(errCreateRoles, err)
@@ -151,11 +171,13 @@ func TestPromoteUserToAdmin(t *testing.T) {
t.Fatalf(errCreateBot, err)
}
// Create an owner
owner, err := bot.getOrCreateUser(config.OwnerTelegramID, "OwnerUser", true)
if err != nil {
t.Fatalf("Failed to create owner: %v", err)
}
// Test promoting a user to admin
regularUser, err := bot.getOrCreateUser(444444444, "RegularUser", false)
if err != nil {
t.Fatalf("Failed to create regular user: %v", err)
@@ -166,6 +188,7 @@ func TestPromoteUserToAdmin(t *testing.T) {
t.Fatalf("Failed to promote user to admin: %v", err)
}
// Refresh user data
promotedUser, err := bot.getOrCreateUser(444444444, "RegularUser", false)
if err != nil {
t.Fatalf("Failed to get promoted user: %v", err)
@@ -176,19 +199,26 @@ func TestPromoteUserToAdmin(t *testing.T) {
}
}
// TestGetOrCreateUser tests the getOrCreateUser method of the Bot.
// It verifies that a new user is created when one does not exist,
// and an existing user is returned when one does exist.
func TestGetOrCreateUser(t *testing.T) {
// Initialize loggers
initLoggers()
// Initialize in-memory database for testing
db, err := gorm.Open(sqlite.Open(memoryDSN), &gorm.Config{})
if err != nil {
t.Fatalf(errOpenDB, err)
}
// Migrate the schema
err = db.AutoMigrate(&BotModel{}, &ConfigModel{}, &Message{}, &User{}, &Role{}, &Scope{})
if err != nil {
t.Fatalf(errMigrateSchema, err)
}
// Create default roles and scopes
err = createDefaultRoles(db)
if err != nil {
t.Fatalf(errCreateRoles, err)
@@ -197,10 +227,12 @@ func TestGetOrCreateUser(t *testing.T) {
t.Fatalf(errCreateScopes, err)
}
// Create a mock clock starting at a fixed time
mockClock := &MockClock{
currentTime: time.Date(2023, 10, 1, 0, 0, 0, 0, time.UTC),
}
// Create a mock configuration
config := BotConfig{
ID: "bot1",
MemorySize: 10,
@@ -212,49 +244,62 @@ func TestGetOrCreateUser(t *testing.T) {
OwnerTelegramID: 123456789,
}
// Initialize MockTelegramClient
mockTGClient := &MockTelegramClient{
SendMessageFunc: func(ctx context.Context, params *bot.SendMessageParams) (*models.Message, error) {
chatID, ok := params.ChatID.(int64)
if !ok {
return nil, fmt.Errorf("ChatID is not of type int64")
}
// Simulate successful message sending
return &models.Message{ID: 1, Chat: models.Chat{ID: chatID}}, nil
},
}
// Create the bot with the mock Telegram client
bot, err := NewBot(db, config, mockClock, mockTGClient)
if err != nil {
t.Fatalf(errCreateBot, err)
}
// Verify that the owner exists
var owner User
err = db.Where("telegram_id = ? AND bot_id = ? AND is_owner = ?", config.OwnerTelegramID, bot.botID, true).First(&owner).Error
if err != nil {
t.Fatalf("Owner was not created: %v", err)
}
// Attempt to create another owner for the same bot
_, err = bot.getOrCreateUser(222222222, "AnotherOwner", true)
if err == nil {
t.Fatalf("Expected error when creating a second owner, but got none")
}
// Create a new user
newUser, err := bot.getOrCreateUser(987654321, "TestUser", false)
if err != nil {
t.Fatalf("Failed to create a new user: %v", err)
}
// Verify that the new user was created
var userInDB User
err = db.Where("telegram_id = ?", newUser.TelegramID).First(&userInDB).Error
if err != nil {
t.Fatalf("New user was not created in the database: %v", err)
}
// Get the existing user
existingUser, err := bot.getOrCreateUser(987654321, "TestUser", false)
if err != nil {
t.Fatalf("Failed to get existing user: %v", err)
}
// Verify that the existing user is the same as the new user
if existingUser.ID != userInDB.ID {
t.Fatalf("Expected to get the existing user, but got a different user")
}
}
// To ensure thread safety and avoid race conditions during testing,
// you can run the tests with the `-race` flag:
// go test -race -v