Telegram bot: Kagi web search tool

- Add Omni/Agent/Tools/WebSearch.hs with Kagi Search API integration
- webSearchTool for agents to search the web
- kagiSearch function for direct API access
- Load KAGI_API_KEY from environment
- Wire web search into Telegram bot tools
- Results formatted with title, URL, and snippet

Closes t-252

2 files changed, 236 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.
diff --git a/Omni/Agent/Tools/WebSearch.hs b/Omni/Agent/Tools/WebSearch.hs
new file mode 100644
index 0000000..f7250b8
--- /dev/null
+++ b/Omni/Agent/Tools/WebSearch.hs
@@ -0,0 +1,212 @@
+{-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE NoImplicitPrelude #-}
+
+-- | Web search tool using Kagi Search API.
+--
+-- Provides web search capabilities for agents.
+--
+-- : out omni-agent-tools-websearch
+-- : dep aeson
+-- : dep http-conduit
+module Omni.Agent.Tools.WebSearch
+  ( -- * Tool
+    webSearchTool,
+
+    -- * Direct API
+    kagiSearch,
+    SearchResult (..),
+
+    -- * Testing
+    main,
+    test,
+  )
+where
+
+import Alpha
+import Data.Aeson ((.=))
+import qualified Data.Aeson as Aeson
+import qualified Data.Aeson.KeyMap as KeyMap
+import qualified Data.Text as Text
+import qualified Data.Text.Encoding as TE
+import qualified Network.HTTP.Simple as HTTP
+import qualified Network.HTTP.Types.URI as URI
+import qualified Omni.Agent.Engine as Engine
+import qualified Omni.Test as Test
+
+main :: IO ()
+main = Test.run test
+
+test :: Test.Tree
+test =
+  Test.group
+    "Omni.Agent.Tools.WebSearch"
+    [ Test.unit "SearchResult JSON parsing" <| do
+        let json =
+              Aeson.object
+                [ "t" .= (0 :: Int),
+                  "url" .= ("https://example.com" :: Text),
+                  "title" .= ("Example Title" :: Text),
+                  "snippet" .= ("This is a snippet" :: Text)
+                ]
+        case parseSearchResult json of
+          Nothing -> Test.assertFailure "Failed to parse search result"
+          Just sr -> do
+            srUrl sr Test.@=? "https://example.com"
+            srTitle sr Test.@=? "Example Title"
+            srSnippet sr Test.@=? Just "This is a snippet",
+      Test.unit "webSearchTool has correct schema" <| do
+        let tool = webSearchTool "test-key"
+        Engine.toolName tool Test.@=? "web_search",
+      Test.unit "formatResultsForLLM formats correctly" <| do
+        let results =
+              [ SearchResult "https://a.com" "Title A" (Just "Snippet A") Nothing,
+                SearchResult "https://b.com" "Title B" Nothing Nothing
+              ]
+            formatted = formatResultsForLLM results
+        ("Title A" `Text.isInfixOf` formatted) Test.@=? True
+        ("https://a.com" `Text.isInfixOf` formatted) Test.@=? True
+    ]
+
+data SearchResult = SearchResult
+  { srUrl :: Text,
+    srTitle :: Text,
+    srSnippet :: Maybe Text,
+    srPublished :: Maybe Text
+  }
+  deriving (Show, Eq, Generic)
+
+instance Aeson.ToJSON SearchResult where
+  toJSON r =
+    Aeson.object
+      [ "url" .= srUrl r,
+        "title" .= srTitle r,
+        "snippet" .= srSnippet r,
+        "published" .= srPublished r
+      ]
+
+parseSearchResult :: Aeson.Value -> Maybe SearchResult
+parseSearchResult val = do
+  Aeson.Object obj <- pure val
+  t <- case KeyMap.lookup "t" obj of
+    Just (Aeson.Number n) -> Just (round n :: Int)
+    _ -> Nothing
+  guard (t == 0)
+  url <- case KeyMap.lookup "url" obj of
+    Just (Aeson.String s) -> Just s
+    _ -> Nothing
+  title <- case KeyMap.lookup "title" obj of
+    Just (Aeson.String s) -> Just s
+    _ -> Nothing
+  let snippet = case KeyMap.lookup "snippet" obj of
+        Just (Aeson.String s) -> Just s
+        _ -> Nothing
+      published = case KeyMap.lookup "published" obj of
+        Just (Aeson.String s) -> Just s
+        _ -> Nothing
+  pure SearchResult {srUrl = url, srTitle = title, srSnippet = snippet, srPublished = published}
+
+kagiSearch :: Text -> Text -> Int -> IO (Either Text [SearchResult])
+kagiSearch apiKey query limit = do
+  let encodedQuery = TE.decodeUtf8 (URI.urlEncode False (TE.encodeUtf8 query))
+      url = "https://kagi.com/api/v0/search?q=" <> Text.unpack encodedQuery <> "&limit=" <> show limit
+  result <-
+    try <| do
+      req0 <- HTTP.parseRequest url
+      let req =
+            HTTP.setRequestMethod "GET"
+              <| HTTP.setRequestHeader "Authorization" ["Bot " <> TE.encodeUtf8 apiKey]
+              <| req0
+      HTTP.httpLBS req
+  case result of
+    Left (e :: SomeException) ->
+      pure (Left ("Kagi API error: " <> tshow e))
+    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 "data" obj of
+            Just (Aeson.Array arr) ->
+              pure (Right (mapMaybe parseSearchResult (toList arr)))
+            _ -> pure (Left "No data in response")
+          _ -> pure (Left "Failed to parse Kagi response")
+        else case Aeson.decode (HTTP.getResponseBody response) of
+          Just (Aeson.Object obj) -> case KeyMap.lookup "error" obj of
+            Just errArr -> pure (Left ("Kagi error: " <> tshow errArr))
+            _ -> pure (Left ("Kagi HTTP error: " <> tshow status))
+          _ -> pure (Left ("Kagi HTTP error: " <> tshow status))
+
+formatResultsForLLM :: [SearchResult] -> Text
+formatResultsForLLM [] = "No results found."
+formatResultsForLLM results =
+  Text.unlines (zipWith formatResult [1 ..] results)
+  where
+    formatResult :: Int -> SearchResult -> Text
+    formatResult n r =
+      tshow n
+        <> ". "
+        <> srTitle r
+        <> "\n   "
+        <> srUrl r
+        <> maybe "" (\s -> "\n   " <> Text.take 200 s) (srSnippet r)
+
+webSearchTool :: Text -> Engine.Tool
+webSearchTool apiKey =
+  Engine.Tool
+    { Engine.toolName = "web_search",
+      Engine.toolDescription =
+        "Search the web using Kagi. Use this to find current information, "
+          <> "verify facts, look up documentation, or research topics. "
+          <> "Returns titles, URLs, and snippets from search results.",
+      Engine.toolJsonSchema =
+        Aeson.object
+          [ "type" .= ("object" :: Text),
+            "properties"
+              .= Aeson.object
+                [ "query"
+                    .= Aeson.object
+                      [ "type" .= ("string" :: Text),
+                        "description" .= ("The search query" :: Text)
+                      ],
+                  "limit"
+                    .= Aeson.object
+                      [ "type" .= ("integer" :: Text),
+                        "description" .= ("Max results to return (default: 5, max: 10)" :: Text)
+                      ]
+                ],
+            "required" .= (["query"] :: [Text])
+          ],
+      Engine.toolExecute = executeWebSearch apiKey
+    }
+
+executeWebSearch :: Text -> Aeson.Value -> IO Aeson.Value
+executeWebSearch apiKey v =
+  case Aeson.fromJSON v of
+    Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e])
+    Aeson.Success (args :: WebSearchArgs) -> do
+      let lim = min 10 (max 1 (wsLimit args))
+      result <- kagiSearch apiKey (wsQuery args) lim
+      case result of
+        Left err ->
+          pure (Aeson.object ["error" .= err])
+        Right results ->
+          pure
+            ( Aeson.object
+                [ "success" .= True,
+                  "count" .= length results,
+                  "results" .= formatResultsForLLM results
+                ]
+            )
+
+data WebSearchArgs = WebSearchArgs
+  { wsQuery :: Text,
+    wsLimit :: Int
+  }
+  deriving (Generic)
+
+instance Aeson.FromJSON WebSearchArgs where
+  parseJSON =
+    Aeson.withObject "WebSearchArgs" <| \v ->
+      (WebSearchArgs </ (v Aeson..: "query"))
+        <*> (v Aeson..:? "limit" Aeson..!= 5)