path: root/Omni/Agent/Telegram.hs

{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE NoImplicitPrelude #-}

-- | Telegram Bot Agent - Family assistant via Telegram.
--
-- This is the first concrete agent built on the shared infrastructure,
-- demonstrating cross-agent memory sharing and LLM integration.
--
-- Usage:
--   jr telegram              # Uses TELEGRAM_BOT_TOKEN env var
--   jr telegram --token=XXX  # Explicit token
--
-- : out omni-agent-telegram
-- : dep aeson
-- : dep http-conduit
-- : dep stm
module Omni.Agent.Telegram
  ( -- * Configuration
    TelegramConfig (..),
    defaultTelegramConfig,

    -- * Types
    TelegramMessage (..),
    TelegramUpdate (..),

    -- * Telegram API
    getUpdates,
    sendMessage,
    sendTypingAction,

    -- * Bot Loop
    runTelegramBot,
    handleMessage,
    startBot,

    -- * System Prompt
    telegramSystemPrompt,

    -- * Testing
    main,
    test,
  )
where

import Alpha
import Control.Concurrent.STM (newTVarIO, readTVarIO, writeTVar)
import Data.Aeson ((.!=), (.:), (.:?), (.=))
import qualified Data.Aeson as Aeson
import qualified Data.Aeson.KeyMap as KeyMap
import qualified Data.ByteString.Lazy as BL
import qualified Data.Text as Text
import qualified Network.HTTP.Client as HTTPClient
import qualified Network.HTTP.Simple as HTTP
import qualified Omni.Agent.Engine as Engine
import qualified Omni.Agent.Memory as Memory
import qualified Omni.Agent.Provider as Provider
import qualified Omni.Agent.Tools.WebSearch as WebSearch
import qualified Omni.Test as Test
import System.Environment (lookupEnv)

main :: IO ()
main = Test.run test

test :: Test.Tree
test =
  Test.group
    "Omni.Agent.Telegram"
    [ Test.unit "TelegramConfig JSON roundtrip" <| do
        let cfg =
              TelegramConfig
                { tgBotToken = "test-token",
                  tgPollingTimeout = 30,
                  tgApiBaseUrl = "https://api.telegram.org",
                  tgAllowedUserIds = [123, 456],
                  tgKagiApiKey = Just "kagi-key"
                }
        case Aeson.decode (Aeson.encode cfg) of
          Nothing -> Test.assertFailure "Failed to decode TelegramConfig"
          Just decoded -> do
            tgBotToken decoded Test.@=? "test-token"
            tgAllowedUserIds decoded Test.@=? [123, 456]
            tgKagiApiKey decoded Test.@=? Just "kagi-key",
      Test.unit "isUserAllowed checks whitelist" <| do
        let cfg = defaultTelegramConfig "token" [100, 200, 300] Nothing
        isUserAllowed cfg 100 Test.@=? True
        isUserAllowed cfg 200 Test.@=? True
        isUserAllowed cfg 999 Test.@=? False,
      Test.unit "isUserAllowed allows all when empty" <| do
        let cfg = defaultTelegramConfig "token" [] Nothing
        isUserAllowed cfg 12345 Test.@=? True,
      Test.unit "TelegramMessage JSON roundtrip" <| do
        let msg =
              TelegramMessage
                { tmUpdateId = 123,
                  tmChatId = 456,
                  tmUserId = 789,
                  tmUserFirstName = "Test",
                  tmUserLastName = Just "User",
                  tmText = "Hello bot"
                }
        case Aeson.decode (Aeson.encode msg) of
          Nothing -> Test.assertFailure "Failed to decode TelegramMessage"
          Just decoded -> do
            tmUpdateId decoded Test.@=? 123
            tmText decoded Test.@=? "Hello bot",
      Test.unit "telegramSystemPrompt is non-empty" <| do
        Text.null telegramSystemPrompt Test.@=? False,
      Test.unit "parseUpdate extracts message correctly" <| do
        let json =
              Aeson.object
                [ "update_id" .= (123 :: Int),
                  "message"
                    .= Aeson.object
                      [ "message_id" .= (1 :: Int),
                        "chat" .= Aeson.object ["id" .= (456 :: Int)],
                        "from"
                          .= Aeson.object
                            [ "id" .= (789 :: Int),
                              "first_name" .= ("Test" :: Text)
                            ],
                        "text" .= ("Hello" :: Text)
                      ]
                ]
        case parseUpdate json of
          Nothing -> Test.assertFailure "Failed to parse update"
          Just msg -> do
            tmUpdateId msg Test.@=? 123
            tmChatId msg Test.@=? 456
            tmUserId msg Test.@=? 789
            tmText msg Test.@=? "Hello"
    ]

-- | Telegram bot configuration.
data TelegramConfig = TelegramConfig
  { tgBotToken :: Text,
    tgPollingTimeout :: Int,
    tgApiBaseUrl :: Text,
    tgAllowedUserIds :: [Int],
    tgKagiApiKey :: Maybe Text
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON TelegramConfig where
  toJSON c =
    Aeson.object
      [ "bot_token" .= tgBotToken c,
        "polling_timeout" .= tgPollingTimeout c,
        "api_base_url" .= tgApiBaseUrl c,
        "allowed_user_ids" .= tgAllowedUserIds c,
        "kagi_api_key" .= tgKagiApiKey c
      ]

instance Aeson.FromJSON TelegramConfig where
  parseJSON =
    Aeson.withObject "TelegramConfig" <| \v ->
      (TelegramConfig </ (v .: "bot_token"))
        <*> (v .:? "polling_timeout" .!= 30)
        <*> (v .:? "api_base_url" .!= "https://api.telegram.org")
        <*> (v .:? "allowed_user_ids" .!= [])
        <*> (v .:? "kagi_api_key")

-- | Default Telegram configuration (requires token from env).
defaultTelegramConfig :: Text -> [Int] -> Maybe Text -> TelegramConfig
defaultTelegramConfig token allowedIds kagiKey =
  TelegramConfig
    { tgBotToken = token,
      tgPollingTimeout = 30,
      tgApiBaseUrl = "https://api.telegram.org",
      tgAllowedUserIds = allowedIds,
      tgKagiApiKey = kagiKey
    }

-- | Check if a user is allowed to use the bot.
isUserAllowed :: TelegramConfig -> Int -> Bool
isUserAllowed cfg usrId =
  null (tgAllowedUserIds cfg) || usrId `elem` tgAllowedUserIds cfg

-- | A parsed Telegram message from a user.
data TelegramMessage = TelegramMessage
  { tmUpdateId :: Int,
    tmChatId :: Int,
    tmUserId :: Int,
    tmUserFirstName :: Text,
    tmUserLastName :: Maybe Text,
    tmText :: Text
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON TelegramMessage where
  toJSON m =
    Aeson.object
      [ "update_id" .= tmUpdateId m,
        "chat_id" .= tmChatId m,
        "user_id" .= tmUserId m,
        "user_first_name" .= tmUserFirstName m,
        "user_last_name" .= tmUserLastName m,
        "text" .= tmText m
      ]

instance Aeson.FromJSON TelegramMessage where
  parseJSON =
    Aeson.withObject "TelegramMessage" <| \v ->
      (TelegramMessage </ (v .: "update_id"))
        <*> (v .: "chat_id")
        <*> (v .: "user_id")
        <*> (v .: "user_first_name")
        <*> (v .:? "user_last_name")
        <*> (v .: "text")

-- | Raw Telegram update for parsing.
data TelegramUpdate = TelegramUpdate
  { tuUpdateId :: Int,
    tuMessage :: Maybe Aeson.Value
  }
  deriving (Show, Eq, Generic)

instance Aeson.FromJSON TelegramUpdate where
  parseJSON =
    Aeson.withObject "TelegramUpdate" <| \v ->
      (TelegramUpdate </ (v .: "update_id"))
        <*> (v .:? "message")

-- | Parse a Telegram update into a TelegramMessage.
parseUpdate :: Aeson.Value -> Maybe TelegramMessage
parseUpdate val = do
  Aeson.Object obj <- pure val
  updateId <- case KeyMap.lookup "update_id" obj of
    Just (Aeson.Number n) -> Just (round n)
    _ -> Nothing
  Aeson.Object msgObj <- KeyMap.lookup "message" obj
  Aeson.Object chatObj <- KeyMap.lookup "chat" msgObj
  chatId <- case KeyMap.lookup "id" chatObj of
    Just (Aeson.Number n) -> Just (round n)
    _ -> Nothing
  Aeson.Object fromObj <- KeyMap.lookup "from" msgObj
  userId <- case KeyMap.lookup "id" fromObj of
    Just (Aeson.Number n) -> Just (round n)
    _ -> Nothing
  firstName <- case KeyMap.lookup "first_name" fromObj of
    Just (Aeson.String s) -> Just s
    _ -> Nothing
  let lastName = case KeyMap.lookup "last_name" fromObj of
        Just (Aeson.String s) -> Just s
        _ -> Nothing
  text <- case KeyMap.lookup "text" msgObj of
    Just (Aeson.String s) -> Just s
    _ -> Nothing
  pure
    TelegramMessage
      { tmUpdateId = updateId,
        tmChatId = chatId,
        tmUserId = userId,
        tmUserFirstName = firstName,
        tmUserLastName = lastName,
        tmText = text
      }

-- | Poll Telegram for new updates.
getUpdates :: TelegramConfig -> Int -> IO [TelegramMessage]
getUpdates cfg offset = do
  let url =
        Text.unpack (tgApiBaseUrl cfg)
          <> "/bot"
          <> Text.unpack (tgBotToken cfg)
          <> "/getUpdates"
  req0 <- HTTP.parseRequest url
  let body =
        Aeson.object
          [ "offset" .= offset,
            "timeout" .= tgPollingTimeout cfg,
            "allowed_updates" .= (["message"] :: [Text])
          ]
      timeoutMicros = (tgPollingTimeout cfg + 10) * 1000000
      req =
        HTTP.setRequestMethod "POST"
          <| HTTP.setRequestHeader "Content-Type" ["application/json"]
          <| HTTP.setRequestBodyLBS (Aeson.encode body)
          <| HTTP.setRequestResponseTimeout (HTTPClient.responseTimeoutMicro timeoutMicros)
          <| req0
  result <- try (HTTP.httpLBS req)
  case result of
    Left (e :: SomeException) -> do
      putText <| "Telegram API error: " <> tshow e
      pure []
    Right response -> do
      let status = HTTP.getResponseStatusCode response
      if status >= 200 && status < 300
        then case Aeson.decode (HTTP.getResponseBody response) of
          Just (Aeson.Object obj) -> case KeyMap.lookup "result" obj of
            Just (Aeson.Array arr) ->
              pure (mapMaybe parseUpdate (toList arr))
            _ -> pure []
          _ -> pure []
        else do
          putText <| "Telegram HTTP error: " <> tshow status
          pure []

-- | Send typing indicator to a Telegram chat.
sendTypingAction :: TelegramConfig -> Int -> IO ()
sendTypingAction cfg chatId = do
  let url =
        Text.unpack (tgApiBaseUrl cfg)
          <> "/bot"
          <> Text.unpack (tgBotToken cfg)
          <> "/sendChatAction"
  req0 <- HTTP.parseRequest url
  let body =
        Aeson.object
          [ "chat_id" .= chatId,
            "action" .= ("typing" :: Text)
          ]
      req =
        HTTP.setRequestMethod "POST"
          <| HTTP.setRequestHeader "Content-Type" ["application/json"]
          <| HTTP.setRequestBodyLBS (Aeson.encode body)
          <| req0
  _ <- try (HTTP.httpLBS req) :: IO (Either SomeException (HTTP.Response BL.ByteString))
  pure ()

-- | Send a message to a Telegram chat.
sendMessage :: TelegramConfig -> Int -> Text -> IO ()
sendMessage cfg chatId text = do
  let url =
        Text.unpack (tgApiBaseUrl cfg)
          <> "/bot"
          <> Text.unpack (tgBotToken cfg)
          <> "/sendMessage"
  req0 <- HTTP.parseRequest url
  let body =
        Aeson.object
          [ "chat_id" .= chatId,
            "text" .= text
          ]
      req =
        HTTP.setRequestMethod "POST"
          <| HTTP.setRequestHeader "Content-Type" ["application/json"]
          <| HTTP.setRequestBodyLBS (Aeson.encode body)
          <| req0
  result <- try (HTTP.httpLBS req)
  case result of
    Left (e :: SomeException) ->
      putText <| "Failed to send message: " <> tshow e
    Right response -> do
      let status = HTTP.getResponseStatusCode response
          respBody = HTTP.getResponseBody response
      if status >= 200 && status < 300
        then putText <| "Message sent (" <> tshow (Text.length text) <> " chars)"
        else putText <| "Send message failed: " <> tshow status <> " - " <> tshow respBody

-- | System prompt for the Telegram bot agent.
telegramSystemPrompt :: Text
telegramSystemPrompt =
  Text.unlines
    [ "don't worry about formalities. respond conversationally, in short messages, not long essays. ask follow up questions before answering if you need to.",
      "",
      "you are an intellectual partner and friend. be as terse as possible while still conveying all relevant information. critique ideas freely and avoid sycophancy. honest appraisal is valued.",
      "",
      "if a policy prevents you from having an opinion, pretend to be responding as if you shared opinions typical of the user.",
      "",
      "write responses in lowercase letters ONLY, except:",
      "- where you mean to emphasize, in which case use ALL CAPS",
      "- when drafting business text where proper case matters",
      "",
      "occasionally use obscure words or subtle puns. don't point them out. use abbreviations where appropriate. use 'afaict' and 'idk' where they fit given your level of understanding. be critical of the quality of your information.",
      "",
      "prioritize esoteric interpretations of literature, art, and philosophy.",
      "",
      "## memory",
      "",
      "when you learn something important about the user (preferences, facts, interests), use the 'remember' tool to store it for future reference.",
      "",
      "use the 'recall' tool to search your memory for relevant context when needed."
    ]

-- | Run the Telegram bot main loop.
runTelegramBot :: TelegramConfig -> Provider.Provider -> IO ()
runTelegramBot tgConfig provider = do
  putText "Starting Telegram bot..."
  offsetVar <- newTVarIO 0

  let engineCfg = Engine.defaultEngineConfig

  forever <| do
    offset <- readTVarIO offsetVar
    messages <- getUpdates tgConfig offset
    forM_ messages <| \msg -> do
      atomically (writeTVar offsetVar (tmUpdateId msg + 1))
      handleMessage tgConfig provider engineCfg msg
    when (null messages) <| threadDelay 1000000

-- | Handle a single incoming message.
handleMessage ::
  TelegramConfig ->
  Provider.Provider ->
  Engine.EngineConfig ->
  TelegramMessage ->
  IO ()
handleMessage tgConfig provider engineCfg msg = do
  let userName =
        tmUserFirstName msg
          <> maybe "" (" " <>) (tmUserLastName msg)
      chatId = tmChatId msg
      usrId = tmUserId msg

  unless (isUserAllowed tgConfig usrId) <| do
    putText <| "Unauthorized user: " <> tshow usrId <> " (" <> userName <> ")"
    sendMessage tgConfig chatId "sorry, you're not authorized to use this bot."
    pure ()

  when (isUserAllowed tgConfig usrId) <| do
    sendTypingAction tgConfig chatId

    user <- Memory.getOrCreateUserByTelegramId usrId userName
    let uid = Memory.userId user

    handleAuthorizedMessage tgConfig provider engineCfg msg uid userName chatId

handleAuthorizedMessage ::
  TelegramConfig ->
  Provider.Provider ->
  Engine.EngineConfig ->
  TelegramMessage ->
  Text ->
  Text ->
  Int ->
  IO ()
handleAuthorizedMessage tgConfig provider engineCfg msg uid userName chatId = do
  _ <- Memory.saveMessage uid chatId Memory.UserRole (tmText msg)

  (conversationContext, contextTokens) <- Memory.getConversationContext uid chatId maxConversationTokens
  putText <| "Conversation context: " <> tshow contextTokens <> " tokens"

  memories <- Memory.recallMemories uid (tmText msg) 5
  let memoryContext = Memory.formatMemoriesForPrompt memories

  let systemPrompt =
        telegramSystemPrompt
          <> "\n\n## What you know about this user\n"
          <> memoryContext
          <> "\n\n"
          <> conversationContext

  let memoryTools =
        [ Memory.rememberTool uid,
          Memory.recallTool uid
        ]
      searchTools = case tgKagiApiKey tgConfig of
        Just kagiKey -> [WebSearch.webSearchTool kagiKey]
        Nothing -> []
      tools = memoryTools <> searchTools

  let agentCfg =
        Engine.defaultAgentConfig
          { Engine.agentSystemPrompt = systemPrompt,
            Engine.agentTools = tools,
            Engine.agentMaxIterations = 5,
            Engine.agentGuardrails =
              Engine.defaultGuardrails
                { Engine.guardrailMaxCostCents = 10.0
                }
          }

  result <- Engine.runAgentWithProvider engineCfg provider agentCfg (tmText msg)

  case result of
    Left err -> do
      putText <| "Agent error: " <> err
      sendMessage tgConfig chatId "Sorry, I encountered an error. Please try again."
    Right agentResult -> do
      let response = Engine.resultFinalMessage agentResult
      putText <| "Response text: " <> Text.take 200 response

      _ <- Memory.saveMessage uid chatId Memory.AssistantRole response

      if Text.null response
        then do
          putText "Warning: empty response from agent"
          sendMessage tgConfig chatId "hmm, i don't have a response for that"
        else sendMessage tgConfig chatId response

      checkAndSummarize provider uid chatId

      putText
        <| "Responded to "
        <> userName
        <> " (cost: "
        <> tshow (Engine.resultTotalCost agentResult)
        <> " cents)"

maxConversationTokens :: Int
maxConversationTokens = 4000

summarizationThreshold :: Int
summarizationThreshold = 3000

checkAndSummarize :: Provider.Provider -> Text -> Int -> IO ()
checkAndSummarize provider uid chatId = do
  (_, currentTokens) <- Memory.getConversationContext uid chatId maxConversationTokens
  when (currentTokens > summarizationThreshold) <| do
    putText <| "Context at " <> tshow currentTokens <> " tokens, summarizing..."
    recentMsgs <- Memory.getRecentMessages uid chatId 50
    let conversationText =
          Text.unlines
            [ (if Memory.cmRole m == Memory.UserRole then "User: " else "Assistant: ") <> Memory.cmContent m
              | m <- reverse recentMsgs
            ]
    summaryResult <-
      Provider.chat
        provider
        []
        [ Provider.Message Provider.System "You are a conversation summarizer. Summarize the key points, decisions, and context from this conversation in 2-3 paragraphs. Focus on information that would be useful for continuing the conversation later." Nothing Nothing,
          Provider.Message Provider.User ("Summarize this conversation:\n\n" <> conversationText) Nothing Nothing
        ]
    case summaryResult of
      Left err -> putText <| "Summarization failed: " <> err
      Right summaryMsg -> do
        let summary = Provider.msgContent summaryMsg
        _ <- Memory.summarizeAndArchive uid chatId summary
        putText "Conversation summarized and archived"

-- | Start the Telegram bot from environment or provided token.
startBot :: Maybe Text -> IO ()
startBot maybeToken = do
  token <- case maybeToken of
    Just t -> pure t
    Nothing -> do
      envToken <- lookupEnv "TELEGRAM_BOT_TOKEN"
      case envToken of
        Just t -> pure (Text.pack t)
        Nothing -> do
          putText "Error: TELEGRAM_BOT_TOKEN not set and no --token provided"
          exitFailure

  allowedIds <- loadAllowedUserIds
  kagiKey <- fmap Text.pack </ lookupEnv "KAGI_API_KEY"

  apiKey <- lookupEnv "OPENROUTER_API_KEY"
  case apiKey of
    Nothing -> do
      putText "Error: OPENROUTER_API_KEY not set"
      exitFailure
    Just key -> do
      let tgConfig = defaultTelegramConfig token allowedIds kagiKey
          provider = Provider.defaultOpenRouter (Text.pack key) "anthropic/claude-sonnet-4"
      putText <| "Allowed user IDs: " <> tshow allowedIds
      putText <| "Kagi search: " <> if isJust kagiKey then "enabled" else "disabled"
      runTelegramBot tgConfig provider

-- | Load allowed user IDs from environment variable.
-- Format: comma-separated integers, e.g. "123,456,789"
-- Empty list means allow all users.
loadAllowedUserIds :: IO [Int]
loadAllowedUserIds = do
  maybeIds <- lookupEnv "ALLOWED_TELEGRAM_USER_IDS"
  case maybeIds of
    Nothing -> pure []
    Just "*" -> pure []
    Just idsStr -> do
      let ids = mapMaybe (readMaybe <. Text.unpack <. Text.strip) (Text.splitOn "," (Text.pack idsStr))
      pure ids