1 files changed, 24 insertions, 11 deletions

diff --git a/Omni/Agent/Telegram.hs b/Omni/Agent/Telegram.hs
index 566377e..1162e25 100644
--- a/Omni/Agent/Telegram.hs
+++ b/Omni/Agent/Telegram.hs
@@ -56,6 +56,7 @@ 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)
 
@@ -72,20 +73,22 @@ test =
                 { tgBotToken = "test-token",
                   tgPollingTimeout = 30,
                   tgApiBaseUrl = "https://api.telegram.org",
-                  tgAllowedUserIds = [123, 456]
+                  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],
+            tgAllowedUserIds decoded Test.@=? [123, 456]
+            tgKagiApiKey decoded Test.@=? Just "kagi-key",
       Test.unit "isUserAllowed checks whitelist" <| do
-        let cfg = defaultTelegramConfig "token" [100, 200, 300]
+        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" []
+        let cfg = defaultTelegramConfig "token" [] Nothing
         isUserAllowed cfg 12345 Test.@=? True,
       Test.unit "TelegramMessage JSON roundtrip" <| do
         let msg =
@@ -134,7 +137,8 @@ data TelegramConfig = TelegramConfig
   { tgBotToken :: Text,
     tgPollingTimeout :: Int,
     tgApiBaseUrl :: Text,
-    tgAllowedUserIds :: [Int]
+    tgAllowedUserIds :: [Int],
+    tgKagiApiKey :: Maybe Text
   }
   deriving (Show, Eq, Generic)
 
@@ -144,7 +148,8 @@ instance Aeson.ToJSON TelegramConfig where
       [ "bot_token" .= tgBotToken c,
         "polling_timeout" .= tgPollingTimeout c,
         "api_base_url" .= tgApiBaseUrl c,
-        "allowed_user_ids" .= tgAllowedUserIds c
+        "allowed_user_ids" .= tgAllowedUserIds c,
+        "kagi_api_key" .= tgKagiApiKey c
       ]
 
 instance Aeson.FromJSON TelegramConfig where
@@ -154,15 +159,17 @@ instance Aeson.FromJSON TelegramConfig where
         <*> (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] -> TelegramConfig
-defaultTelegramConfig token allowedIds =
+defaultTelegramConfig :: Text -> [Int] -> Maybe Text -> TelegramConfig
+defaultTelegramConfig token allowedIds kagiKey =
   TelegramConfig
     { tgBotToken = token,
       tgPollingTimeout = 30,
       tgApiBaseUrl = "https://api.telegram.org",
-      tgAllowedUserIds = allowedIds
+      tgAllowedUserIds = allowedIds,
+      tgKagiApiKey = kagiKey
     }
 
 -- | Check if a user is allowed to use the bot.
@@ -435,10 +442,14 @@ handleAuthorizedMessage tgConfig provider engineCfg msg uid userName chatId = do
           <> "\n\n"
           <> conversationContext
 
-  let tools =
+  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
@@ -523,6 +534,7 @@ startBot maybeToken = do
           exitFailure
 
   allowedIds <- loadAllowedUserIds
+  kagiKey <- fmap Text.pack </ lookupEnv "KAGI_API_KEY"
 
   apiKey <- lookupEnv "OPENROUTER_API_KEY"
   case apiKey of
@@ -530,9 +542,10 @@ startBot maybeToken = do
       putText "Error: OPENROUTER_API_KEY not set"
       exitFailure
     Just key -> do
-      let tgConfig = defaultTelegramConfig token allowedIds
+      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.