mirror of
https://github.com/HugeFrog24/go-telegram-bot.git
synced 2026-08-28 22:11:38 +00:00
Compare commits
7
Commits
0543283b8a
..
v1.0.0
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
44fcd02d9a | ||
|
|
d8d0da4704 | ||
|
|
c8af457af1 | ||
|
|
e5532df7f9 | ||
|
|
0ab56448c7 | ||
|
|
9f2b3df4c8 | ||
|
|
41c9b8075b |
@@ -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.
|
||||
@@ -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
@@ -15,4 +15,4 @@ bot.db
|
||||
|
||||
# All config files except for the default
|
||||
config/*
|
||||
!config/default.json
|
||||
!config/default.json
|
||||
@@ -0,0 +1,3 @@
|
||||
{
|
||||
"mcpServers": {}
|
||||
}
|
||||
@@ -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.
|
||||
|
||||
@@ -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
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
@@ -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")
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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
@@ -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
@@ -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
@@ -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 {
|
||||
|
||||
@@ -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.
@@ -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
|
||||
)
|
||||
|
||||
@@ -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
@@ -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
@@ -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(®ularUser).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(®ularUser).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{
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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.")
|
||||
|
||||
@@ -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"
|
||||
}
|
||||
|
||||
@@ -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
@@ -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
@@ -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
@@ -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)
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user