1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
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: 10, max: 20)" :: 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 20 (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..!= 10)
|
