diff options
Diffstat (limited to 'Omni/Agent/Telegram/Messages.hs')
| -rw-r--r-- | Omni/Agent/Telegram/Messages.hs | 54 |
1 files changed, 35 insertions, 19 deletions
diff --git a/Omni/Agent/Telegram/Messages.hs b/Omni/Agent/Telegram/Messages.hs index dfa3a3d..eab9668 100644 --- a/Omni/Agent/Telegram/Messages.hs +++ b/Omni/Agent/Telegram/Messages.hs @@ -128,6 +128,7 @@ data ScheduledMessage = ScheduledMessage { smId :: Text, smUserId :: Maybe Text, smChatId :: Int, + smThreadId :: Maybe Int, smContent :: Text, smSendAt :: UTCTime, smCreatedAt :: UTCTime, @@ -147,6 +148,7 @@ instance Aeson.ToJSON ScheduledMessage where [ "id" .= smId m, "user_id" .= smUserId m, "chat_id" .= smChatId m, + "thread_id" .= smThreadId m, "content" .= smContent m, "send_at" .= smSendAt m, "created_at" .= smCreatedAt m, @@ -164,6 +166,7 @@ instance SQL.FromRow ScheduledMessage where id' <- SQL.field userId <- SQL.field chatId <- SQL.field + threadId <- SQL.field content <- SQL.field sendAt <- SQL.field createdAt <- SQL.field @@ -180,6 +183,7 @@ instance SQL.FromRow ScheduledMessage where { smId = id', smUserId = userId, smChatId = chatId, + smThreadId = threadId, smContent = content, smSendAt = sendAt, smCreatedAt = createdAt, @@ -199,13 +203,14 @@ maxRetries :: Int maxRetries = 5 initScheduledMessagesTable :: SQL.Connection -> IO () -initScheduledMessagesTable conn = +initScheduledMessagesTable conn = do SQL.execute_ conn "CREATE TABLE IF NOT EXISTS scheduled_messages (\ \ id TEXT PRIMARY KEY,\ \ user_id TEXT,\ \ chat_id INTEGER NOT NULL,\ + \ thread_id INTEGER,\ \ content TEXT NOT NULL,\ \ send_at TIMESTAMP NOT NULL,\ \ created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,\ @@ -217,16 +222,25 @@ initScheduledMessagesTable conn = \ correlation_id TEXT,\ \ telegram_message_id INTEGER\ \)" + migrateAddThreadId conn + +migrateAddThreadId :: SQL.Connection -> IO () +migrateAddThreadId conn = do + result <- try @SomeException <| SQL.execute_ conn "ALTER TABLE scheduled_messages ADD COLUMN thread_id INTEGER" + case result of + Left _ -> pure () + Right () -> pure () queueMessage :: Maybe Text -> Int -> + Maybe Int -> Text -> UTCTime -> Maybe Text -> Maybe Text -> IO Text -queueMessage mUserId chatId content sendAt msgType correlationId = do +queueMessage mUserId chatId mThreadId content sendAt msgType correlationId = do uuid <- UUID.nextRandom now <- getCurrentTime let msgId = UUID.toText uuid @@ -235,34 +249,36 @@ queueMessage mUserId chatId content sendAt msgType correlationId = do SQL.execute conn "INSERT INTO scheduled_messages \ - \(id, user_id, chat_id, content, send_at, created_at, status, retry_count, message_type, correlation_id) \ - \VALUES (?, ?, ?, ?, ?, ?, 'pending', 0, ?, ?)" - (msgId, mUserId, chatId, content, sendAt, now, msgType, correlationId) + \(id, user_id, chat_id, thread_id, content, send_at, created_at, status, retry_count, message_type, correlation_id) \ + \VALUES (?, ?, ?, ?, ?, ?, ?, 'pending', 0, ?, ?)" + (msgId, mUserId, chatId, mThreadId, content, sendAt, now, msgType, correlationId) pure msgId enqueueImmediate :: Maybe Text -> Int -> + Maybe Int -> Text -> Maybe Text -> Maybe Text -> IO Text -enqueueImmediate mUserId chatId content msgType correlationId = do +enqueueImmediate mUserId chatId mThreadId content msgType correlationId = do now <- getCurrentTime - queueMessage mUserId chatId content now msgType correlationId + queueMessage mUserId chatId mThreadId content now msgType correlationId enqueueDelayed :: Maybe Text -> Int -> + Maybe Int -> Text -> NominalDiffTime -> Maybe Text -> Maybe Text -> IO Text -enqueueDelayed mUserId chatId content delay msgType correlationId = do +enqueueDelayed mUserId chatId mThreadId content delay msgType correlationId = do now <- getCurrentTime let sendAt = addUTCTime delay now - queueMessage mUserId chatId content sendAt msgType correlationId + queueMessage mUserId chatId mThreadId content sendAt msgType correlationId fetchDueMessages :: UTCTime -> Int -> IO [ScheduledMessage] fetchDueMessages now batchSize = @@ -270,7 +286,7 @@ fetchDueMessages now batchSize = initScheduledMessagesTable conn SQL.query conn - "SELECT id, user_id, chat_id, content, send_at, created_at, status, \ + "SELECT id, user_id, chat_id, thread_id, content, send_at, created_at, status, \ \retry_count, last_attempt_at, last_error, message_type, correlation_id, telegram_message_id \ \FROM scheduled_messages \ \WHERE status = 'pending' AND send_at <= ? \ @@ -286,7 +302,7 @@ listPendingMessages mUserId chatId = Just uid -> SQL.query conn - "SELECT id, user_id, chat_id, content, send_at, created_at, status, \ + "SELECT id, user_id, chat_id, thread_id, content, send_at, created_at, status, \ \retry_count, last_attempt_at, last_error, message_type, correlation_id, telegram_message_id \ \FROM scheduled_messages \ \WHERE user_id = ? AND chat_id = ? AND status = 'pending' AND send_at > datetime('now') \ @@ -295,7 +311,7 @@ listPendingMessages mUserId chatId = Nothing -> SQL.query conn - "SELECT id, user_id, chat_id, content, send_at, created_at, status, \ + "SELECT id, user_id, chat_id, thread_id, content, send_at, created_at, status, \ \retry_count, last_attempt_at, last_error, message_type, correlation_id, telegram_message_id \ \FROM scheduled_messages \ \WHERE chat_id = ? AND status = 'pending' AND send_at > datetime('now') \ @@ -309,7 +325,7 @@ getMessageById msgId = results <- SQL.query conn - "SELECT id, user_id, chat_id, content, send_at, created_at, status, \ + "SELECT id, user_id, chat_id, thread_id, content, send_at, created_at, status, \ \retry_count, last_attempt_at, last_error, message_type, correlation_id, telegram_message_id \ \FROM scheduled_messages \ \WHERE id = ?" @@ -380,7 +396,7 @@ cancelMessage msgId = changes <- SQL.changes conn pure (changes > 0) -messageDispatchLoop :: (Int -> Text -> IO (Maybe Int)) -> IO () +messageDispatchLoop :: (Int -> Maybe Int -> Text -> IO (Maybe Int)) -> IO () messageDispatchLoop sendFn = forever <| do now <- getCurrentTime @@ -391,11 +407,11 @@ messageDispatchLoop sendFn = forM_ due <| \m -> dispatchOne sendFn m when (length due < 10) <| threadDelay 1000000 -dispatchOne :: (Int -> Text -> IO (Maybe Int)) -> ScheduledMessage -> IO () +dispatchOne :: (Int -> Maybe Int -> Text -> IO (Maybe Int)) -> ScheduledMessage -> IO () dispatchOne sendFn m = do now <- getCurrentTime markSending (smId m) now - result <- try (sendFn (smChatId m) (smContent m)) + result <- try (sendFn (smChatId m) (smThreadId m) (smContent m)) case result of Left (e :: SomeException) -> do let err = "Exception sending Telegram message: " <> tshow e @@ -409,8 +425,8 @@ dispatchOne sendFn m = do markSent (smId m) (Just telegramMsgId) now' putText <| "Sent message " <> smId m <> " -> telegram_id " <> tshow telegramMsgId -sendMessageTool :: Text -> Int -> Engine.Tool -sendMessageTool uid chatId = +sendMessageTool :: Text -> Int -> Maybe Int -> Engine.Tool +sendMessageTool uid chatId mThreadId = Engine.Tool { Engine.toolName = "send_message", Engine.toolDescription = @@ -453,7 +469,7 @@ sendMessageTool uid chatId = let delay = fromIntegral (fromMaybe 0 delaySeconds) now <- getCurrentTime let sendAt = addUTCTime delay now - msgId <- queueMessage (Just uid) chatId text sendAt (Just "agent_tool") Nothing + msgId <- queueMessage (Just uid) chatId mThreadId text sendAt (Just "agent_tool") Nothing pure <| Aeson.object [ "status" .= ("queued" :: Text), |