1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
|
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE NoImplicitPrelude #-}
-- | Telegram Incoming Message Queue - Batches incoming messages by chat.
--
-- Messages are queued in-memory and batched by chat_id with a configurable
-- window (default 1s). This prevents confusion when messages arrive
-- simultaneously from different chats.
--
-- : out omni-agent-telegram-incoming-queue
-- : dep stm
module Omni.Agent.Telegram.IncomingQueue
( -- * Types
IncomingQueues,
ChatQueue (..),
QueuedMsg (..),
-- * Queue Operations
newIncomingQueues,
enqueueIncoming,
-- * Batch Processing
flushReadyBatches,
startIncomingBatcher,
-- * Batch Formatting
formatBatch,
-- * Configuration
defaultBatchWindowSeconds,
-- * Testing
main,
test,
)
where
import Alpha
import Control.Concurrent.STM (TVar, newTVarIO, readTVar, readTVarIO, writeTVar)
import qualified Data.Map.Strict as Map
import qualified Data.Text as Text
import Data.Time (NominalDiffTime, UTCTime, addUTCTime, getCurrentTime)
import qualified Omni.Agent.Telegram.Types as Types
import qualified Omni.Test as Test
main :: IO ()
main = Test.run test
test :: Test.Tree
test =
Test.group
"Omni.Agent.Telegram.IncomingQueue"
[ Test.unit "newIncomingQueues creates empty map" <| do
queues <- newIncomingQueues
qs <- readTVarIO queues
Map.null qs Test.@=? True,
Test.unit "formatBatch single message no attribution in DM" <| do
now <- getCurrentTime
let msg = mkTestMessage 123 456 Types.Private "hello"
qmsg = QueuedMsg now msg
result = formatBatch [qmsg]
result Test.@=? "hello",
Test.unit "formatBatch multiple messages numbered" <| do
now <- getCurrentTime
let msg1 = mkTestMessage 123 456 Types.Private "first"
msg2 = mkTestMessage 123 456 Types.Private "second"
qmsgs = [QueuedMsg now msg1, QueuedMsg now msg2]
result = formatBatch qmsgs
("1. first" `Text.isInfixOf` result) Test.@=? True
("2. second" `Text.isInfixOf` result) Test.@=? True,
Test.unit "formatBatch group chat has sender attribution" <| do
now <- getCurrentTime
let msg = mkTestMessage 123 456 Types.Group "hello"
qmsg = QueuedMsg now msg
result = formatBatch [qmsg]
("[Test] hello" `Text.isInfixOf` result) Test.@=? True,
Test.unit "enqueueIncoming adds to queue" <| do
queues <- newIncomingQueues
let msg = mkTestMessage 123 456 Types.Private "test"
enqueueIncoming queues 1.0 msg
qs <- readTVarIO queues
Map.member 123 qs Test.@=? True,
Test.unit "flushReadyBatches returns due batches" <| do
queues <- newIncomingQueues
t <- getCurrentTime
let msg = mkTestMessage 123 456 Types.Private "test"
atomically <| do
let qmsg = QueuedMsg t msg
queue = ChatQueue [qmsg] t
writeTVar queues (Map.singleton 123 queue)
threadDelay 10000
batches <- flushReadyBatches queues
length batches Test.@=? 1
]
mkTestMessage :: Int -> Int -> Types.ChatType -> Text -> Types.TelegramMessage
mkTestMessage chatId usrId chatType txt =
Types.TelegramMessage
{ Types.tmUpdateId = 1,
Types.tmChatId = chatId,
Types.tmChatType = chatType,
Types.tmUserId = usrId,
Types.tmUserFirstName = "Test",
Types.tmUserLastName = Nothing,
Types.tmText = txt,
Types.tmDocument = Nothing,
Types.tmPhoto = Nothing,
Types.tmVoice = Nothing,
Types.tmReplyTo = Nothing
}
data QueuedMsg = QueuedMsg
{ qmReceivedAt :: UTCTime,
qmMsg :: Types.TelegramMessage
}
deriving (Show, Eq)
data ChatQueue = ChatQueue
{ cqMessages :: [QueuedMsg],
cqDeadline :: UTCTime
}
deriving (Show, Eq)
type ChatId = Int
type IncomingQueues = TVar (Map.Map ChatId ChatQueue)
defaultBatchWindowSeconds :: NominalDiffTime
defaultBatchWindowSeconds = 3.0
newIncomingQueues :: IO IncomingQueues
newIncomingQueues = newTVarIO Map.empty
enqueueIncoming :: IncomingQueues -> NominalDiffTime -> Types.TelegramMessage -> IO ()
enqueueIncoming queuesVar windowSeconds msg = do
now <- getCurrentTime
let chatId = Types.tmChatId msg
newDeadline = addUTCTime windowSeconds now
qMsg = QueuedMsg now msg
atomically <| do
qs <- readTVar queuesVar
let qs' = Map.alter (insertOrUpdate newDeadline qMsg) chatId qs
writeTVar queuesVar qs'
where
insertOrUpdate deadline qMsg Nothing =
Just ChatQueue {cqMessages = [qMsg], cqDeadline = deadline}
insertOrUpdate deadline qMsg (Just q) =
Just
q
{ cqMessages = cqMessages q <> [qMsg],
cqDeadline = deadline
}
flushReadyBatches :: IncomingQueues -> IO [(ChatId, [QueuedMsg])]
flushReadyBatches queuesVar = do
now <- getCurrentTime
atomically <| do
qs <- readTVar queuesVar
let (ready, pending) = Map.partition (\q -> cqDeadline q <= now) qs
batches =
[ (chatId, cqMessages q)
| (chatId, q) <- Map.toList ready
]
writeTVar queuesVar pending
pure batches
startIncomingBatcher ::
IncomingQueues ->
(Types.TelegramMessage -> Text -> IO ()) ->
IO ()
startIncomingBatcher queuesVar processFn =
void <| forkIO <| forever <| do
batches <- flushReadyBatches queuesVar
forM_ batches <| \(_chatId, qmsgs) -> do
case qmsgs of
[] -> pure ()
(firstQm : _) -> do
let baseMsg = qmMsg firstQm
batchedTxt = formatBatch qmsgs
processFn baseMsg batchedTxt
threadDelay 200000
formatBatch :: [QueuedMsg] -> Text
formatBatch [] = ""
formatBatch [single] = formatOne False 1 single
formatBatch qmsgs = Text.intercalate "\n\n" (zipWith (formatOne True) [1 ..] qmsgs)
formatOne :: Bool -> Int -> QueuedMsg -> Text
formatOne numbered idx (QueuedMsg _ msg) =
let baseText = Types.tmText msg
sender = senderLabel msg
media = mediaSuffix msg
reply = replySuffix msg
prefix =
if numbered
then tshow idx <> ". "
else ""
in Text.concat [prefix, sender, baseText, reply, media]
senderLabel :: Types.TelegramMessage -> Text
senderLabel msg
| Types.isGroupChat msg =
let firstName = Types.tmUserFirstName msg
lastName = fromMaybe "" (Types.tmUserLastName msg)
name = Text.strip (firstName <> " " <> lastName)
in "[" <> name <> "] "
| otherwise = ""
mediaSuffix :: Types.TelegramMessage -> Text
mediaSuffix msg =
Text.concat
<| [ " [document: " <> fromMaybe "unnamed" (Types.tdFileName d) <> "]"
| Just d <- [Types.tmDocument msg]
]
<> [" [photo attached]" | isJust (Types.tmPhoto msg)]
<> [" [voice message]" | isJust (Types.tmVoice msg)]
replySuffix :: Types.TelegramMessage -> Text
replySuffix msg =
case Types.tmReplyTo msg of
Nothing -> ""
Just r ->
let fn = fromMaybe "someone" (Types.trFromFirstName r)
ln = fromMaybe "" (Types.trFromLastName r)
name = Text.strip (fn <> " " <> ln)
snippet = Text.take 80 (Types.trText r)
in " (replying to " <> name <> ": \"" <> snippet <> "\")"
|
