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

1 files changed, 212 insertions, 0 deletions

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)