diff options
Diffstat (limited to 'Omni/Agent/Tools')
| -rw-r--r-- | Omni/Agent/Tools/Calendar.hs | 322 | ||||
| -rw-r--r-- | Omni/Agent/Tools/Email.hs | 675 | ||||
| -rw-r--r-- | Omni/Agent/Tools/Feedback.hs | 204 | ||||
| -rw-r--r-- | Omni/Agent/Tools/Hledger.hs | 489 | ||||
| -rw-r--r-- | Omni/Agent/Tools/Http.hs | 338 | ||||
| -rw-r--r-- | Omni/Agent/Tools/Notes.hs | 357 | ||||
| -rw-r--r-- | Omni/Agent/Tools/Outreach.hs | 513 | ||||
| -rw-r--r-- | Omni/Agent/Tools/Pdf.hs | 180 | ||||
| -rw-r--r-- | Omni/Agent/Tools/Python.hs | 217 | ||||
| -rw-r--r-- | Omni/Agent/Tools/Todos.hs | 527 | ||||
| -rw-r--r-- | Omni/Agent/Tools/WebReader.hs | 308 | ||||
| -rw-r--r-- | Omni/Agent/Tools/WebReaderTest.hs | 53 | ||||
| -rw-r--r-- | Omni/Agent/Tools/WebSearch.hs | 212 |
13 files changed, 4395 insertions, 0 deletions
diff --git a/Omni/Agent/Tools/Calendar.hs b/Omni/Agent/Tools/Calendar.hs new file mode 100644 index 0000000..805916f --- /dev/null +++ b/Omni/Agent/Tools/Calendar.hs @@ -0,0 +1,322 @@ +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE NoImplicitPrelude #-} + +-- | Calendar tool using khal CLI. +-- +-- Provides calendar access for agents via local khal/CalDAV. +-- +-- : out omni-agent-tools-calendar +-- : dep aeson +-- : dep process +module Omni.Agent.Tools.Calendar + ( -- * Tools + calendarListTool, + calendarAddTool, + calendarSearchTool, + + -- * Direct API + listEvents, + addEvent, + searchEvents, + listCalendars, + + -- * Testing + main, + test, + ) +where + +import Alpha +import Data.Aeson ((.!=), (.:), (.:?), (.=)) +import qualified Data.Aeson as Aeson +import qualified Data.Text as Text +import qualified Omni.Agent.Engine as Engine +import qualified Omni.Test as Test +import System.Process (readProcessWithExitCode) + +main :: IO () +main = Test.run test + +test :: Test.Tree +test = + Test.group + "Omni.Agent.Tools.Calendar" + [ Test.unit "calendarListTool has correct schema" <| do + let tool = calendarListTool + Engine.toolName tool Test.@=? "calendar_list", + Test.unit "calendarAddTool has correct schema" <| do + let tool = calendarAddTool + Engine.toolName tool Test.@=? "calendar_add", + Test.unit "calendarSearchTool has correct schema" <| do + let tool = calendarSearchTool + Engine.toolName tool Test.@=? "calendar_search", + Test.unit "listCalendars returns calendars" <| do + result <- listCalendars + case result of + Left _ -> pure () + Right cals -> (not (null cals) || null cals) Test.@=? True + ] + +defaultCalendars :: [String] +defaultCalendars = ["BenSimaShared", "Kate"] + +listEvents :: Text -> Maybe Text -> IO (Either Text Text) +listEvents range maybeCalendar = do + let rangeArg = if Text.null range then "today 7d" else Text.unpack range + calArgs = case maybeCalendar of + Just cal -> ["-a", Text.unpack cal] + Nothing -> concatMap (\c -> ["-a", c]) defaultCalendars + formatArg = ["-f", "[{calendar}] {title} | {start-time} - {end-time}"] + result <- + try <| readProcessWithExitCode "khal" (["list"] <> calArgs <> formatArg <> [rangeArg, "-o"]) "" + case result of + Left (e :: SomeException) -> + pure (Left ("khal error: " <> tshow e)) + Right (exitCode, stdoutStr, stderrStr) -> + case exitCode of + ExitSuccess -> pure (Right (Text.pack stdoutStr)) + ExitFailure code -> + pure (Left ("khal failed (" <> tshow code <> "): " <> Text.pack stderrStr)) + +addEvent :: Text -> Text -> Maybe Text -> Maybe Text -> Maybe Text -> IO (Either Text Text) +addEvent calendarName eventSpec location alarm description = do + let baseArgs = ["new", "-a", Text.unpack calendarName] + locArgs = maybe [] (\l -> ["-l", Text.unpack l]) location + alarmArgs = maybe [] (\a -> ["-m", Text.unpack a]) alarm + specParts = Text.unpack eventSpec + descParts = maybe [] (\d -> ["::", Text.unpack d]) description + allArgs = baseArgs <> locArgs <> alarmArgs <> [specParts] <> descParts + result <- try <| readProcessWithExitCode "khal" allArgs "" + case result of + Left (e :: SomeException) -> + pure (Left ("khal error: " <> tshow e)) + Right (exitCode, stdoutStr, stderrStr) -> + case exitCode of + ExitSuccess -> + pure (Right ("Event created: " <> Text.pack stdoutStr)) + ExitFailure code -> + pure (Left ("khal failed (" <> tshow code <> "): " <> Text.pack stderrStr)) + +searchEvents :: Text -> IO (Either Text Text) +searchEvents query = do + let calArgs = concatMap (\c -> ["-a", c]) defaultCalendars + result <- + try <| readProcessWithExitCode "khal" (["search"] <> calArgs <> [Text.unpack query]) "" + case result of + Left (e :: SomeException) -> + pure (Left ("khal error: " <> tshow e)) + Right (exitCode, stdoutStr, stderrStr) -> + case exitCode of + ExitSuccess -> pure (Right (Text.pack stdoutStr)) + ExitFailure code -> + pure (Left ("khal failed (" <> tshow code <> "): " <> Text.pack stderrStr)) + +listCalendars :: IO (Either Text [Text]) +listCalendars = do + result <- + try <| readProcessWithExitCode "khal" ["printcalendars"] "" + case result of + Left (e :: SomeException) -> + pure (Left ("khal error: " <> tshow e)) + Right (exitCode, stdoutStr, stderrStr) -> + case exitCode of + ExitSuccess -> + pure (Right (filter (not <. Text.null) (Text.lines (Text.pack stdoutStr)))) + ExitFailure code -> + pure (Left ("khal failed (" <> tshow code <> "): " <> Text.pack stderrStr)) + +calendarListTool :: Engine.Tool +calendarListTool = + Engine.Tool + { Engine.toolName = "calendar_list", + Engine.toolDescription = + "List upcoming calendar events. Use to check what's scheduled. " + <> "Range can be like 'today', 'tomorrow', 'today 7d', 'next week', etc. " + <> "Available calendars: BenSimaShared, Kate.", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "range" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Time range like 'today 7d', 'tomorrow', 'next week' (default: today 7d)" :: Text) + ], + "calendar" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Filter to specific calendar: 'BenSimaShared' or 'Kate' (default: both)" :: Text) + ] + ], + "required" .= ([] :: [Text]) + ], + Engine.toolExecute = executeCalendarList + } + +executeCalendarList :: Aeson.Value -> IO Aeson.Value +executeCalendarList v = + case Aeson.fromJSON v of + Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e]) + Aeson.Success (args :: CalendarListArgs) -> do + result <- listEvents (clRange args) (clCalendar args) + case result of + Left err -> + pure (Aeson.object ["error" .= err]) + Right events -> + pure + ( Aeson.object + [ "success" .= True, + "events" .= events + ] + ) + +data CalendarListArgs = CalendarListArgs + { clRange :: Text, + clCalendar :: Maybe Text + } + deriving (Generic) + +instance Aeson.FromJSON CalendarListArgs where + parseJSON = + Aeson.withObject "CalendarListArgs" <| \v -> + (CalendarListArgs </ (v .:? "range" .!= "today 7d")) + <*> (v .:? "calendar") + +calendarAddTool :: Engine.Tool +calendarAddTool = + Engine.Tool + { Engine.toolName = "calendar_add", + Engine.toolDescription = + "Add a new calendar event. The event_spec format is: " + <> "'START [END] SUMMARY' where START/END are dates or times. " + <> "Examples: '2024-12-25 Christmas', 'tomorrow 10:00 11:00 Meeting', " + <> "'friday 14:00 1h Doctor appointment'.", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "calendar" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Calendar name to add to (e.g., 'BenSimaShared', 'Kate')" :: Text) + ], + "event_spec" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Event specification: 'START [END] SUMMARY' (e.g., 'tomorrow 10:00 11:00 Team meeting')" :: Text) + ], + "location" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Location of the event (optional)" :: Text) + ], + "alarm" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Alarm time before event, e.g., '15m', '1h', '1d' (optional)" :: Text) + ], + "description" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Detailed description of the event (optional)" :: Text) + ] + ], + "required" .= (["calendar", "event_spec"] :: [Text]) + ], + Engine.toolExecute = executeCalendarAdd + } + +executeCalendarAdd :: Aeson.Value -> IO Aeson.Value +executeCalendarAdd v = + case Aeson.fromJSON v of + Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e]) + Aeson.Success (args :: CalendarAddArgs) -> do + result <- + addEvent + (caCalendar args) + (caEventSpec args) + (caLocation args) + (caAlarm args) + (caDescription args) + case result of + Left err -> + pure (Aeson.object ["error" .= err]) + Right msg -> + pure + ( Aeson.object + [ "success" .= True, + "message" .= msg + ] + ) + +data CalendarAddArgs = CalendarAddArgs + { caCalendar :: Text, + caEventSpec :: Text, + caLocation :: Maybe Text, + caAlarm :: Maybe Text, + caDescription :: Maybe Text + } + deriving (Generic) + +instance Aeson.FromJSON CalendarAddArgs where + parseJSON = + Aeson.withObject "CalendarAddArgs" <| \v -> + (CalendarAddArgs </ (v .: "calendar")) + <*> (v .: "event_spec") + <*> (v .:? "location") + <*> (v .:? "alarm") + <*> (v .:? "description") + +calendarSearchTool :: Engine.Tool +calendarSearchTool = + Engine.Tool + { Engine.toolName = "calendar_search", + Engine.toolDescription = + "Search for calendar events by text. Finds events matching the query " + <> "in title, description, or location.", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "query" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Search text to find in events" :: Text) + ] + ], + "required" .= (["query"] :: [Text]) + ], + Engine.toolExecute = executeCalendarSearch + } + +executeCalendarSearch :: Aeson.Value -> IO Aeson.Value +executeCalendarSearch v = + case Aeson.fromJSON v of + Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e]) + Aeson.Success (args :: CalendarSearchArgs) -> do + result <- searchEvents (csQuery args) + case result of + Left err -> + pure (Aeson.object ["error" .= err]) + Right events -> + pure + ( Aeson.object + [ "success" .= True, + "results" .= events + ] + ) + +newtype CalendarSearchArgs = CalendarSearchArgs + { csQuery :: Text + } + deriving (Generic) + +instance Aeson.FromJSON CalendarSearchArgs where + parseJSON = + Aeson.withObject "CalendarSearchArgs" <| \v -> + CalendarSearchArgs </ (v .: "query") diff --git a/Omni/Agent/Tools/Email.hs b/Omni/Agent/Tools/Email.hs new file mode 100644 index 0000000..7a9bc64 --- /dev/null +++ b/Omni/Agent/Tools/Email.hs @@ -0,0 +1,675 @@ +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE NoImplicitPrelude #-} + +-- | Email tools for IMAP and SMTP access via Telegram bot. +-- +-- Provides email management for agents: +-- - Check for urgent/time-sensitive emails +-- - Identify emails needing response vs FYI +-- - Auto-unsubscribe from marketing +-- - Send approved outreach emails via SMTP +-- +-- Uses HaskellNet for IMAP/SMTP client support. +-- Password retrieved via `pass ben@bensima.com`. +-- +-- : out omni-agent-tools-email +-- : dep aeson +-- : dep process +-- : dep regex-applicative +-- : dep http-conduit +-- : dep HaskellNet +-- : dep HaskellNet-SSL +module Omni.Agent.Tools.Email + ( -- * Tools + emailCheckTool, + emailReadTool, + emailUnsubscribeTool, + emailArchiveTool, + emailSendTool, + + -- * All tools + allEmailTools, + + -- * Direct API + checkNewEmails, + readEmail, + unsubscribeFromEmail, + archiveEmail, + getPassword, + sendApprovedEmail, + + -- * Scheduled Check + emailCheckLoop, + performScheduledCheck, + + -- * 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.ByteString.Char8 as BS8 +import qualified Data.List as List +import qualified Data.Text as Text +import qualified Data.Text.Lazy as LText +import Data.Time (NominalDiffTime, UTCTime, addUTCTime, getCurrentTime) +import Data.Time.Format (defaultTimeLocale, formatTime, parseTimeM) +import Data.Time.LocalTime (TimeZone (..), utcToZonedTime) +import qualified Network.HTTP.Simple as HTTP +import qualified Network.HaskellNet.IMAP as IMAP +import Network.HaskellNet.IMAP.Connection (IMAPConnection) +import qualified Network.HaskellNet.IMAP.SSL as IMAPSSL +import qualified Network.HaskellNet.SMTP as SMTP +import qualified Network.HaskellNet.SMTP.SSL as SMTPSSL +import Network.Mail.Mime (Address (..), simpleMail') +import qualified Omni.Agent.Engine as Engine +import qualified Omni.Agent.Tools.Outreach as Outreach +import qualified Omni.Test as Test +import System.Process (readProcessWithExitCode) +import Text.Regex.Applicative (RE, anySym, few, (=~)) +import qualified Text.Regex.Applicative as RE + +main :: IO () +main = Test.run test + +test :: Test.Tree +test = + Test.group + "Omni.Agent.Tools.Email" + [ Test.unit "emailCheckTool has correct name" <| do + Engine.toolName emailCheckTool Test.@=? "email_check", + Test.unit "emailReadTool has correct name" <| do + Engine.toolName emailReadTool Test.@=? "email_read", + Test.unit "emailUnsubscribeTool has correct name" <| do + Engine.toolName emailUnsubscribeTool Test.@=? "email_unsubscribe", + Test.unit "emailArchiveTool has correct name" <| do + Engine.toolName emailArchiveTool Test.@=? "email_archive", + Test.unit "emailSendTool has correct name" <| do + Engine.toolName emailSendTool Test.@=? "email_send", + Test.unit "allEmailTools has 5 tools" <| do + length allEmailTools Test.@=? 5, + Test.unit "parseEmailHeaders extracts fields" <| do + let headers = + "From: test@example.com\r\n\ + \Subject: Test Subject\r\n\ + \Date: Mon, 1 Jan 2024 12:00:00 +0000\r\n\ + \\r\n" + case parseEmailHeaders headers of + Nothing -> Test.assertFailure "Failed to parse headers" + Just email -> do + emailFrom email Test.@=? "test@example.com" + emailSubject email Test.@=? "Test Subject", + Test.unit "parseUnsubscribeHeader extracts URL" <| do + let header = "<https://example.com/unsubscribe>, <mailto:unsub@example.com>" + case parseUnsubscribeUrl header of + Nothing -> Test.assertFailure "Failed to parse unsubscribe URL" + Just url -> ("https://example.com" `Text.isPrefixOf` url) Test.@=? True + ] + +imapServer :: String +imapServer = "bensima.com" + +imapUser :: String +imapUser = "ben@bensima.com" + +getPassword :: IO (Either Text Text) +getPassword = do + result <- try <| readProcessWithExitCode "pass" ["ben@bensima.com"] "" + case result of + Left (e :: SomeException) -> + pure (Left ("Failed to get password: " <> tshow e)) + Right (exitCode, stdoutStr, stderrStr) -> + case exitCode of + ExitSuccess -> pure (Right (Text.strip (Text.pack stdoutStr))) + ExitFailure code -> + pure (Left ("pass failed (" <> tshow code <> "): " <> Text.pack stderrStr)) + +withImapConnection :: (IMAPConnection -> IO a) -> IO (Either Text a) +withImapConnection action = do + pwResult <- getPassword + case pwResult of + Left err -> pure (Left err) + Right pw -> do + result <- + try <| do + conn <- IMAPSSL.connectIMAPSSL imapServer + IMAP.login conn imapUser (Text.unpack pw) + r <- action conn + IMAP.logout conn + pure r + case result of + Left (e :: SomeException) -> pure (Left ("IMAP error: " <> tshow e)) + Right r -> pure (Right r) + +data EmailSummary = EmailSummary + { emailUid :: Int, + emailFrom :: Text, + emailSubject :: Text, + emailDate :: Text, + emailUnsubscribe :: Maybe Text + } + deriving (Show, Generic) + +instance Aeson.ToJSON EmailSummary where + toJSON e = + Aeson.object + [ "uid" .= emailUid e, + "from" .= emailFrom e, + "subject" .= emailSubject e, + "date" .= formatDateAsEst (emailDate e), + "has_unsubscribe" .= isJust (emailUnsubscribe e) + ] + +estTimezone :: TimeZone +estTimezone = TimeZone (-300) False "EST" + +formatDateAsEst :: Text -> Text +formatDateAsEst dateStr = + case parseEmailDate dateStr of + Nothing -> dateStr + Just utcTime -> + let zonedTime = utcToZonedTime estTimezone utcTime + in Text.pack (formatTime defaultTimeLocale "%a %b %d %H:%M EST" zonedTime) + +parseEmailHeaders :: Text -> Maybe EmailSummary +parseEmailHeaders raw = do + let headerLines = Text.lines raw + fromLine = findHeader "From:" headerLines + subjectLine = findHeader "Subject:" headerLines + dateLine = findHeader "Date:" headerLines + unsubLine = findHeader "List-Unsubscribe:" headerLines + fromVal <- fromLine + subject <- subjectLine + dateVal <- dateLine + pure + EmailSummary + { emailUid = 0, + emailFrom = Text.strip (Text.drop 5 fromVal), + emailSubject = Text.strip (Text.drop 8 subject), + emailDate = Text.strip (Text.drop 5 dateVal), + emailUnsubscribe = (parseUnsubscribeUrl <. Text.drop 16) =<< unsubLine + } + where + findHeader :: Text -> [Text] -> Maybe Text + findHeader prefix = List.find (prefix `Text.isPrefixOf`) + +parseUnsubscribeUrl :: Text -> Maybe Text +parseUnsubscribeUrl header = + let text = Text.unpack header + in case text =~ urlInBrackets of + Just url | "http" `List.isPrefixOf` url -> Just (Text.pack url) + _ -> Nothing + where + urlInBrackets :: RE Char String + urlInBrackets = few anySym *> RE.sym '<' *> few anySym <* RE.sym '>' + +checkNewEmails :: Maybe Int -> Maybe Int -> IO (Either Text [EmailSummary]) +checkNewEmails maybeLimit maybeHours = do + withImapConnection <| \conn -> do + IMAP.select conn "INBOX" + uids <- IMAP.search conn [IMAP.UNFLAG IMAP.Seen] + let limit = fromMaybe 20 maybeLimit + recentUids = take limit (reverse (map fromIntegral uids)) + if null recentUids + then pure [] + else do + emails <- + forM recentUids <| \uid -> do + headerBytes <- IMAP.fetchHeader conn (fromIntegral uid) + let headerText = Text.pack (BS8.unpack headerBytes) + pure (parseEmailHeaders headerText, uid) + let parsed = + [ e {emailUid = uid} + | (Just e, uid) <- emails + ] + case maybeHours of + Nothing -> pure parsed + Just hours -> do + now <- getCurrentTime + let cutoff = addUTCTime (negate (fromIntegral hours * 3600 :: NominalDiffTime)) now + pure (filter (isAfterCutoff cutoff) parsed) + +isAfterCutoff :: UTCTime -> EmailSummary -> Bool +isAfterCutoff cutoff email = + case parseEmailDate (emailDate email) of + Nothing -> False + Just emailTime -> emailTime >= cutoff + +parseEmailDate :: Text -> Maybe UTCTime +parseEmailDate dateStr = + let cleaned = stripParenTz (Text.strip dateStr) + formats = + [ "%a, %d %b %Y %H:%M:%S %z", + "%a, %d %b %Y %H:%M:%S %Z", + "%d %b %Y %H:%M:%S %z", + "%a, %d %b %Y %H:%M %z", + "%a, %d %b %Y %H:%M:%S %z (%Z)" + ] + tryParse [] = Nothing + tryParse (fmt : rest) = + case parseTimeM True defaultTimeLocale fmt (Text.unpack cleaned) of + Just t -> Just t + Nothing -> tryParse rest + in tryParse formats + +stripParenTz :: Text -> Text +stripParenTz t = + case Text.breakOn " (" t of + (before, after) + | Text.null after -> t + | ")" `Text.isSuffixOf` after -> before + | otherwise -> t + +readEmail :: Int -> IO (Either Text Text) +readEmail uid = + withImapConnection <| \conn -> do + IMAP.select conn "INBOX" + bodyBytes <- IMAP.fetch conn (fromIntegral uid) + let bodyText = Text.pack (BS8.unpack bodyBytes) + pure (Text.take 10000 bodyText) + +unsubscribeFromEmail :: Int -> IO (Either Text Text) +unsubscribeFromEmail uid = do + headerResult <- + withImapConnection <| \conn -> do + IMAP.select conn "INBOX" + headerBytes <- IMAP.fetchHeader conn (fromIntegral uid) + pure (Text.pack (BS8.unpack headerBytes)) + case headerResult of + Left err -> pure (Left err) + Right headerText -> + case extractUnsubscribeUrl headerText of + Nothing -> pure (Left "No unsubscribe URL found in this email") + Just url -> do + clickResult <- clickUnsubscribeLink url + case clickResult of + Left err -> pure (Left ("Failed to unsubscribe: " <> err)) + Right () -> do + _ <- archiveEmail uid + pure (Right ("Unsubscribed and archived email " <> tshow uid)) + +extractUnsubscribeUrl :: Text -> Maybe Text +extractUnsubscribeUrl headerText = + let unsubLine = List.find ("List-Unsubscribe:" `Text.isInfixOf`) (Text.lines headerText) + in (parseUnsubscribeUrl <. Text.drop 16 <. Text.strip) =<< unsubLine + +clickUnsubscribeLink :: Text -> IO (Either Text ()) +clickUnsubscribeLink url = do + result <- + try <| do + req <- HTTP.parseRequest (Text.unpack url) + _ <- HTTP.httpLBS req + pure () + case result of + Left (e :: SomeException) -> pure (Left (tshow e)) + Right () -> pure (Right ()) + +archiveEmail :: Int -> IO (Either Text Text) +archiveEmail uid = + withImapConnection <| \conn -> do + IMAP.select conn "INBOX" + IMAP.copy conn (fromIntegral uid) "Archives.2025" + IMAP.store conn (fromIntegral uid) (IMAP.PlusFlags [IMAP.Deleted]) + _ <- IMAP.expunge conn + pure ("Archived email " <> tshow uid) + +allEmailTools :: [Engine.Tool] +allEmailTools = + [ emailCheckTool, + emailReadTool, + emailUnsubscribeTool, + emailArchiveTool, + emailSendTool + ] + +emailCheckTool :: Engine.Tool +emailCheckTool = + Engine.Tool + { Engine.toolName = "email_check", + Engine.toolDescription = + "Check for new/unread emails. Returns a summary of recent unread emails " + <> "including sender, subject, date, and whether they have an unsubscribe link. " + <> "Use this to identify urgent items or emails needing response. " + <> "Use 'hours' to filter to emails received in the last N hours (e.g., hours=6 for last 6 hours).", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "limit" + .= Aeson.object + [ "type" .= ("integer" :: Text), + "description" .= ("Max emails to return (default: 20)" :: Text) + ], + "hours" + .= Aeson.object + [ "type" .= ("integer" :: Text), + "description" .= ("Only return emails from the last N hours (e.g., 6 for last 6 hours)" :: Text) + ] + ], + "required" .= ([] :: [Text]) + ], + Engine.toolExecute = executeEmailCheck + } + +executeEmailCheck :: Aeson.Value -> IO Aeson.Value +executeEmailCheck v = do + let (limit, hours) = case v of + Aeson.Object obj -> + let l = case KeyMap.lookup "limit" obj of + Just (Aeson.Number n) -> Just (round n :: Int) + _ -> Nothing + h = case KeyMap.lookup "hours" obj of + Just (Aeson.Number n) -> Just (round n :: Int) + _ -> Nothing + in (l, h) + _ -> (Nothing, Nothing) + result <- checkNewEmails limit hours + case result of + Left err -> pure (Aeson.object ["error" .= err]) + Right emails -> + pure + ( Aeson.object + [ "success" .= True, + "count" .= length emails, + "emails" .= emails + ] + ) + +emailReadTool :: Engine.Tool +emailReadTool = + Engine.Tool + { Engine.toolName = "email_read", + Engine.toolDescription = + "Read the full content of an email by its UID. " + <> "Use after email_check to read emails that seem important or need a response.", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "uid" + .= Aeson.object + [ "type" .= ("integer" :: Text), + "description" .= ("Email UID from email_check" :: Text) + ] + ], + "required" .= (["uid"] :: [Text]) + ], + Engine.toolExecute = executeEmailRead + } + +executeEmailRead :: Aeson.Value -> IO Aeson.Value +executeEmailRead v = do + let uidM = case v of + Aeson.Object obj -> case KeyMap.lookup "uid" obj of + Just (Aeson.Number n) -> Just (round n :: Int) + _ -> Nothing + _ -> Nothing + case uidM of + Nothing -> pure (Aeson.object ["error" .= ("Missing uid parameter" :: Text)]) + Just uid -> do + result <- readEmail uid + case result of + Left err -> pure (Aeson.object ["error" .= err]) + Right body -> + pure + ( Aeson.object + [ "success" .= True, + "uid" .= uid, + "body" .= body + ] + ) + +emailUnsubscribeTool :: Engine.Tool +emailUnsubscribeTool = + Engine.Tool + { Engine.toolName = "email_unsubscribe", + Engine.toolDescription = + "Unsubscribe from a mailing list by clicking the List-Unsubscribe link. " + <> "Use for marketing/newsletter emails. Automatically archives the email after unsubscribing.", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "uid" + .= Aeson.object + [ "type" .= ("integer" :: Text), + "description" .= ("Email UID to unsubscribe from" :: Text) + ] + ], + "required" .= (["uid"] :: [Text]) + ], + Engine.toolExecute = executeEmailUnsubscribe + } + +executeEmailUnsubscribe :: Aeson.Value -> IO Aeson.Value +executeEmailUnsubscribe v = do + let uidM = case v of + Aeson.Object obj -> case KeyMap.lookup "uid" obj of + Just (Aeson.Number n) -> Just (round n :: Int) + _ -> Nothing + _ -> Nothing + case uidM of + Nothing -> pure (Aeson.object ["error" .= ("Missing uid parameter" :: Text)]) + Just uid -> do + result <- unsubscribeFromEmail uid + case result of + Left err -> pure (Aeson.object ["error" .= err]) + Right msg -> + pure + ( Aeson.object + [ "success" .= True, + "message" .= msg + ] + ) + +emailArchiveTool :: Engine.Tool +emailArchiveTool = + Engine.Tool + { Engine.toolName = "email_archive", + Engine.toolDescription = + "Archive an email (move to Archives.2025 folder). " + <> "Use for emails that don't need a response and are just FYI.", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "uid" + .= Aeson.object + [ "type" .= ("integer" :: Text), + "description" .= ("Email UID to archive" :: Text) + ] + ], + "required" .= (["uid"] :: [Text]) + ], + Engine.toolExecute = executeEmailArchive + } + +executeEmailArchive :: Aeson.Value -> IO Aeson.Value +executeEmailArchive v = do + let uidM = case v of + Aeson.Object obj -> case KeyMap.lookup "uid" obj of + Just (Aeson.Number n) -> Just (round n :: Int) + _ -> Nothing + _ -> Nothing + case uidM of + Nothing -> pure (Aeson.object ["error" .= ("Missing uid parameter" :: Text)]) + Just uid -> do + result <- archiveEmail uid + case result of + Left err -> pure (Aeson.object ["error" .= err]) + Right msg -> + pure + ( Aeson.object + [ "success" .= True, + "message" .= msg + ] + ) + +emailCheckLoop :: (Int -> Maybe Int -> Text -> IO (Maybe Int)) -> Int -> IO () +emailCheckLoop sendFn chatId = + forever <| do + let sixHours = 6 * 60 * 60 * 1000000 + threadDelay sixHours + performScheduledCheck sendFn chatId + +performScheduledCheck :: (Int -> Maybe Int -> Text -> IO (Maybe Int)) -> Int -> IO () +performScheduledCheck sendFn chatId = do + putText "Running scheduled email check..." + result <- checkNewEmails (Just 50) (Just 6) + case result of + Left err -> putText ("Email check failed: " <> err) + Right emails -> do + let urgent = filter isUrgent emails + needsResponse = filter needsResponsePred emails + marketing = filter hasUnsubscribe emails + when (not (null urgent) || not (null needsResponse)) <| do + let msg = formatEmailSummary urgent needsResponse (length marketing) + _ <- sendFn chatId Nothing msg + pure () + where + isUrgent :: EmailSummary -> Bool + isUrgent email = + let subj = Text.toLower (emailSubject email) + in "urgent" + `Text.isInfixOf` subj + || "asap" + `Text.isInfixOf` subj + || "important" + `Text.isInfixOf` subj + || "action required" + `Text.isInfixOf` subj + + needsResponsePred :: EmailSummary -> Bool + needsResponsePred email = + let sender = Text.toLower (emailFrom email) + subj = Text.toLower (emailSubject email) + in not (hasUnsubscribe email) + && not (isUrgent email) + && not ("noreply" `Text.isInfixOf` sender) + && not ("no-reply" `Text.isInfixOf` sender) + && ("?" `Text.isInfixOf` subj || "reply" `Text.isInfixOf` subj || "response" `Text.isInfixOf` subj) + + hasUnsubscribe :: EmailSummary -> Bool + hasUnsubscribe = isJust <. emailUnsubscribe + + formatEmailSummary :: [EmailSummary] -> [EmailSummary] -> Int -> Text + formatEmailSummary urgent needs marketingCount = + Text.unlines + <| ["📧 *email check*", ""] + <> (if null urgent then [] else ["*urgent:*"] <> map formatOne urgent <> [""]) + <> (if null needs then [] else ["*may need response:*"] <> map formatOne needs <> [""]) + <> [tshow marketingCount <> " marketing emails (use email_check to review)"] + + formatOne :: EmailSummary -> Text + formatOne e = + "• " <> emailSubject e <> " (from: " <> emailFrom e <> ", uid: " <> tshow (emailUid e) <> ")" + +smtpServer :: String +smtpServer = "bensima.com" + +smtpUser :: String +smtpUser = "ben@bensima.com" + +withSmtpConnection :: (SMTP.SMTPConnection -> IO a) -> IO (Either Text a) +withSmtpConnection action = do + pwResult <- getPassword + case pwResult of + Left err -> pure (Left err) + Right pw -> do + result <- + try <| do + conn <- SMTPSSL.connectSMTPSSL smtpServer + authSuccess <- SMTP.authenticate SMTP.LOGIN smtpUser (Text.unpack pw) conn + if authSuccess + then do + r <- action conn + SMTP.closeSMTP conn + pure r + else do + SMTP.closeSMTP conn + panic "SMTP authentication failed" + case result of + Left (e :: SomeException) -> pure (Left ("SMTP error: " <> tshow e)) + Right r -> pure (Right r) + +sendApprovedEmail :: Text -> IO (Either Text Text) +sendApprovedEmail draftId = do + mDraft <- Outreach.getDraft draftId + case mDraft of + Nothing -> pure (Left "Draft not found") + Just draft -> do + case Outreach.draftStatus draft of + Outreach.Approved -> do + let recipientAddr = Address Nothing (Outreach.draftRecipient draft) + senderAddr = Address (Just "Ben Sima") "ben@bensima.com" + subject = fromMaybe "" (Outreach.draftSubject draft) + body = LText.fromStrict (Outreach.draftBody draft) + footer = "\n\n---\nSent by Ava on behalf of Ben" + fullBody = body <> footer + mail = simpleMail' recipientAddr senderAddr subject fullBody + sendResult <- + withSmtpConnection <| \conn -> do + SMTP.sendMail mail conn + case sendResult of + Left err -> pure (Left err) + Right () -> do + _ <- Outreach.markSent draftId + pure (Right ("Email sent to " <> Outreach.draftRecipient draft)) + Outreach.Pending -> pure (Left "Draft is still pending approval") + Outreach.Rejected -> pure (Left "Draft was rejected") + Outreach.Sent -> pure (Left "Draft was already sent") + +emailSendTool :: Engine.Tool +emailSendTool = + Engine.Tool + { Engine.toolName = "email_send", + Engine.toolDescription = + "Send an approved outreach email. Only sends emails that have been approved " + <> "by Ben in the outreach queue. Use outreach_draft to create drafts first, " + <> "wait for approval, then use this to send.", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "draft_id" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("ID of the approved draft to send" :: Text) + ] + ], + "required" .= (["draft_id"] :: [Text]) + ], + Engine.toolExecute = executeEmailSend + } + +executeEmailSend :: Aeson.Value -> IO Aeson.Value +executeEmailSend v = do + let draftIdM = case v of + Aeson.Object obj -> case KeyMap.lookup "draft_id" obj of + Just (Aeson.String s) -> Just s + _ -> Nothing + _ -> Nothing + case draftIdM of + Nothing -> pure (Aeson.object ["error" .= ("Missing draft_id parameter" :: Text)]) + Just draftId -> do + result <- sendApprovedEmail draftId + case result of + Left err -> pure (Aeson.object ["error" .= err]) + Right msg -> + pure + ( Aeson.object + [ "success" .= True, + "message" .= msg + ] + ) diff --git a/Omni/Agent/Tools/Feedback.hs b/Omni/Agent/Tools/Feedback.hs new file mode 100644 index 0000000..1ec684c --- /dev/null +++ b/Omni/Agent/Tools/Feedback.hs @@ -0,0 +1,204 @@ +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE NoImplicitPrelude #-} + +-- | Feedback query tool for PodcastItLater user research. +-- +-- Allows the agent to query collected feedback from the PIL database. +-- Feedback is submitted via /feedback on the PIL web app. +-- +-- : out omni-agent-tools-feedback +-- : dep aeson +-- : dep http-conduit +module Omni.Agent.Tools.Feedback + ( -- * Tools + feedbackListTool, + allFeedbackTools, + + -- * Types + FeedbackEntry (..), + ListFeedbackArgs (..), + + -- * Testing + main, + test, + ) +where + +import Alpha +import Data.Aeson ((.!=), (.:), (.:?), (.=)) +import qualified Data.Aeson as Aeson +import qualified Data.Text as Text +import qualified Network.HTTP.Simple as HTTP +import qualified Omni.Agent.Engine as Engine +import qualified Omni.Test as Test +import System.Environment (lookupEnv) + +main :: IO () +main = Test.run test + +test :: Test.Tree +test = + Test.group + "Omni.Agent.Tools.Feedback" + [ Test.unit "feedbackListTool has correct name" <| do + Engine.toolName feedbackListTool Test.@=? "feedback_list", + Test.unit "allFeedbackTools has 1 tool" <| do + length allFeedbackTools Test.@=? 1, + Test.unit "ListFeedbackArgs parses correctly" <| do + let json = Aeson.object ["limit" .= (10 :: Int)] + case Aeson.fromJSON json of + Aeson.Success (args :: ListFeedbackArgs) -> lfaLimit args Test.@=? 10 + Aeson.Error e -> Test.assertFailure e, + Test.unit "ListFeedbackArgs parses with since" <| do + let json = + Aeson.object + [ "limit" .= (20 :: Int), + "since" .= ("2024-01-01" :: Text) + ] + case Aeson.fromJSON json of + Aeson.Success (args :: ListFeedbackArgs) -> do + lfaLimit args Test.@=? 20 + lfaSince args Test.@=? Just "2024-01-01" + Aeson.Error e -> Test.assertFailure e, + Test.unit "FeedbackEntry JSON roundtrip" <| do + let entry = + FeedbackEntry + { feId = "abc123", + feEmail = Just "test@example.com", + feSource = Just "outreach", + feCampaignId = Nothing, + feRating = Just 4, + feFeedbackText = Just "Great product!", + feUseCase = Just "Commute listening", + feCreatedAt = "2024-01-15T10:00:00Z" + } + case Aeson.decode (Aeson.encode entry) of + Nothing -> Test.assertFailure "Failed to decode FeedbackEntry" + Just decoded -> do + feId decoded Test.@=? "abc123" + feEmail decoded Test.@=? Just "test@example.com" + feRating decoded Test.@=? Just 4 + ] + +data FeedbackEntry = FeedbackEntry + { feId :: Text, + feEmail :: Maybe Text, + feSource :: Maybe Text, + feCampaignId :: Maybe Text, + feRating :: Maybe Int, + feFeedbackText :: Maybe Text, + feUseCase :: Maybe Text, + feCreatedAt :: Text + } + deriving (Show, Eq, Generic) + +instance Aeson.ToJSON FeedbackEntry where + toJSON e = + Aeson.object + [ "id" .= feId e, + "email" .= feEmail e, + "source" .= feSource e, + "campaign_id" .= feCampaignId e, + "rating" .= feRating e, + "feedback_text" .= feFeedbackText e, + "use_case" .= feUseCase e, + "created_at" .= feCreatedAt e + ] + +instance Aeson.FromJSON FeedbackEntry where + parseJSON = + Aeson.withObject "FeedbackEntry" <| \v -> + (FeedbackEntry </ (v .: "id")) + <*> (v .:? "email") + <*> (v .:? "source") + <*> (v .:? "campaign_id") + <*> (v .:? "rating") + <*> (v .:? "feedback_text") + <*> (v .:? "use_case") + <*> (v .: "created_at") + +data ListFeedbackArgs = ListFeedbackArgs + { lfaLimit :: Int, + lfaSince :: Maybe Text + } + deriving (Show, Eq, Generic) + +instance Aeson.FromJSON ListFeedbackArgs where + parseJSON = + Aeson.withObject "ListFeedbackArgs" <| \v -> + (ListFeedbackArgs </ (v .:? "limit" .!= 20)) + <*> (v .:? "since") + +allFeedbackTools :: [Engine.Tool] +allFeedbackTools = [feedbackListTool] + +feedbackListTool :: Engine.Tool +feedbackListTool = + Engine.Tool + { Engine.toolName = "feedback_list", + Engine.toolDescription = + "List feedback entries from PodcastItLater users. " + <> "Use to review user research data and understand what potential " + <> "customers want from the product.", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "limit" + .= Aeson.object + [ "type" .= ("integer" :: Text), + "description" .= ("Max entries to return (default: 20)" :: Text) + ], + "since" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("ISO date to filter by (entries after this date)" :: Text) + ] + ], + "required" .= ([] :: [Text]) + ], + Engine.toolExecute = executeFeedbackList + } + +executeFeedbackList :: Aeson.Value -> IO Aeson.Value +executeFeedbackList v = + case Aeson.fromJSON v of + Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e]) + Aeson.Success (args :: ListFeedbackArgs) -> do + mBaseUrl <- lookupEnv "PIL_BASE_URL" + let baseUrl = maybe "http://localhost:8000" Text.pack mBaseUrl + limit = min 100 (max 1 (lfaLimit args)) + sinceParam = case lfaSince args of + Nothing -> "" + Just since -> "&since=" <> since + url = baseUrl <> "/api/feedback?limit=" <> tshow limit <> sinceParam + result <- fetchFeedback url + case result of + Left err -> pure (Aeson.object ["error" .= err]) + Right entries -> + pure + ( Aeson.object + [ "success" .= True, + "count" .= length entries, + "entries" .= entries + ] + ) + +fetchFeedback :: Text -> IO (Either Text [FeedbackEntry]) +fetchFeedback url = do + result <- + try <| do + req <- HTTP.parseRequest (Text.unpack url) + resp <- HTTP.httpLBS req + pure (HTTP.getResponseStatusCode resp, HTTP.getResponseBody resp) + case result of + Left (e :: SomeException) -> pure (Left ("Request failed: " <> tshow e)) + Right (status, body) -> + if status /= 200 + then pure (Left ("HTTP " <> tshow status)) + else case Aeson.decode body of + Nothing -> pure (Left "Failed to parse response") + Just entries -> pure (Right entries) diff --git a/Omni/Agent/Tools/Hledger.hs b/Omni/Agent/Tools/Hledger.hs new file mode 100644 index 0000000..59e0c05 --- /dev/null +++ b/Omni/Agent/Tools/Hledger.hs @@ -0,0 +1,489 @@ +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE NoImplicitPrelude #-} + +-- | Hledger tools for personal finance queries and transaction entry. +-- +-- Provides hledger access for agents via the nix-shell in ~/fund. +-- +-- : out omni-agent-tools-hledger +-- : dep aeson +-- : dep process +-- : dep directory +module Omni.Agent.Tools.Hledger + ( -- * Tools + hledgerBalanceTool, + hledgerRegisterTool, + hledgerAddTool, + hledgerIncomeStatementTool, + hledgerBalanceSheetTool, + + -- * All tools (for easy import) + allHledgerTools, + + -- * Direct API + queryBalance, + queryRegister, + addTransaction, + incomeStatement, + balanceSheet, + + -- * Testing + main, + test, + ) +where + +import Alpha +import Data.Aeson ((.:), (.:?), (.=)) +import qualified Data.Aeson as Aeson +import qualified Data.List as List +import qualified Data.Text as Text +import qualified Data.Text.IO as TextIO +import Data.Time (getCurrentTime, utcToLocalTime) +import Data.Time.Format (defaultTimeLocale, formatTime) +import Data.Time.LocalTime (getCurrentTimeZone) +import qualified Omni.Agent.Engine as Engine +import qualified Omni.Test as Test +import System.Directory (doesFileExist) +import System.Process (readProcessWithExitCode) + +main :: IO () +main = Test.run test + +test :: Test.Tree +test = + Test.group + "Omni.Agent.Tools.Hledger" + [ Test.unit "hledgerBalanceTool has correct name" <| do + Engine.toolName hledgerBalanceTool Test.@=? "hledger_balance", + Test.unit "hledgerRegisterTool has correct name" <| do + Engine.toolName hledgerRegisterTool Test.@=? "hledger_register", + Test.unit "hledgerAddTool has correct name" <| do + Engine.toolName hledgerAddTool Test.@=? "hledger_add", + Test.unit "hledgerIncomeStatementTool has correct name" <| do + Engine.toolName hledgerIncomeStatementTool Test.@=? "hledger_income_statement", + Test.unit "hledgerBalanceSheetTool has correct name" <| do + Engine.toolName hledgerBalanceSheetTool Test.@=? "hledger_balance_sheet", + Test.unit "allHledgerTools has 5 tools" <| do + length allHledgerTools Test.@=? 5 + ] + +fundDir :: FilePath +fundDir = "/home/ben/fund" + +journalFile :: FilePath +journalFile = fundDir <> "/ledger.journal" + +transactionsFile :: FilePath +transactionsFile = fundDir <> "/telegram-transactions.journal" + +runHledgerInFund :: [String] -> IO (Either Text Text) +runHledgerInFund args = do + let fullArgs :: [String] + fullArgs = ["-f", journalFile] <> args + hledgerCmd :: String + hledgerCmd = "hledger " ++ List.unwords fullArgs + cmd :: String + cmd = "cd " ++ fundDir ++ " && " ++ hledgerCmd + result <- + try <| readProcessWithExitCode "nix-shell" [fundDir ++ "/shell.nix", "--run", cmd] "" + case result of + Left (e :: SomeException) -> + pure (Left ("hledger error: " <> tshow e)) + Right (exitCode, stdoutStr, stderrStr) -> + case exitCode of + ExitSuccess -> pure (Right (Text.pack stdoutStr)) + ExitFailure code -> + pure (Left ("hledger failed (" <> tshow code <> "): " <> Text.pack stderrStr)) + +allHledgerTools :: [Engine.Tool] +allHledgerTools = + [ hledgerBalanceTool, + hledgerRegisterTool, + hledgerAddTool, + hledgerIncomeStatementTool, + hledgerBalanceSheetTool + ] + +queryBalance :: Maybe Text -> Maybe Text -> Maybe Text -> IO (Either Text Text) +queryBalance maybePattern maybePeriod maybeCurrency = do + let patternArg = maybe [] (\p -> [Text.unpack p]) maybePattern + periodArg = maybe [] (\p -> ["-p", "'" ++ Text.unpack p ++ "'"]) maybePeriod + currency = maybe "USD" Text.unpack maybeCurrency + currencyArg = ["-X", currency] + runHledgerInFund (["bal", "-1", "--flat"] <> currencyArg <> patternArg <> periodArg) + +queryRegister :: Text -> Maybe Int -> Maybe Text -> Maybe Text -> IO (Either Text Text) +queryRegister accountPattern maybeLimit maybeCurrency maybePeriod = do + let limitArg = maybe ["-n", "10"] (\n -> ["-n", show n]) maybeLimit + currency = maybe "USD" Text.unpack maybeCurrency + currencyArg = ["-X", currency] + periodArg = maybe [] (\p -> ["-p", "'" ++ Text.unpack p ++ "'"]) maybePeriod + runHledgerInFund (["reg", Text.unpack accountPattern] <> currencyArg <> periodArg <> limitArg) + +incomeStatement :: Maybe Text -> Maybe Text -> IO (Either Text Text) +incomeStatement maybePeriod maybeCurrency = do + let periodArg = maybe ["-p", "thismonth"] (\p -> ["-p", "'" ++ Text.unpack p ++ "'"]) maybePeriod + currency = maybe "USD" Text.unpack maybeCurrency + currencyArg = ["-X", currency] + runHledgerInFund (["is"] <> currencyArg <> periodArg) + +balanceSheet :: Maybe Text -> IO (Either Text Text) +balanceSheet maybeCurrency = do + let currency = maybe "USD" Text.unpack maybeCurrency + currencyArg = ["-X", currency] + runHledgerInFund (["bs"] <> currencyArg) + +addTransaction :: Text -> Text -> Text -> Text -> Maybe Text -> IO (Either Text Text) +addTransaction description fromAccount toAccount amount maybeDate = do + now <- getCurrentTime + tz <- getCurrentTimeZone + let localTime = utcToLocalTime tz now + todayStr = formatTime defaultTimeLocale "%Y-%m-%d" localTime + dateStr = maybe todayStr Text.unpack maybeDate + transaction = + Text.unlines + [ "", + Text.pack dateStr <> " " <> description, + " " <> toAccount <> " " <> amount, + " " <> fromAccount + ] + exists <- doesFileExist transactionsFile + unless exists <| do + TextIO.writeFile transactionsFile "; Transactions added via Telegram bot\n" + TextIO.appendFile transactionsFile transaction + pure (Right ("Transaction added:\n" <> transaction)) + +hledgerBalanceTool :: Engine.Tool +hledgerBalanceTool = + Engine.Tool + { Engine.toolName = "hledger_balance", + Engine.toolDescription = + "Query account balances from hledger. " + <> "Account patterns: 'as' (assets), 'li' (liabilities), 'ex' (expenses), 'in' (income), 'eq' (equity). " + <> "Can drill down like 'as:me:cash' or 'ex:us:need'. " + <> "Currency defaults to USD but can be changed (e.g., 'BTC', 'ETH'). " + <> "Period uses hledger syntax: 'thismonth', 'lastmonth', 'thisyear', '2024', '2024-06', " + <> "'from 2024-01-01 to 2024-06-30', 'from 2024-06-01'.", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "account_pattern" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Account pattern to filter (e.g., 'as:me:cash', 'ex', 'li:us:cred')" :: Text) + ], + "period" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("hledger period: 'thismonth', 'lastmonth', '2024', '2024-06', 'from 2024-01-01 to 2024-06-30'" :: Text) + ], + "currency" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Currency to display values in (default: 'USD'). Examples: 'BTC', 'ETH', 'EUR'" :: Text) + ] + ], + "required" .= ([] :: [Text]) + ], + Engine.toolExecute = executeBalance + } + +executeBalance :: Aeson.Value -> IO Aeson.Value +executeBalance v = + case Aeson.fromJSON v of + Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e]) + Aeson.Success (args :: BalanceArgs) -> do + result <- queryBalance (baPattern args) (baPeriod args) (baCurrency args) + case result of + Left err -> pure (Aeson.object ["error" .= err]) + Right output -> + pure + ( Aeson.object + [ "success" .= True, + "balances" .= output + ] + ) + +data BalanceArgs = BalanceArgs + { baPattern :: Maybe Text, + baPeriod :: Maybe Text, + baCurrency :: Maybe Text + } + deriving (Generic) + +instance Aeson.FromJSON BalanceArgs where + parseJSON = + Aeson.withObject "BalanceArgs" <| \v -> + (BalanceArgs </ (v .:? "account_pattern")) + <*> (v .:? "period") + <*> (v .:? "currency") + +hledgerRegisterTool :: Engine.Tool +hledgerRegisterTool = + Engine.Tool + { Engine.toolName = "hledger_register", + Engine.toolDescription = + "Show recent transactions for an account. " + <> "Useful for seeing transaction history and checking recent spending. " + <> "Currency defaults to USD.", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "account_pattern" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Account pattern to show transactions for (e.g., 'ex:us:need:grocery')" :: Text) + ], + "limit" + .= Aeson.object + [ "type" .= ("integer" :: Text), + "description" .= ("Max transactions to show (default: 10)" :: Text) + ], + "currency" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Currency to display values in (default: 'USD')" :: Text) + ], + "period" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("hledger period: 'thismonth', 'lastmonth', '2024', '2024-06', 'from 2024-06-01 to 2024-12-31'" :: Text) + ] + ], + "required" .= (["account_pattern"] :: [Text]) + ], + Engine.toolExecute = executeRegister + } + +executeRegister :: Aeson.Value -> IO Aeson.Value +executeRegister v = + case Aeson.fromJSON v of + Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e]) + Aeson.Success (args :: RegisterArgs) -> do + result <- queryRegister (raPattern args) (raLimit args) (raCurrency args) (raPeriod args) + case result of + Left err -> pure (Aeson.object ["error" .= err]) + Right output -> + pure + ( Aeson.object + [ "success" .= True, + "transactions" .= output + ] + ) + +data RegisterArgs = RegisterArgs + { raPattern :: Text, + raLimit :: Maybe Int, + raCurrency :: Maybe Text, + raPeriod :: Maybe Text + } + deriving (Generic) + +instance Aeson.FromJSON RegisterArgs where + parseJSON = + Aeson.withObject "RegisterArgs" <| \v -> + (RegisterArgs </ (v .: "account_pattern")) + <*> (v .:? "limit") + <*> (v .:? "currency") + <*> (v .:? "period") + +hledgerAddTool :: Engine.Tool +hledgerAddTool = + Engine.Tool + { Engine.toolName = "hledger_add", + Engine.toolDescription = + "Add a new transaction to the ledger. " + <> "Use for recording expenses like 'I spent $30 at the barber'. " + <> "Account naming: ex:me:want (personal discretionary), ex:us:need (shared necessities), " + <> "as:me:cash:checking (bank account), li:us:cred:chase (credit card). " + <> "Common expense accounts: ex:us:need:grocery, ex:us:need:utilities, ex:me:want:dining, ex:me:want:grooming.", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "description" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Transaction description (e.g., 'Haircut at Joe's Barber')" :: Text) + ], + "from_account" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Account paying (e.g., 'as:me:cash:checking', 'li:us:cred:chase')" :: Text) + ], + "to_account" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Account receiving (e.g., 'ex:me:want:grooming')" :: Text) + ], + "amount" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Amount with currency (e.g., '$30.00', '30 USD')" :: Text) + ], + "date" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Transaction date YYYY-MM-DD (default: today)" :: Text) + ] + ], + "required" .= (["description", "from_account", "to_account", "amount"] :: [Text]) + ], + Engine.toolExecute = executeAdd + } + +executeAdd :: Aeson.Value -> IO Aeson.Value +executeAdd v = + case Aeson.fromJSON v of + Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e]) + Aeson.Success (args :: AddArgs) -> do + result <- + addTransaction + (aaDescription args) + (aaFromAccount args) + (aaToAccount args) + (aaAmount args) + (aaDate args) + case result of + Left err -> pure (Aeson.object ["error" .= err]) + Right msg -> + pure + ( Aeson.object + [ "success" .= True, + "message" .= msg + ] + ) + +data AddArgs = AddArgs + { aaDescription :: Text, + aaFromAccount :: Text, + aaToAccount :: Text, + aaAmount :: Text, + aaDate :: Maybe Text + } + deriving (Generic) + +instance Aeson.FromJSON AddArgs where + parseJSON = + Aeson.withObject "AddArgs" <| \v -> + (AddArgs </ (v .: "description")) + <*> (v .: "from_account") + <*> (v .: "to_account") + <*> (v .: "amount") + <*> (v .:? "date") + +hledgerIncomeStatementTool :: Engine.Tool +hledgerIncomeStatementTool = + Engine.Tool + { Engine.toolName = "hledger_income_statement", + Engine.toolDescription = + "Show income statement (income vs expenses) for a period. " + <> "Good for seeing 'how much did I spend this month' or 'what's my net income'. " + <> "Currency defaults to USD. " + <> "Period uses hledger syntax: 'thismonth', 'lastmonth', '2024', 'from 2024-01-01 to 2024-06-30'.", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "period" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("hledger period (default: 'thismonth'): 'lastmonth', '2024', '2024-06', 'from 2024-01-01 to 2024-06-30'" :: Text) + ], + "currency" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Currency to display values in (default: 'USD')" :: Text) + ] + ], + "required" .= ([] :: [Text]) + ], + Engine.toolExecute = executeIncomeStatement + } + +executeIncomeStatement :: Aeson.Value -> IO Aeson.Value +executeIncomeStatement v = + case Aeson.fromJSON v of + Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e]) + Aeson.Success (args :: IncomeStatementArgs) -> do + result <- incomeStatement (isaPeriod args) (isaCurrency args) + case result of + Left err -> pure (Aeson.object ["error" .= err]) + Right output -> + pure + ( Aeson.object + [ "success" .= True, + "income_statement" .= output + ] + ) + +data IncomeStatementArgs = IncomeStatementArgs + { isaPeriod :: Maybe Text, + isaCurrency :: Maybe Text + } + deriving (Generic) + +instance Aeson.FromJSON IncomeStatementArgs where + parseJSON = + Aeson.withObject "IncomeStatementArgs" <| \v -> + (IncomeStatementArgs </ (v .:? "period")) + <*> (v .:? "currency") + +hledgerBalanceSheetTool :: Engine.Tool +hledgerBalanceSheetTool = + Engine.Tool + { Engine.toolName = "hledger_balance_sheet", + Engine.toolDescription = + "Show current balance sheet (assets, liabilities, net worth). " + <> "Good for seeing 'what's my net worth' or 'how much do I have'. " + <> "Currency defaults to USD.", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "currency" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Currency to display values in (default: 'USD')" :: Text) + ] + ], + "required" .= ([] :: [Text]) + ], + Engine.toolExecute = executeBalanceSheet + } + +executeBalanceSheet :: Aeson.Value -> IO Aeson.Value +executeBalanceSheet v = + case Aeson.fromJSON v of + Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e]) + Aeson.Success (args :: BalanceSheetArgs) -> do + result <- balanceSheet (bsCurrency args) + case result of + Left err -> pure (Aeson.object ["error" .= err]) + Right output -> + pure + ( Aeson.object + [ "success" .= True, + "balance_sheet" .= output + ] + ) + +newtype BalanceSheetArgs = BalanceSheetArgs + { bsCurrency :: Maybe Text + } + deriving (Generic) + +instance Aeson.FromJSON BalanceSheetArgs where + parseJSON = + Aeson.withObject "BalanceSheetArgs" <| \v -> + BalanceSheetArgs </ (v .:? "currency") diff --git a/Omni/Agent/Tools/Http.hs b/Omni/Agent/Tools/Http.hs new file mode 100644 index 0000000..d996ff5 --- /dev/null +++ b/Omni/Agent/Tools/Http.hs @@ -0,0 +1,338 @@ +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE NoImplicitPrelude #-} + +-- | HTTP request tools for agent API interactions. +-- +-- Provides http_get and http_post tools for making HTTP requests. +-- Supports headers, query params, and JSON body. +-- +-- : out omni-agent-tools-http +-- : dep aeson +-- : dep http-conduit +module Omni.Agent.Tools.Http + ( -- * Tools + httpGetTool, + httpPostTool, + allHttpTools, + + -- * Types + HttpGetArgs (..), + HttpPostArgs (..), + HttpResult (..), + + -- * Testing + main, + test, + ) +where + +import Alpha +import Data.Aeson ((.:), (.:?), (.=)) +import qualified Data.Aeson as Aeson +import qualified Data.Aeson.Key as Key +import qualified Data.Aeson.KeyMap as KeyMap +import qualified Data.ByteString.Lazy as BL +import qualified Data.CaseInsensitive as CI +import qualified Data.Text as Text +import qualified Data.Text.Encoding as TE +import qualified Network.HTTP.Client as HTTPClient +import qualified Network.HTTP.Simple as HTTP +import qualified Omni.Agent.Engine as Engine +import qualified Omni.Test as Test +import System.Timeout (timeout) + +main :: IO () +main = Test.run test + +test :: Test.Tree +test = + Test.group + "Omni.Agent.Tools.Http" + [ Test.unit "httpGetTool has correct name" <| do + Engine.toolName httpGetTool Test.@=? "http_get", + Test.unit "httpPostTool has correct name" <| do + Engine.toolName httpPostTool Test.@=? "http_post", + Test.unit "allHttpTools has 2 tools" <| do + length allHttpTools Test.@=? 2, + Test.unit "HttpGetArgs parses correctly" <| do + let json = Aeson.object ["url" .= ("https://example.com" :: Text)] + case Aeson.fromJSON json of + Aeson.Success (args :: HttpGetArgs) -> httpGetUrl args Test.@=? "https://example.com" + Aeson.Error e -> Test.assertFailure e, + Test.unit "HttpGetArgs parses with headers" <| do + let json = + Aeson.object + [ "url" .= ("https://api.example.com" :: Text), + "headers" .= Aeson.object ["Authorization" .= ("Bearer token" :: Text)] + ] + case Aeson.fromJSON json of + Aeson.Success (args :: HttpGetArgs) -> do + httpGetUrl args Test.@=? "https://api.example.com" + isJust (httpGetHeaders args) Test.@=? True + Aeson.Error e -> Test.assertFailure e, + Test.unit "HttpPostArgs parses correctly" <| do + let json = + Aeson.object + [ "url" .= ("https://api.example.com" :: Text), + "body" .= Aeson.object ["key" .= ("value" :: Text)] + ] + case Aeson.fromJSON json of + Aeson.Success (args :: HttpPostArgs) -> do + httpPostUrl args Test.@=? "https://api.example.com" + isJust (httpPostBody args) Test.@=? True + Aeson.Error e -> Test.assertFailure e, + Test.unit "HttpResult JSON roundtrip" <| do + let result = + HttpResult + { httpResultStatus = 200, + httpResultHeaders = Aeson.object ["Content-Type" .= ("application/json" :: Text)], + httpResultBody = "{\"ok\": true}" + } + case Aeson.decode (Aeson.encode result) of + Nothing -> Test.assertFailure "Failed to decode HttpResult" + Just decoded -> httpResultStatus decoded Test.@=? 200, + Test.unit "http_get fetches real URL" <| do + let args = Aeson.object ["url" .= ("https://httpbin.org/get" :: Text)] + result <- Engine.toolExecute httpGetTool args + case Aeson.fromJSON result of + Aeson.Success (r :: HttpResult) -> do + httpResultStatus r Test.@=? 200 + ("httpbin.org" `Text.isInfixOf` httpResultBody r) Test.@=? True + Aeson.Error e -> Test.assertFailure e, + Test.unit "http_post with JSON body" <| do + let args = + Aeson.object + [ "url" .= ("https://httpbin.org/post" :: Text), + "body" .= Aeson.object ["test" .= ("value" :: Text)] + ] + result <- Engine.toolExecute httpPostTool args + case Aeson.fromJSON result of + Aeson.Success (r :: HttpResult) -> do + httpResultStatus r Test.@=? 200 + ("test" `Text.isInfixOf` httpResultBody r) Test.@=? True + Aeson.Error e -> Test.assertFailure e + ] + +data HttpGetArgs = HttpGetArgs + { httpGetUrl :: Text, + httpGetHeaders :: Maybe Aeson.Object, + httpGetParams :: Maybe Aeson.Object + } + deriving (Show, Eq, Generic) + +instance Aeson.FromJSON HttpGetArgs where + parseJSON = + Aeson.withObject "HttpGetArgs" <| \v -> + (HttpGetArgs </ (v .: "url")) + <*> (v .:? "headers") + <*> (v .:? "params") + +data HttpPostArgs = HttpPostArgs + { httpPostUrl :: Text, + httpPostHeaders :: Maybe Aeson.Object, + httpPostBody :: Maybe Aeson.Value, + httpPostContentType :: Maybe Text + } + deriving (Show, Eq, Generic) + +instance Aeson.FromJSON HttpPostArgs where + parseJSON = + Aeson.withObject "HttpPostArgs" <| \v -> + (HttpPostArgs </ (v .: "url")) + <*> (v .:? "headers") + <*> (v .:? "body") + <*> (v .:? "content_type") + +data HttpResult = HttpResult + { httpResultStatus :: Int, + httpResultHeaders :: Aeson.Value, + httpResultBody :: Text + } + deriving (Show, Eq, Generic) + +instance Aeson.ToJSON HttpResult where + toJSON r = + Aeson.object + [ "status" .= httpResultStatus r, + "headers" .= httpResultHeaders r, + "body" .= httpResultBody r + ] + +instance Aeson.FromJSON HttpResult where + parseJSON = + Aeson.withObject "HttpResult" <| \v -> + (HttpResult </ (v .: "status")) + <*> (v .: "headers") + <*> (v .: "body") + +allHttpTools :: [Engine.Tool] +allHttpTools = [httpGetTool, httpPostTool] + +httpGetTool :: Engine.Tool +httpGetTool = + Engine.Tool + { Engine.toolName = "http_get", + Engine.toolDescription = + "Make an HTTP GET request. Returns status code, headers, and response body. " + <> "Use for fetching data from APIs or web pages.", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "url" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("The URL to request" :: Text) + ], + "headers" + .= Aeson.object + [ "type" .= ("object" :: Text), + "description" .= ("Optional headers as key-value pairs" :: Text) + ], + "params" + .= Aeson.object + [ "type" .= ("object" :: Text), + "description" .= ("Optional query parameters as key-value pairs" :: Text) + ] + ], + "required" .= (["url"] :: [Text]) + ], + Engine.toolExecute = executeHttpGet + } + +executeHttpGet :: Aeson.Value -> IO Aeson.Value +executeHttpGet v = + case Aeson.fromJSON v of + Aeson.Error e -> pure <| mkError ("Invalid arguments: " <> Text.pack e) + Aeson.Success args -> do + let urlWithParams = case httpGetParams args of + Nothing -> httpGetUrl args + Just params -> + let paramList = [(k, v') | (k, v') <- KeyMap.toList params] + paramStr = Text.intercalate "&" [Key.toText k <> "=" <> valueToText v' | (k, v') <- paramList] + in if Text.null paramStr + then httpGetUrl args + else httpGetUrl args <> "?" <> paramStr + doHttpRequest "GET" urlWithParams (httpGetHeaders args) Nothing + +httpPostTool :: Engine.Tool +httpPostTool = + Engine.Tool + { Engine.toolName = "http_post", + Engine.toolDescription = + "Make an HTTP POST request. Returns status code, headers, and response body. " + <> "Use for submitting data to APIs or forms.", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "url" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("The URL to request" :: Text) + ], + "headers" + .= Aeson.object + [ "type" .= ("object" :: Text), + "description" .= ("Optional headers as key-value pairs" :: Text) + ], + "body" + .= Aeson.object + [ "type" .= ("object" :: Text), + "description" .= ("Optional JSON body (object or string)" :: Text) + ], + "content_type" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Content type (default: application/json)" :: Text) + ] + ], + "required" .= (["url"] :: [Text]) + ], + Engine.toolExecute = executeHttpPost + } + +executeHttpPost :: Aeson.Value -> IO Aeson.Value +executeHttpPost v = + case Aeson.fromJSON v of + Aeson.Error e -> pure <| mkError ("Invalid arguments: " <> Text.pack e) + Aeson.Success args -> do + let contentType = fromMaybe "application/json" (httpPostContentType args) + body = case httpPostBody args of + Nothing -> Nothing + Just b -> Just (contentType, BL.toStrict (Aeson.encode b)) + doHttpRequest "POST" (httpPostUrl args) (httpPostHeaders args) body + +doHttpRequest :: + ByteString -> + Text -> + Maybe Aeson.Object -> + Maybe (Text, ByteString) -> + IO Aeson.Value +doHttpRequest method url mHeaders mBody = do + let timeoutMicros = 30 * 1000000 + result <- + try <| do + req0 <- HTTP.parseRequest (Text.unpack url) + let req1 = + HTTP.setRequestMethod method + <| HTTP.setRequestHeader "User-Agent" ["OmniAgent/1.0"] + <| HTTP.setRequestResponseTimeout (HTTPClient.responseTimeoutMicro timeoutMicros) + <| req0 + req2 = case mHeaders of + Nothing -> req1 + Just hdrs -> foldr addHeader req1 (KeyMap.toList hdrs) + req3 = case mBody of + Nothing -> req2 + Just (ct, bodyBytes) -> + HTTP.setRequestHeader "Content-Type" [TE.encodeUtf8 ct] + <| HTTP.setRequestBodyLBS (BL.fromStrict bodyBytes) + <| req2 + mResp <- timeout timeoutMicros (HTTP.httpLBS req3) + case mResp of + Nothing -> pure (Left "Request timed out after 30 seconds") + Just resp -> pure (Right resp) + case result of + Left (e :: SomeException) -> pure <| mkError ("Request failed: " <> tshow e) + Right (Left err) -> pure <| mkError err + Right (Right response) -> do + let status = HTTP.getResponseStatusCode response + respHeaders = HTTP.getResponseHeaders response + headerObj = + Aeson.object + [ Key.fromText (TE.decodeUtf8 (CI.original k)) .= TE.decodeUtf8 v + | (k, v) <- respHeaders + ] + body = TE.decodeUtf8With (\_ _ -> Just '?') (BL.toStrict (HTTP.getResponseBody response)) + pure + <| Aeson.toJSON + <| HttpResult + { httpResultStatus = status, + httpResultHeaders = headerObj, + httpResultBody = body + } + where + addHeader :: (Aeson.Key, Aeson.Value) -> HTTP.Request -> HTTP.Request + addHeader (k, v) req = + let headerName = CI.mk (TE.encodeUtf8 (Key.toText k)) + headerValue = TE.encodeUtf8 (valueToText v) + in HTTP.addRequestHeader headerName headerValue req + +valueToText :: Aeson.Value -> Text +valueToText (Aeson.String s) = s +valueToText (Aeson.Number n) = tshow n +valueToText (Aeson.Bool b) = if b then "true" else "false" +valueToText Aeson.Null = "" +valueToText other = TE.decodeUtf8 (BL.toStrict (Aeson.encode other)) + +mkError :: Text -> Aeson.Value +mkError err = + Aeson.object + [ "status" .= (-1 :: Int), + "headers" .= Aeson.object [], + "body" .= err + ] diff --git a/Omni/Agent/Tools/Notes.hs b/Omni/Agent/Tools/Notes.hs new file mode 100644 index 0000000..e3cef5d --- /dev/null +++ b/Omni/Agent/Tools/Notes.hs @@ -0,0 +1,357 @@ +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE NoImplicitPrelude #-} + +-- | Quick notes tool for agents. +-- +-- Provides simple CRUD for tagged notes stored in memory.db. +-- +-- : out omni-agent-tools-notes +-- : dep aeson +-- : dep sqlite-simple +module Omni.Agent.Tools.Notes + ( -- * Tools + noteAddTool, + noteListTool, + noteDeleteTool, + + -- * Direct API + Note (..), + createNote, + listNotes, + listNotesByTopic, + deleteNote, + + -- * Database + initNotesTable, + + -- * Testing + main, + test, + ) +where + +import Alpha +import Data.Aeson ((.!=), (.:), (.:?), (.=)) +import qualified Data.Aeson as Aeson +import qualified Data.Text as Text +import Data.Time (UTCTime, getCurrentTime) +import qualified Database.SQLite.Simple as SQL +import qualified Omni.Agent.Engine as Engine +import qualified Omni.Agent.Memory as Memory +import qualified Omni.Test as Test + +main :: IO () +main = Test.run test + +test :: Test.Tree +test = + Test.group + "Omni.Agent.Tools.Notes" + [ Test.unit "noteAddTool has correct schema" <| do + let tool = noteAddTool "test-user-id" + Engine.toolName tool Test.@=? "note_add", + Test.unit "noteListTool has correct schema" <| do + let tool = noteListTool "test-user-id" + Engine.toolName tool Test.@=? "note_list", + Test.unit "noteDeleteTool has correct schema" <| do + let tool = noteDeleteTool "test-user-id" + Engine.toolName tool Test.@=? "note_delete", + Test.unit "Note JSON roundtrip" <| do + now <- getCurrentTime + let n = + Note + { noteId = 1, + noteUserId = "user-123", + noteTopic = "groceries", + noteContent = "Buy milk", + noteCreatedAt = now + } + case Aeson.decode (Aeson.encode n) of + Nothing -> Test.assertFailure "Failed to decode Note" + Just decoded -> do + noteContent decoded Test.@=? "Buy milk" + noteTopic decoded Test.@=? "groceries" + ] + +data Note = Note + { noteId :: Int, + noteUserId :: Text, + noteTopic :: Text, + noteContent :: Text, + noteCreatedAt :: UTCTime + } + deriving (Show, Eq, Generic) + +instance Aeson.ToJSON Note where + toJSON n = + Aeson.object + [ "id" .= noteId n, + "user_id" .= noteUserId n, + "topic" .= noteTopic n, + "content" .= noteContent n, + "created_at" .= noteCreatedAt n + ] + +instance Aeson.FromJSON Note where + parseJSON = + Aeson.withObject "Note" <| \v -> + (Note </ (v .: "id")) + <*> (v .: "user_id") + <*> (v .: "topic") + <*> (v .: "content") + <*> (v .: "created_at") + +instance SQL.FromRow Note where + fromRow = + (Note </ SQL.field) + <*> SQL.field + <*> SQL.field + <*> SQL.field + <*> SQL.field + +initNotesTable :: SQL.Connection -> IO () +initNotesTable conn = do + SQL.execute_ + conn + "CREATE TABLE IF NOT EXISTS notes (\ + \ id INTEGER PRIMARY KEY AUTOINCREMENT,\ + \ user_id TEXT NOT NULL,\ + \ topic TEXT NOT NULL,\ + \ content TEXT NOT NULL,\ + \ created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP\ + \)" + SQL.execute_ + conn + "CREATE INDEX IF NOT EXISTS idx_notes_user ON notes(user_id)" + SQL.execute_ + conn + "CREATE INDEX IF NOT EXISTS idx_notes_topic ON notes(user_id, topic)" + +createNote :: Text -> Text -> Text -> IO Note +createNote uid topic content = do + now <- getCurrentTime + Memory.withMemoryDb <| \conn -> do + initNotesTable conn + SQL.execute + conn + "INSERT INTO notes (user_id, topic, content, created_at) VALUES (?, ?, ?, ?)" + (uid, topic, content, now) + rowId <- SQL.lastInsertRowId conn + pure + Note + { noteId = fromIntegral rowId, + noteUserId = uid, + noteTopic = topic, + noteContent = content, + noteCreatedAt = now + } + +listNotes :: Text -> Int -> IO [Note] +listNotes uid limit = + Memory.withMemoryDb <| \conn -> do + initNotesTable conn + SQL.query + conn + "SELECT id, user_id, topic, content, created_at \ + \FROM notes WHERE user_id = ? \ + \ORDER BY created_at DESC LIMIT ?" + (uid, limit) + +listNotesByTopic :: Text -> Text -> Int -> IO [Note] +listNotesByTopic uid topic limit = + Memory.withMemoryDb <| \conn -> do + initNotesTable conn + SQL.query + conn + "SELECT id, user_id, topic, content, created_at \ + \FROM notes WHERE user_id = ? AND topic = ? \ + \ORDER BY created_at DESC LIMIT ?" + (uid, topic, limit) + +deleteNote :: Text -> Int -> IO Bool +deleteNote uid nid = + Memory.withMemoryDb <| \conn -> do + initNotesTable conn + SQL.execute + conn + "DELETE FROM notes WHERE id = ? AND user_id = ?" + (nid, uid) + changes <- SQL.changes conn + pure (changes > 0) + +noteAddTool :: Text -> Engine.Tool +noteAddTool uid = + Engine.Tool + { Engine.toolName = "note_add", + Engine.toolDescription = + "Add a quick note on a topic. Use for reminders, lists, ideas, or anything " + <> "the user wants to jot down. Topics help organize notes (e.g., 'groceries', " + <> "'ideas', 'todo', 'recipes').", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "topic" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Topic/category for the note (e.g., 'groceries', 'todo')" :: Text) + ], + "content" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("The note content" :: Text) + ] + ], + "required" .= (["topic", "content"] :: [Text]) + ], + Engine.toolExecute = executeNoteAdd uid + } + +executeNoteAdd :: Text -> Aeson.Value -> IO Aeson.Value +executeNoteAdd uid v = + case Aeson.fromJSON v of + Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e]) + Aeson.Success (args :: NoteAddArgs) -> do + newNote <- createNote uid (naTopic args) (naContent args) + pure + ( Aeson.object + [ "success" .= True, + "note_id" .= noteId newNote, + "message" .= ("Added note to '" <> noteTopic newNote <> "': " <> noteContent newNote) + ] + ) + +data NoteAddArgs = NoteAddArgs + { naTopic :: Text, + naContent :: Text + } + deriving (Generic) + +instance Aeson.FromJSON NoteAddArgs where + parseJSON = + Aeson.withObject "NoteAddArgs" <| \v -> + (NoteAddArgs </ (v .: "topic")) + <*> (v .: "content") + +noteListTool :: Text -> Engine.Tool +noteListTool uid = + Engine.Tool + { Engine.toolName = "note_list", + Engine.toolDescription = + "List notes, optionally filtered by topic. Use to show the user their " + <> "saved notes or check what's on a specific list.", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "topic" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Filter by topic (optional, omit to list all)" :: Text) + ], + "limit" + .= Aeson.object + [ "type" .= ("integer" :: Text), + "description" .= ("Max notes to return (default: 20)" :: Text) + ] + ], + "required" .= ([] :: [Text]) + ], + Engine.toolExecute = executeNoteList uid + } + +executeNoteList :: Text -> Aeson.Value -> IO Aeson.Value +executeNoteList uid v = + case Aeson.fromJSON v of + Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e]) + Aeson.Success (args :: NoteListArgs) -> do + let lim = min 50 (max 1 (nlLimit args)) + notes <- case nlTopic args of + Just topic -> listNotesByTopic uid topic lim + Nothing -> listNotes uid lim + pure + ( Aeson.object + [ "success" .= True, + "count" .= length notes, + "notes" .= formatNotesForLLM notes + ] + ) + +formatNotesForLLM :: [Note] -> Text +formatNotesForLLM [] = "No notes found." +formatNotesForLLM notes = + Text.unlines (map formatNote notes) + where + formatNote n = + "[" <> noteTopic n <> "] " <> noteContent n <> " (id: " <> tshow (noteId n) <> ")" + +data NoteListArgs = NoteListArgs + { nlTopic :: Maybe Text, + nlLimit :: Int + } + deriving (Generic) + +instance Aeson.FromJSON NoteListArgs where + parseJSON = + Aeson.withObject "NoteListArgs" <| \v -> + (NoteListArgs </ (v .:? "topic")) + <*> (v .:? "limit" .!= 20) + +noteDeleteTool :: Text -> Engine.Tool +noteDeleteTool uid = + Engine.Tool + { Engine.toolName = "note_delete", + Engine.toolDescription = + "Delete a note by its ID. Use after the user says they've completed an item " + <> "or no longer need a note.", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "note_id" + .= Aeson.object + [ "type" .= ("integer" :: Text), + "description" .= ("The ID of the note to delete" :: Text) + ] + ], + "required" .= (["note_id"] :: [Text]) + ], + Engine.toolExecute = executeNoteDelete uid + } + +executeNoteDelete :: Text -> Aeson.Value -> IO Aeson.Value +executeNoteDelete uid v = + case Aeson.fromJSON v of + Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e]) + Aeson.Success (args :: NoteDeleteArgs) -> do + deleted <- deleteNote uid (ndNoteId args) + if deleted + then + pure + ( Aeson.object + [ "success" .= True, + "message" .= ("Note deleted" :: Text) + ] + ) + else + pure + ( Aeson.object + [ "success" .= False, + "error" .= ("Note not found or already deleted" :: Text) + ] + ) + +newtype NoteDeleteArgs = NoteDeleteArgs + { ndNoteId :: Int + } + deriving (Generic) + +instance Aeson.FromJSON NoteDeleteArgs where + parseJSON = + Aeson.withObject "NoteDeleteArgs" <| \v -> + NoteDeleteArgs </ (v .: "note_id") diff --git a/Omni/Agent/Tools/Outreach.hs b/Omni/Agent/Tools/Outreach.hs new file mode 100644 index 0000000..e576cbd --- /dev/null +++ b/Omni/Agent/Tools/Outreach.hs @@ -0,0 +1,513 @@ +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE NoImplicitPrelude #-} + +-- | Outreach approval queue for agent use. +-- +-- Provides tools for creating and tracking outreach drafts that require +-- human approval before sending (emails, messages, etc). +-- +-- Drafts flow: pending -> approved -> sent (or rejected) +-- +-- : out omni-agent-tools-outreach +-- : dep aeson +-- : dep uuid +-- : dep directory +module Omni.Agent.Tools.Outreach + ( -- * Tools + outreachDraftTool, + outreachListTool, + outreachStatusTool, + allOutreachTools, + + -- * Types + OutreachDraft (..), + OutreachType (..), + OutreachStatus (..), + + -- * Direct API + createDraft, + listDrafts, + getDraft, + approveDraft, + rejectDraft, + markSent, + getPendingCount, + + -- * Paths + outreachDir, + pendingDir, + approvedDir, + rejectedDir, + sentDir, + + -- * Testing + main, + test, + ) +where + +import Alpha +import Control.Monad.Fail (MonadFail (fail)) +import Data.Aeson ((.!=), (.:), (.:?), (.=)) +import qualified Data.Aeson as Aeson +import qualified Data.ByteString.Lazy as BL +import qualified Data.Text as Text +import qualified Data.Text.Encoding as TE +import qualified Data.Text.IO as TextIO +import Data.Time (UTCTime, getCurrentTime) +import qualified Data.UUID as UUID +import qualified Data.UUID.V4 as UUID +import qualified Omni.Agent.Engine as Engine +import qualified Omni.Agent.Paths as Paths +import qualified Omni.Test as Test +import qualified System.Directory as Directory +import System.FilePath ((</>)) + +main :: IO () +main = Test.run test + +test :: Test.Tree +test = + Test.group + "Omni.Agent.Tools.Outreach" + [ Test.unit "outreachDraftTool has correct name" <| do + Engine.toolName outreachDraftTool Test.@=? "outreach_draft", + Test.unit "outreachListTool has correct name" <| do + Engine.toolName outreachListTool Test.@=? "outreach_list", + Test.unit "outreachStatusTool has correct name" <| do + Engine.toolName outreachStatusTool Test.@=? "outreach_status", + Test.unit "allOutreachTools has 3 tools" <| do + length allOutreachTools Test.@=? 3, + Test.unit "OutreachDraft JSON roundtrip" <| do + now <- getCurrentTime + let draft = + OutreachDraft + { draftId = "test-id-123", + draftType = Email, + draftCreatedAt = now, + draftSubject = Just "Test subject", + draftRecipient = "test@example.com", + draftBody = "Hello, this is a test.", + draftContext = "Testing outreach system", + draftStatus = Pending, + draftRejectReason = Nothing + } + case Aeson.decode (Aeson.encode draft) of + Nothing -> Test.assertFailure "Failed to decode OutreachDraft" + Just decoded -> do + draftId decoded Test.@=? "test-id-123" + draftType decoded Test.@=? Email + draftRecipient decoded Test.@=? "test@example.com", + Test.unit "OutreachType JSON roundtrip" <| do + case Aeson.decode (Aeson.encode Email) of + Just Email -> pure () + _ -> Test.assertFailure "Failed to decode Email" + case Aeson.decode (Aeson.encode Message) of + Just Message -> pure () + _ -> Test.assertFailure "Failed to decode Message", + Test.unit "OutreachStatus JSON roundtrip" <| do + let statuses = [Pending, Approved, Rejected, Sent] + forM_ statuses <| \s -> + case Aeson.decode (Aeson.encode s) of + Nothing -> Test.assertFailure ("Failed to decode " <> show s) + Just decoded -> decoded Test.@=? s + ] + +outreachDir :: FilePath +outreachDir = Paths.outreachDir + +pendingDir :: FilePath +pendingDir = outreachDir </> "pending" + +approvedDir :: FilePath +approvedDir = outreachDir </> "approved" + +rejectedDir :: FilePath +rejectedDir = outreachDir </> "rejected" + +sentDir :: FilePath +sentDir = outreachDir </> "sent" + +data OutreachType = Email | Message + deriving (Show, Eq, Generic) + +instance Aeson.ToJSON OutreachType where + toJSON Email = Aeson.String "email" + toJSON Message = Aeson.String "message" + +instance Aeson.FromJSON OutreachType where + parseJSON = + Aeson.withText "OutreachType" <| \t -> + case Text.toLower t of + "email" -> pure Email + "message" -> pure Message + _ -> fail "OutreachType must be 'email' or 'message'" + +data OutreachStatus = Pending | Approved | Rejected | Sent + deriving (Show, Eq, Generic) + +instance Aeson.ToJSON OutreachStatus where + toJSON Pending = Aeson.String "pending" + toJSON Approved = Aeson.String "approved" + toJSON Rejected = Aeson.String "rejected" + toJSON Sent = Aeson.String "sent" + +instance Aeson.FromJSON OutreachStatus where + parseJSON = + Aeson.withText "OutreachStatus" <| \t -> + case Text.toLower t of + "pending" -> pure Pending + "approved" -> pure Approved + "rejected" -> pure Rejected + "sent" -> pure Sent + _ -> fail "OutreachStatus must be 'pending', 'approved', 'rejected', or 'sent'" + +data OutreachDraft = OutreachDraft + { draftId :: Text, + draftType :: OutreachType, + draftCreatedAt :: UTCTime, + draftSubject :: Maybe Text, + draftRecipient :: Text, + draftBody :: Text, + draftContext :: Text, + draftStatus :: OutreachStatus, + draftRejectReason :: Maybe Text + } + deriving (Show, Eq, Generic) + +instance Aeson.ToJSON OutreachDraft where + toJSON d = + Aeson.object + [ "id" .= draftId d, + "type" .= draftType d, + "created_at" .= draftCreatedAt d, + "subject" .= draftSubject d, + "recipient" .= draftRecipient d, + "body" .= draftBody d, + "context" .= draftContext d, + "status" .= draftStatus d, + "reject_reason" .= draftRejectReason d + ] + +instance Aeson.FromJSON OutreachDraft where + parseJSON = + Aeson.withObject "OutreachDraft" <| \v -> + (OutreachDraft </ (v .: "id")) + <*> (v .: "type") + <*> (v .: "created_at") + <*> (v .:? "subject") + <*> (v .: "recipient") + <*> (v .: "body") + <*> (v .: "context") + <*> (v .: "status") + <*> (v .:? "reject_reason") + +ensureDirs :: IO () +ensureDirs = do + Directory.createDirectoryIfMissing True pendingDir + Directory.createDirectoryIfMissing True approvedDir + Directory.createDirectoryIfMissing True rejectedDir + Directory.createDirectoryIfMissing True sentDir + +draftPath :: FilePath -> Text -> FilePath +draftPath dir draftId' = dir </> (Text.unpack draftId' <> ".json") + +saveDraft :: OutreachDraft -> IO () +saveDraft draft = do + ensureDirs + let dir = case draftStatus draft of + Pending -> pendingDir + Approved -> approvedDir + Rejected -> rejectedDir + Sent -> sentDir + path = draftPath dir (draftId draft) + TextIO.writeFile path (TE.decodeUtf8 (BL.toStrict (Aeson.encode draft))) + +createDraft :: OutreachType -> Text -> Maybe Text -> Text -> Text -> IO OutreachDraft +createDraft otype recipient subject body context = do + uuid <- UUID.nextRandom + now <- getCurrentTime + let draft = + OutreachDraft + { draftId = UUID.toText uuid, + draftType = otype, + draftCreatedAt = now, + draftSubject = subject, + draftRecipient = recipient, + draftBody = body, + draftContext = context, + draftStatus = Pending, + draftRejectReason = Nothing + } + saveDraft draft + pure draft + +listDrafts :: OutreachStatus -> IO [OutreachDraft] +listDrafts status = do + ensureDirs + let dir = case status of + Pending -> pendingDir + Approved -> approvedDir + Rejected -> rejectedDir + Sent -> sentDir + files <- Directory.listDirectory dir + let jsonFiles = filter (".json" `isSuffixOf`) files + drafts <- + forM jsonFiles <| \f -> do + content <- TextIO.readFile (dir </> f) + pure (Aeson.decode (BL.fromStrict (TE.encodeUtf8 content))) + pure (catMaybes drafts) + +getDraft :: Text -> IO (Maybe OutreachDraft) +getDraft draftId' = do + ensureDirs + let dirs = [pendingDir, approvedDir, rejectedDir, sentDir] + findFirst dirs + where + findFirst [] = pure Nothing + findFirst (dir : rest) = do + let path = draftPath dir draftId' + exists <- Directory.doesFileExist path + if exists + then do + content <- TextIO.readFile path + pure (Aeson.decode (BL.fromStrict (TE.encodeUtf8 content))) + else findFirst rest + +moveDraft :: Text -> OutreachStatus -> OutreachStatus -> Maybe Text -> IO (Either Text OutreachDraft) +moveDraft draftId' fromStatus toStatus reason = do + ensureDirs + let fromDir = case fromStatus of + Pending -> pendingDir + Approved -> approvedDir + Rejected -> rejectedDir + Sent -> sentDir + fromPath = draftPath fromDir draftId' + exists <- Directory.doesFileExist fromPath + if not exists + then pure (Left ("Draft not found in " <> tshow fromStatus <> " queue")) + else do + content <- TextIO.readFile fromPath + case Aeson.decode (BL.fromStrict (TE.encodeUtf8 content)) of + Nothing -> pure (Left "Failed to parse draft") + Just draft -> do + let updated = draft {draftStatus = toStatus, draftRejectReason = reason} + Directory.removeFile fromPath + saveDraft updated + pure (Right updated) + +approveDraft :: Text -> IO (Either Text OutreachDraft) +approveDraft draftId' = moveDraft draftId' Pending Approved Nothing + +rejectDraft :: Text -> Maybe Text -> IO (Either Text OutreachDraft) +rejectDraft draftId' = moveDraft draftId' Pending Rejected + +markSent :: Text -> IO (Either Text OutreachDraft) +markSent draftId' = moveDraft draftId' Approved Sent Nothing + +getPendingCount :: IO Int +getPendingCount = do + ensureDirs + files <- Directory.listDirectory pendingDir + pure (length (filter (".json" `isSuffixOf`) files)) + +allOutreachTools :: [Engine.Tool] +allOutreachTools = + [ outreachDraftTool, + outreachListTool, + outreachStatusTool + ] + +outreachDraftTool :: Engine.Tool +outreachDraftTool = + Engine.Tool + { Engine.toolName = "outreach_draft", + Engine.toolDescription = + "Create a new outreach draft for Ben to review before sending. " + <> "Use this when you want to send an email or message on behalf of the business. " + <> "All outreach requires approval before it goes out.", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "type" + .= Aeson.object + [ "type" .= ("string" :: Text), + "enum" .= (["email", "message"] :: [Text]), + "description" .= ("Type of outreach: 'email' or 'message'" :: Text) + ], + "recipient" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Email address or identifier of the recipient" :: Text) + ], + "subject" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Subject line (required for emails)" :: Text) + ], + "body" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("The message content" :: Text) + ], + "context" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Explain why you're sending this - helps Ben review" :: Text) + ] + ], + "required" .= (["type", "recipient", "body", "context"] :: [Text]) + ], + Engine.toolExecute = executeOutreachDraft + } + +executeOutreachDraft :: Aeson.Value -> IO Aeson.Value +executeOutreachDraft v = + case Aeson.fromJSON v of + Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e]) + Aeson.Success (args :: DraftArgs) -> do + let otype = case daType args of + "email" -> Email + _ -> Message + draft <- createDraft otype (daRecipient args) (daSubject args) (daBody args) (daContext args) + pure + ( Aeson.object + [ "success" .= True, + "draft_id" .= draftId draft, + "message" .= ("Draft created and queued for review. ID: " <> draftId draft) + ] + ) + +data DraftArgs = DraftArgs + { daType :: Text, + daRecipient :: Text, + daSubject :: Maybe Text, + daBody :: Text, + daContext :: Text + } + deriving (Generic) + +instance Aeson.FromJSON DraftArgs where + parseJSON = + Aeson.withObject "DraftArgs" <| \v -> + (DraftArgs </ (v .: "type")) + <*> (v .: "recipient") + <*> (v .:? "subject") + <*> (v .: "body") + <*> (v .: "context") + +outreachListTool :: Engine.Tool +outreachListTool = + Engine.Tool + { Engine.toolName = "outreach_list", + Engine.toolDescription = + "List outreach drafts by status. Use to check what's pending approval, " + <> "what's been approved, or review past outreach.", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "status" + .= Aeson.object + [ "type" .= ("string" :: Text), + "enum" .= (["pending", "approved", "rejected", "sent"] :: [Text]), + "description" .= ("Filter by status (default: pending)" :: Text) + ], + "limit" + .= Aeson.object + [ "type" .= ("integer" :: Text), + "description" .= ("Max drafts to return (default: 20)" :: Text) + ] + ], + "required" .= ([] :: [Text]) + ], + Engine.toolExecute = executeOutreachList + } + +executeOutreachList :: Aeson.Value -> IO Aeson.Value +executeOutreachList v = + case Aeson.fromJSON v of + Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e]) + Aeson.Success (args :: ListArgs) -> do + let status = case laStatus args of + Just "approved" -> Approved + Just "rejected" -> Rejected + Just "sent" -> Sent + _ -> Pending + limit = min 50 (max 1 (laLimit args)) + drafts <- listDrafts status + let limited = take limit drafts + pure + ( Aeson.object + [ "success" .= True, + "status" .= tshow status, + "count" .= length limited, + "drafts" .= limited + ] + ) + +data ListArgs = ListArgs + { laStatus :: Maybe Text, + laLimit :: Int + } + deriving (Generic) + +instance Aeson.FromJSON ListArgs where + parseJSON = + Aeson.withObject "ListArgs" <| \v -> + (ListArgs </ (v .:? "status")) + <*> (v .:? "limit" .!= 20) + +outreachStatusTool :: Engine.Tool +outreachStatusTool = + Engine.Tool + { Engine.toolName = "outreach_status", + Engine.toolDescription = + "Check the status of a specific outreach draft by ID.", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "draft_id" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("The draft ID to check" :: Text) + ] + ], + "required" .= (["draft_id"] :: [Text]) + ], + Engine.toolExecute = executeOutreachStatus + } + +executeOutreachStatus :: Aeson.Value -> IO Aeson.Value +executeOutreachStatus v = + case Aeson.fromJSON v of + Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e]) + Aeson.Success (args :: StatusArgs) -> do + mDraft <- getDraft (saId args) + case mDraft of + Nothing -> + pure (Aeson.object ["error" .= ("Draft not found" :: Text)]) + Just draft -> + pure + ( Aeson.object + [ "success" .= True, + "draft" .= draft + ] + ) + +newtype StatusArgs = StatusArgs + { saId :: Text + } + deriving (Generic) + +instance Aeson.FromJSON StatusArgs where + parseJSON = + Aeson.withObject "StatusArgs" <| \v -> + StatusArgs </ (v .: "draft_id") diff --git a/Omni/Agent/Tools/Pdf.hs b/Omni/Agent/Tools/Pdf.hs new file mode 100644 index 0000000..7687234 --- /dev/null +++ b/Omni/Agent/Tools/Pdf.hs @@ -0,0 +1,180 @@ +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE NoImplicitPrelude #-} + +-- | PDF extraction tool using poppler-utils (pdftotext). +-- +-- Extracts text from PDF files for LLM consumption. +-- +-- : out omni-agent-tools-pdf +-- : dep aeson +-- : dep http-conduit +-- : dep directory +-- : dep process +module Omni.Agent.Tools.Pdf + ( -- * Tool + pdfTool, + + -- * Direct API + extractPdfText, + downloadAndExtract, + + -- * Testing + main, + test, + ) +where + +import Alpha +import Data.Aeson ((.=)) +import qualified Data.Aeson as Aeson +import qualified Data.ByteString.Lazy as BL +import qualified Data.Text as Text +import qualified Network.HTTP.Simple as HTTP +import qualified Omni.Agent.Engine as Engine +import qualified Omni.Test as Test +import System.IO (hClose) +import System.IO.Temp (withSystemTempFile) +import System.Process (readProcessWithExitCode) + +main :: IO () +main = Test.run test + +test :: Test.Tree +test = + Test.group + "Omni.Agent.Tools.Pdf" + [ Test.unit "pdfTool has correct schema" <| do + let tool = pdfTool + Engine.toolName tool Test.@=? "read_pdf", + Test.unit "extractPdfText handles missing file" <| do + result <- extractPdfText "/nonexistent/file.pdf" + case result of + Left err -> ("No such file" `Text.isInfixOf` err || "pdftotext" `Text.isInfixOf` err) Test.@=? True + Right _ -> Test.assertFailure "Expected error for missing file", + Test.unit "chunkText splits correctly" <| do + let text = Text.replicate 5000 "a" + chunks = chunkText 1000 text + length chunks Test.@=? 5 + all (\c -> Text.length c <= 1000) chunks Test.@=? True, + Test.unit "chunkText handles small text" <| do + let text = "small text" + chunks = chunkText 1000 text + chunks Test.@=? ["small text"] + ] + +data PdfArgs = PdfArgs + { pdfPath :: Text, + pdfMaxChars :: Maybe Int + } + deriving (Generic) + +instance Aeson.FromJSON PdfArgs where + parseJSON = + Aeson.withObject "PdfArgs" <| \v -> + (PdfArgs </ (v Aeson..: "path")) + <*> (v Aeson..:? "max_chars") + +pdfTool :: Engine.Tool +pdfTool = + Engine.Tool + { Engine.toolName = "read_pdf", + Engine.toolDescription = + "Extract text from a PDF file. Use this when you receive a PDF document " + <> "and need to read its contents. Returns the extracted text.", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "path" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Path to the PDF file" :: Text) + ], + "max_chars" + .= Aeson.object + [ "type" .= ("integer" :: Text), + "description" .= ("Maximum characters to return (default: 50000)" :: Text) + ] + ], + "required" .= (["path"] :: [Text]) + ], + Engine.toolExecute = executePdf + } + +executePdf :: Aeson.Value -> IO Aeson.Value +executePdf v = + case Aeson.fromJSON v of + Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e]) + Aeson.Success (args :: PdfArgs) -> do + let maxChars = maybe 50000 (min 100000 <. max 1000) (pdfMaxChars args) + result <- extractPdfText (Text.unpack (pdfPath args)) + case result of + Left err -> + pure (Aeson.object ["error" .= err]) + Right text -> do + let truncated = Text.take maxChars text + wasTruncated = Text.length text > maxChars + pure + ( Aeson.object + [ "success" .= True, + "text" .= truncated, + "chars" .= Text.length truncated, + "truncated" .= wasTruncated + ] + ) + +extractPdfText :: FilePath -> IO (Either Text Text) +extractPdfText path = do + result <- + try <| readProcessWithExitCode "pdftotext" ["-layout", path, "-"] "" + case result of + Left (e :: SomeException) -> + pure (Left ("pdftotext error: " <> tshow e)) + Right (exitCode, stdoutStr, stderrStr) -> + case exitCode of + ExitSuccess -> pure (Right (Text.pack stdoutStr)) + ExitFailure code -> + pure (Left ("pdftotext failed (" <> tshow code <> "): " <> Text.pack stderrStr)) + +downloadAndExtract :: Text -> Text -> Text -> IO (Either Text Text) +downloadAndExtract botToken filePath maxCharsText = do + let url = + "https://api.telegram.org/file/bot" + <> Text.unpack botToken + <> "/" + <> Text.unpack filePath + maxChars = maybe 50000 identity (readMaybe (Text.unpack maxCharsText) :: Maybe Int) + withSystemTempFile "telegram_pdf.pdf" <| \tmpPath tmpHandle -> do + hClose tmpHandle + downloadResult <- + try <| do + req <- HTTP.parseRequest url + response <- HTTP.httpLBS req + let status = HTTP.getResponseStatusCode response + if status >= 200 && status < 300 + then do + BL.writeFile tmpPath (HTTP.getResponseBody response) + pure (Right ()) + else pure (Left ("Download failed: HTTP " <> tshow status)) + case downloadResult of + Left (e :: SomeException) -> + pure (Left ("Download error: " <> tshow e)) + Right (Left err) -> pure (Left err) + Right (Right ()) -> do + result <- extractPdfText tmpPath + case result of + Left err -> pure (Left err) + Right text -> do + let truncated = Text.take maxChars text + pure (Right truncated) + +chunkText :: Int -> Text -> [Text] +chunkText chunkSize text + | Text.null text = [] + | Text.length text <= chunkSize = [text] + | otherwise = + let (chunk, rest) = Text.splitAt chunkSize text + in chunk : chunkText chunkSize rest diff --git a/Omni/Agent/Tools/Python.hs b/Omni/Agent/Tools/Python.hs new file mode 100644 index 0000000..99f3f7d --- /dev/null +++ b/Omni/Agent/Tools/Python.hs @@ -0,0 +1,217 @@ +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE NoImplicitPrelude #-} + +-- | Python execution tool for agent use. +-- +-- Executes Python snippets via subprocess with timeout support. +-- Writes code to temp file, executes with python3, cleans up after. +-- +-- Available stdlib: requests, json, csv, re, datetime, urllib +-- +-- : out omni-agent-tools-python +-- : dep aeson +-- : dep process +-- : dep directory +-- : dep temporary +module Omni.Agent.Tools.Python + ( pythonExecTool, + PythonExecArgs (..), + PythonResult (..), + main, + test, + ) +where + +import Alpha +import Data.Aeson ((.:), (.:?), (.=)) +import qualified Data.Aeson as Aeson +import qualified Data.Text as Text +import qualified Data.Text.IO as TextIO +import qualified Omni.Agent.Engine as Engine +import qualified Omni.Test as Test +import qualified System.Directory as Directory +import qualified System.Exit as Exit +import qualified System.Process as Process +import System.Timeout (timeout) + +main :: IO () +main = Test.run test + +test :: Test.Tree +test = + Test.group + "Omni.Agent.Tools.Python" + [ Test.unit "pythonExecTool has correct name" <| do + Engine.toolName pythonExecTool Test.@=? "python_exec", + Test.unit "pythonExecTool schema is valid" <| do + let schema = Engine.toolJsonSchema pythonExecTool + case schema of + Aeson.Object _ -> pure () + _ -> Test.assertFailure "Schema should be an object", + Test.unit "PythonExecArgs parses correctly" <| do + let json = Aeson.object ["code" .= ("print('hello')" :: Text)] + case Aeson.fromJSON json of + Aeson.Success (args :: PythonExecArgs) -> pythonCode args Test.@=? "print('hello')" + Aeson.Error e -> Test.assertFailure e, + Test.unit "PythonExecArgs parses with timeout" <| do + let json = Aeson.object ["code" .= ("x = 1" :: Text), "timeout" .= (10 :: Int)] + case Aeson.fromJSON json of + Aeson.Success (args :: PythonExecArgs) -> do + pythonCode args Test.@=? "x = 1" + pythonTimeout args Test.@=? Just 10 + Aeson.Error e -> Test.assertFailure e, + Test.unit "simple print statement" <| do + let args = Aeson.object ["code" .= ("print('hello world')" :: Text)] + result <- Engine.toolExecute pythonExecTool args + case Aeson.fromJSON result of + Aeson.Success (r :: PythonResult) -> do + pythonResultExitCode r Test.@=? 0 + ("hello world" `Text.isInfixOf` pythonResultStdout r) Test.@=? True + Aeson.Error e -> Test.assertFailure e, + Test.unit "syntax error handling" <| do + let args = Aeson.object ["code" .= ("def broken(" :: Text)] + result <- Engine.toolExecute pythonExecTool args + case Aeson.fromJSON result of + Aeson.Success (r :: PythonResult) -> do + (pythonResultExitCode r /= 0) Test.@=? True + not (Text.null (pythonResultStderr r)) Test.@=? True + Aeson.Error e -> Test.assertFailure e, + Test.unit "import json works" <| do + let code = "import json\nprint(json.dumps({'a': 1}))" + args = Aeson.object ["code" .= (code :: Text)] + result <- Engine.toolExecute pythonExecTool args + case Aeson.fromJSON result of + Aeson.Success (r :: PythonResult) -> do + pythonResultExitCode r Test.@=? 0 + ("{\"a\": 1}" `Text.isInfixOf` pythonResultStdout r) Test.@=? True + Aeson.Error e -> Test.assertFailure e, + Test.unit "timeout handling" <| do + let code = "import time\ntime.sleep(5)" + args = Aeson.object ["code" .= (code :: Text), "timeout" .= (1 :: Int)] + result <- Engine.toolExecute pythonExecTool args + case Aeson.fromJSON result of + Aeson.Success (r :: PythonResult) -> do + pythonResultExitCode r Test.@=? (-1) + ("timeout" `Text.isInfixOf` Text.toLower (pythonResultStderr r)) Test.@=? True + Aeson.Error e -> Test.assertFailure e + ] + +data PythonExecArgs = PythonExecArgs + { pythonCode :: Text, + pythonTimeout :: Maybe Int + } + deriving (Show, Eq, Generic) + +instance Aeson.FromJSON PythonExecArgs where + parseJSON = + Aeson.withObject "PythonExecArgs" <| \v -> + (PythonExecArgs </ (v .: "code")) + <*> (v .:? "timeout") + +data PythonResult = PythonResult + { pythonResultStdout :: Text, + pythonResultStderr :: Text, + pythonResultExitCode :: Int + } + deriving (Show, Eq, Generic) + +instance Aeson.ToJSON PythonResult where + toJSON r = + Aeson.object + [ "stdout" .= pythonResultStdout r, + "stderr" .= pythonResultStderr r, + "exit_code" .= pythonResultExitCode r + ] + +instance Aeson.FromJSON PythonResult where + parseJSON = + Aeson.withObject "PythonResult" <| \v -> + (PythonResult </ (v .: "stdout")) + <*> (v .: "stderr") + <*> (v .: "exit_code") + +pythonExecTool :: Engine.Tool +pythonExecTool = + Engine.Tool + { Engine.toolName = "python_exec", + Engine.toolDescription = + "Execute Python code and return the output. " + <> "Use for data processing, API calls, calculations, or any task requiring Python. " + <> "Available libraries: requests, json, csv, re, datetime, urllib. " + <> "Code runs in a subprocess with a 30 second default timeout.", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "code" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Python code to execute" :: Text) + ], + "timeout" + .= Aeson.object + [ "type" .= ("integer" :: Text), + "description" .= ("Timeout in seconds (default: 30)" :: Text) + ] + ], + "required" .= (["code"] :: [Text]) + ], + Engine.toolExecute = executePythonExec + } + +executePythonExec :: Aeson.Value -> IO Aeson.Value +executePythonExec v = + case Aeson.fromJSON v of + Aeson.Error e -> pure <| mkError ("Invalid arguments: " <> Text.pack e) + Aeson.Success args -> do + let code = pythonCode args + timeoutSecs = fromMaybe 30 (pythonTimeout args) + timeoutMicros = timeoutSecs * 1000000 + tmpDir <- Directory.getTemporaryDirectory + let tmpFile = tmpDir <> "/python_exec_" <> show (codeHash code) <> ".py" + result <- + try <| do + TextIO.writeFile tmpFile code + let proc = Process.proc "python3" [tmpFile] + mResult <- timeout timeoutMicros <| Process.readCreateProcessWithExitCode proc "" + Directory.removeFile tmpFile + pure mResult + case result of + Left (e :: SomeException) -> do + _ <- try @SomeException <| Directory.removeFile tmpFile + pure <| mkError ("Execution failed: " <> tshow e) + Right Nothing -> do + _ <- try @SomeException <| Directory.removeFile tmpFile + pure + <| Aeson.toJSON + <| PythonResult + { pythonResultStdout = "", + pythonResultStderr = "Timeout: execution exceeded " <> tshow timeoutSecs <> " seconds", + pythonResultExitCode = -1 + } + Right (Just (exitCode, stdoutStr, stderrStr)) -> + pure + <| Aeson.toJSON + <| PythonResult + { pythonResultStdout = Text.pack stdoutStr, + pythonResultStderr = Text.pack stderrStr, + pythonResultExitCode = exitCodeToInt exitCode + } + +exitCodeToInt :: Exit.ExitCode -> Int +exitCodeToInt Exit.ExitSuccess = 0 +exitCodeToInt (Exit.ExitFailure n) = n + +mkError :: Text -> Aeson.Value +mkError err = + Aeson.toJSON + <| PythonResult + { pythonResultStdout = "", + pythonResultStderr = err, + pythonResultExitCode = -1 + } + +codeHash :: Text -> Int +codeHash = Text.foldl' (\h c -> 31 * h + fromEnum c) 0 diff --git a/Omni/Agent/Tools/Todos.hs b/Omni/Agent/Tools/Todos.hs new file mode 100644 index 0000000..2aacacc --- /dev/null +++ b/Omni/Agent/Tools/Todos.hs @@ -0,0 +1,527 @@ +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE NoImplicitPrelude #-} + +-- | Todo tool with due dates and reminders. +-- +-- Provides user-scoped todos with optional due dates. +-- +-- : out omni-agent-tools-todos +-- : dep aeson +-- : dep sqlite-simple +-- : dep time +module Omni.Agent.Tools.Todos + ( -- * Tools + todoAddTool, + todoListTool, + todoCompleteTool, + todoDeleteTool, + + -- * Direct API + Todo (..), + createTodo, + listTodos, + listPendingTodos, + listOverdueTodos, + completeTodo, + deleteTodo, + + -- * Reminders + listTodosDueForReminder, + markReminderSent, + reminderInterval, + + -- * Database + initTodosTable, + + -- * Testing + main, + test, + ) +where + +import Alpha +import Data.Aeson ((.!=), (.:), (.:?), (.=)) +import qualified Data.Aeson as Aeson +import qualified Data.Text as Text +import Data.Time (LocalTime, NominalDiffTime, TimeZone, UTCTime, addUTCTime, getCurrentTime, localTimeToUTC, minutesToTimeZone, utcToLocalTime) +import Data.Time.Format (defaultTimeLocale, formatTime, parseTimeM) +import qualified Database.SQLite.Simple as SQL +import qualified Omni.Agent.Engine as Engine +import qualified Omni.Agent.Memory as Memory +import qualified Omni.Test as Test + +main :: IO () +main = Test.run test + +test :: Test.Tree +test = + Test.group + "Omni.Agent.Tools.Todos" + [ Test.unit "todoAddTool has correct schema" <| do + let tool = todoAddTool "test-user-id" + Engine.toolName tool Test.@=? "todo_add", + Test.unit "todoListTool has correct schema" <| do + let tool = todoListTool "test-user-id" + Engine.toolName tool Test.@=? "todo_list", + Test.unit "todoCompleteTool has correct schema" <| do + let tool = todoCompleteTool "test-user-id" + Engine.toolName tool Test.@=? "todo_complete", + Test.unit "todoDeleteTool has correct schema" <| do + let tool = todoDeleteTool "test-user-id" + Engine.toolName tool Test.@=? "todo_delete", + Test.unit "Todo JSON roundtrip" <| do + now <- getCurrentTime + let td = + Todo + { todoId = 1, + todoUserId = "user-123", + todoTitle = "Buy milk", + todoDueDate = Just now, + todoCompleted = False, + todoCreatedAt = now, + todoLastRemindedAt = Nothing + } + case Aeson.decode (Aeson.encode td) of + Nothing -> Test.assertFailure "Failed to decode Todo" + Just decoded -> do + todoTitle decoded Test.@=? "Buy milk" + todoCompleted decoded Test.@=? False, + Test.unit "parseDueDate handles various formats" <| do + isJust (parseDueDate "2024-12-25") Test.@=? True + isJust (parseDueDate "2024-12-25 14:00") Test.@=? True + ] + +data Todo = Todo + { todoId :: Int, + todoUserId :: Text, + todoTitle :: Text, + todoDueDate :: Maybe UTCTime, + todoCompleted :: Bool, + todoCreatedAt :: UTCTime, + todoLastRemindedAt :: Maybe UTCTime + } + deriving (Show, Eq, Generic) + +instance Aeson.ToJSON Todo where + toJSON td = + Aeson.object + [ "id" .= todoId td, + "user_id" .= todoUserId td, + "title" .= todoTitle td, + "due_date" .= todoDueDate td, + "completed" .= todoCompleted td, + "created_at" .= todoCreatedAt td, + "last_reminded_at" .= todoLastRemindedAt td + ] + +instance Aeson.FromJSON Todo where + parseJSON = + Aeson.withObject "Todo" <| \v -> + (Todo </ (v .: "id")) + <*> (v .: "user_id") + <*> (v .: "title") + <*> (v .:? "due_date") + <*> (v .: "completed") + <*> (v .: "created_at") + <*> (v .:? "last_reminded_at") + +instance SQL.FromRow Todo where + fromRow = + (Todo </ SQL.field) + <*> SQL.field + <*> SQL.field + <*> SQL.field + <*> SQL.field + <*> SQL.field + <*> SQL.field + +initTodosTable :: SQL.Connection -> IO () +initTodosTable conn = do + SQL.execute_ + conn + "CREATE TABLE IF NOT EXISTS todos (\ + \ id INTEGER PRIMARY KEY AUTOINCREMENT,\ + \ user_id TEXT NOT NULL,\ + \ title TEXT NOT NULL,\ + \ due_date TIMESTAMP,\ + \ completed INTEGER NOT NULL DEFAULT 0,\ + \ created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,\ + \ last_reminded_at TIMESTAMP\ + \)" + SQL.execute_ + conn + "CREATE INDEX IF NOT EXISTS idx_todos_user ON todos(user_id)" + SQL.execute_ + conn + "CREATE INDEX IF NOT EXISTS idx_todos_due ON todos(user_id, due_date)" + migrateTodosTable conn + +migrateTodosTable :: SQL.Connection -> IO () +migrateTodosTable conn = do + cols <- SQL.query_ conn "PRAGMA table_info(todos)" :: IO [(Int, Text, Text, Int, Maybe Text, Int)] + let colNames = map (\(_, name, _, _, _, _) -> name) cols + unless ("last_reminded_at" `elem` colNames) <| do + SQL.execute_ conn "ALTER TABLE todos ADD COLUMN last_reminded_at TIMESTAMP" + +easternTimeZone :: TimeZone +easternTimeZone = minutesToTimeZone (-300) + +parseDueDate :: Text -> Maybe UTCTime +parseDueDate txt = + let s = Text.unpack txt + parseLocal :: Maybe LocalTime + parseLocal = + parseTimeM True defaultTimeLocale "%Y-%m-%d %H:%M" s + <|> parseTimeM True defaultTimeLocale "%Y-%m-%d" s + <|> parseTimeM True defaultTimeLocale "%Y-%m-%dT%H:%M:%S" s + in fmap (localTimeToUTC easternTimeZone) parseLocal + <|> parseTimeM True defaultTimeLocale "%Y-%m-%dT%H:%M:%SZ" s + +createTodo :: Text -> Text -> Maybe Text -> IO Todo +createTodo uid title maybeDueDateStr = do + now <- getCurrentTime + let dueDate = maybeDueDateStr +> parseDueDate + Memory.withMemoryDb <| \conn -> do + initTodosTable conn + SQL.execute + conn + "INSERT INTO todos (user_id, title, due_date, completed, created_at) VALUES (?, ?, ?, 0, ?)" + (uid, title, dueDate, now) + rowId <- SQL.lastInsertRowId conn + pure + Todo + { todoId = fromIntegral rowId, + todoUserId = uid, + todoTitle = title, + todoDueDate = dueDate, + todoCompleted = False, + todoCreatedAt = now, + todoLastRemindedAt = Nothing + } + +listTodos :: Text -> Int -> IO [Todo] +listTodos uid limit = + Memory.withMemoryDb <| \conn -> do + initTodosTable conn + SQL.query + conn + "SELECT id, user_id, title, due_date, completed, created_at, last_reminded_at \ + \FROM todos WHERE user_id = ? \ + \ORDER BY completed ASC, due_date ASC NULLS LAST, created_at DESC LIMIT ?" + (uid, limit) + +listPendingTodos :: Text -> Int -> IO [Todo] +listPendingTodos uid limit = + Memory.withMemoryDb <| \conn -> do + initTodosTable conn + SQL.query + conn + "SELECT id, user_id, title, due_date, completed, created_at, last_reminded_at \ + \FROM todos WHERE user_id = ? AND completed = 0 \ + \ORDER BY due_date ASC NULLS LAST, created_at DESC LIMIT ?" + (uid, limit) + +listOverdueTodos :: Text -> IO [Todo] +listOverdueTodos uid = do + now <- getCurrentTime + Memory.withMemoryDb <| \conn -> do + initTodosTable conn + SQL.query + conn + "SELECT id, user_id, title, due_date, completed, created_at, last_reminded_at \ + \FROM todos WHERE user_id = ? AND completed = 0 AND due_date < ? \ + \ORDER BY due_date ASC" + (uid, now) + +reminderInterval :: NominalDiffTime +reminderInterval = 24 * 60 * 60 + +listTodosDueForReminder :: IO [Todo] +listTodosDueForReminder = do + now <- getCurrentTime + let cutoff = addUTCTime (negate reminderInterval) now + Memory.withMemoryDb <| \conn -> do + initTodosTable conn + SQL.query + conn + "SELECT id, user_id, title, due_date, completed, created_at, last_reminded_at \ + \FROM todos \ + \WHERE completed = 0 \ + \ AND due_date IS NOT NULL \ + \ AND due_date < ? \ + \ AND (last_reminded_at IS NULL OR last_reminded_at < ?)" + (now, cutoff) + +markReminderSent :: Int -> IO () +markReminderSent tid = do + now <- getCurrentTime + Memory.withMemoryDb <| \conn -> do + initTodosTable conn + SQL.execute + conn + "UPDATE todos SET last_reminded_at = ? WHERE id = ?" + (now, tid) + +completeTodo :: Text -> Int -> IO Bool +completeTodo uid tid = + Memory.withMemoryDb <| \conn -> do + initTodosTable conn + SQL.execute + conn + "UPDATE todos SET completed = 1 WHERE id = ? AND user_id = ?" + (tid, uid) + changes <- SQL.changes conn + pure (changes > 0) + +deleteTodo :: Text -> Int -> IO Bool +deleteTodo uid tid = + Memory.withMemoryDb <| \conn -> do + initTodosTable conn + SQL.execute + conn + "DELETE FROM todos WHERE id = ? AND user_id = ?" + (tid, uid) + changes <- SQL.changes conn + pure (changes > 0) + +todoAddTool :: Text -> Engine.Tool +todoAddTool uid = + Engine.Tool + { Engine.toolName = "todo_add", + Engine.toolDescription = + "Add a todo item with optional due date. Use for tasks, reminders, " + <> "or anything the user needs to remember to do. " + <> "Due date format: 'YYYY-MM-DD' or 'YYYY-MM-DD HH:MM'.", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "title" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("What needs to be done" :: Text) + ], + "due_date" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Optional due date in Eastern time: 'YYYY-MM-DD' or 'YYYY-MM-DD HH:MM'" :: Text) + ] + ], + "required" .= (["title"] :: [Text]) + ], + Engine.toolExecute = executeTodoAdd uid + } + +executeTodoAdd :: Text -> Aeson.Value -> IO Aeson.Value +executeTodoAdd uid v = + case Aeson.fromJSON v of + Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e]) + Aeson.Success (args :: TodoAddArgs) -> do + td <- createTodo uid (taTitle args) (taDueDate args) + let dueDateMsg = case todoDueDate td of + Just d -> + let localTime = utcToLocalTime easternTimeZone d + in " (due: " <> Text.pack (formatTime defaultTimeLocale "%Y-%m-%d %H:%M ET" localTime) <> ")" + Nothing -> "" + pure + ( Aeson.object + [ "success" .= True, + "todo_id" .= todoId td, + "message" .= ("Added todo: " <> todoTitle td <> dueDateMsg) + ] + ) + +data TodoAddArgs = TodoAddArgs + { taTitle :: Text, + taDueDate :: Maybe Text + } + deriving (Generic) + +instance Aeson.FromJSON TodoAddArgs where + parseJSON = + Aeson.withObject "TodoAddArgs" <| \v -> + (TodoAddArgs </ (v .: "title")) + <*> (v .:? "due_date") + +todoListTool :: Text -> Engine.Tool +todoListTool uid = + Engine.Tool + { Engine.toolName = "todo_list", + Engine.toolDescription = + "List todos. By default shows pending (incomplete) todos. " + <> "Can show all todos or just overdue ones.", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "filter" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Filter: 'pending' (default), 'all', or 'overdue'" :: Text) + ], + "limit" + .= Aeson.object + [ "type" .= ("integer" :: Text), + "description" .= ("Max todos to return (default: 20)" :: Text) + ] + ], + "required" .= ([] :: [Text]) + ], + Engine.toolExecute = executeTodoList uid + } + +executeTodoList :: Text -> Aeson.Value -> IO Aeson.Value +executeTodoList uid v = + case Aeson.fromJSON v of + Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e]) + Aeson.Success (args :: TodoListArgs) -> do + let lim = min 50 (max 1 (tlLimit args)) + todos <- case tlFilter args of + "all" -> listTodos uid lim + "overdue" -> listOverdueTodos uid + _ -> listPendingTodos uid lim + pure + ( Aeson.object + [ "success" .= True, + "count" .= length todos, + "todos" .= formatTodosForLLM todos + ] + ) + +formatTodosForLLM :: [Todo] -> Text +formatTodosForLLM [] = "No todos found." +formatTodosForLLM todos = + Text.unlines (map formatTodo todos) + where + formatTodo td = + let status = if todoCompleted td then "[x]" else "[ ]" + dueStr = case todoDueDate td of + Just d -> + let localTime = utcToLocalTime easternTimeZone d + in " (due: " <> Text.pack (formatTime defaultTimeLocale "%Y-%m-%d %H:%M ET" localTime) <> ")" + Nothing -> "" + in status <> " " <> todoTitle td <> dueStr <> " (id: " <> tshow (todoId td) <> ")" + +data TodoListArgs = TodoListArgs + { tlFilter :: Text, + tlLimit :: Int + } + deriving (Generic) + +instance Aeson.FromJSON TodoListArgs where + parseJSON = + Aeson.withObject "TodoListArgs" <| \v -> + (TodoListArgs </ (v .:? "filter" .!= "pending")) + <*> (v .:? "limit" .!= 20) + +todoCompleteTool :: Text -> Engine.Tool +todoCompleteTool uid = + Engine.Tool + { Engine.toolName = "todo_complete", + Engine.toolDescription = + "Mark a todo as completed. Use when the user says they finished something.", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "todo_id" + .= Aeson.object + [ "type" .= ("integer" :: Text), + "description" .= ("The ID of the todo to complete" :: Text) + ] + ], + "required" .= (["todo_id"] :: [Text]) + ], + Engine.toolExecute = executeTodoComplete uid + } + +executeTodoComplete :: Text -> Aeson.Value -> IO Aeson.Value +executeTodoComplete uid v = + case Aeson.fromJSON v of + Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e]) + Aeson.Success (args :: TodoCompleteArgs) -> do + completed <- completeTodo uid (tcTodoId args) + if completed + then + pure + ( Aeson.object + [ "success" .= True, + "message" .= ("Todo marked as complete" :: Text) + ] + ) + else + pure + ( Aeson.object + [ "success" .= False, + "error" .= ("Todo not found" :: Text) + ] + ) + +newtype TodoCompleteArgs = TodoCompleteArgs + { tcTodoId :: Int + } + deriving (Generic) + +instance Aeson.FromJSON TodoCompleteArgs where + parseJSON = + Aeson.withObject "TodoCompleteArgs" <| \v -> + TodoCompleteArgs </ (v .: "todo_id") + +todoDeleteTool :: Text -> Engine.Tool +todoDeleteTool uid = + Engine.Tool + { Engine.toolName = "todo_delete", + Engine.toolDescription = + "Delete a todo permanently. Use when a todo is no longer needed.", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "todo_id" + .= Aeson.object + [ "type" .= ("integer" :: Text), + "description" .= ("The ID of the todo to delete" :: Text) + ] + ], + "required" .= (["todo_id"] :: [Text]) + ], + Engine.toolExecute = executeTodoDelete uid + } + +executeTodoDelete :: Text -> Aeson.Value -> IO Aeson.Value +executeTodoDelete uid v = + case Aeson.fromJSON v of + Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e]) + Aeson.Success (args :: TodoDeleteArgs) -> do + deleted <- deleteTodo uid (tdTodoId args) + if deleted + then + pure + ( Aeson.object + [ "success" .= True, + "message" .= ("Todo deleted" :: Text) + ] + ) + else + pure + ( Aeson.object + [ "success" .= False, + "error" .= ("Todo not found" :: Text) + ] + ) + +newtype TodoDeleteArgs = TodoDeleteArgs + { tdTodoId :: Int + } + deriving (Generic) + +instance Aeson.FromJSON TodoDeleteArgs where + parseJSON = + Aeson.withObject "TodoDeleteArgs" <| \v -> + TodoDeleteArgs </ (v .: "todo_id") diff --git a/Omni/Agent/Tools/WebReader.hs b/Omni/Agent/Tools/WebReader.hs new file mode 100644 index 0000000..a69e3cf --- /dev/null +++ b/Omni/Agent/Tools/WebReader.hs @@ -0,0 +1,308 @@ +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE NoImplicitPrelude #-} + +-- | Web page reader tool - fetches and summarizes web pages. +-- +-- : out omni-agent-tools-webreader +-- : dep aeson +-- : dep http-conduit +-- : run trafilatura +module Omni.Agent.Tools.WebReader + ( -- * Tool + webReaderTool, + + -- * Direct API + fetchWebpage, + extractText, + fetchAndSummarize, + + -- * Testing + main, + test, + ) +where + +import Alpha +import qualified Control.Concurrent.Sema as Sema +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 Data.Text.Encoding as TE +import qualified Data.Text.IO as TIO +import Data.Time.Clock (diffUTCTime, getCurrentTime) +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.Provider as Provider +import qualified Omni.Test as Test +import qualified System.Exit as Exit +import qualified System.IO as IO +import qualified System.Process as Process +import qualified System.Timeout as Timeout + +main :: IO () +main = Test.run test + +test :: Test.Tree +test = + Test.group + "Omni.Agent.Tools.WebReader" + [ Test.unit "extractText removes HTML tags" <| do + let html = "<html><body><p>Hello world</p></body></html>" + result = extractText html + ("Hello world" `Text.isInfixOf` result) Test.@=? True, + Test.unit "extractText removes script tags" <| do + let html = "<html><script>alert('hi')</script><p>Content</p></html>" + result = extractText html + ("alert" `Text.isInfixOf` result) Test.@=? False + ("Content" `Text.isInfixOf` result) Test.@=? True, + Test.unit "webReaderTool has correct schema" <| do + let tool = webReaderTool "test-key" + Engine.toolName tool Test.@=? "read_webpages" + ] + +-- | Fetch timeout in microseconds (15 seconds - short because blocked sites won't respond anyway) +fetchTimeoutMicros :: Int +fetchTimeoutMicros = 15 * 1000000 + +-- | Summarization timeout in microseconds (30 seconds) +summarizeTimeoutMicros :: Int +summarizeTimeoutMicros = 30 * 1000000 + +-- | Maximum concurrent fetches +maxConcurrentFetches :: Int +maxConcurrentFetches = 10 + +-- | Simple debug logging to stderr +dbg :: Text -> IO () +dbg = TIO.hPutStrLn IO.stderr + +fetchWebpage :: Text -> IO (Either Text Text) +fetchWebpage url = do + dbg ("[webreader] Fetching: " <> url) + result <- + Timeout.timeout fetchTimeoutMicros <| do + innerResult <- + try <| do + req0 <- HTTP.parseRequest (Text.unpack url) + let req = + HTTP.setRequestMethod "GET" + <| HTTP.setRequestHeader "User-Agent" ["Mozilla/5.0 (compatible; OmniBot/1.0)"] + <| HTTP.setRequestHeader "Accept" ["text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8"] + <| HTTP.setRequestResponseTimeout (HTTPClient.responseTimeoutMicro fetchTimeoutMicros) + <| req0 + HTTP.httpLBS req + case innerResult of + Left (e :: SomeException) -> do + dbg ("[webreader] Fetch error: " <> url <> " - " <> tshow e) + pure (Left ("Failed to fetch URL: " <> tshow e)) + Right response -> do + let status = HTTP.getResponseStatusCode response + if status >= 200 && status < 300 + then do + let body = HTTP.getResponseBody response + text = TE.decodeUtf8With (\_ _ -> Just '?') (BL.toStrict body) + len = Text.length text + dbg ("[webreader] Fetched: " <> url <> " (" <> tshow len <> " chars)") + pure (Right text) + else do + dbg ("[webreader] HTTP " <> tshow status <> ": " <> url) + pure (Left ("HTTP error: " <> tshow status)) + case result of + Nothing -> do + dbg ("[webreader] Timeout: " <> url) + pure (Left ("Timeout fetching " <> url)) + Just r -> pure r + +-- | Fast single-pass text extraction from HTML +-- Strips all tags in one pass, no expensive operations +extractText :: Text -> Text +extractText html = collapseWhitespace (stripAllTags html) + where + -- Single pass: accumulate text outside of tags + stripAllTags :: Text -> Text + stripAllTags txt = Text.pack (go (Text.unpack txt) False []) + where + go :: [Char] -> Bool -> [Char] -> [Char] + go [] _ acc = reverse acc + go ('<' : rest) _ acc = go rest True acc -- Enter tag + go ('>' : rest) True acc = go rest False (' ' : acc) -- Exit tag, add space + go (_ : rest) True acc = go rest True acc -- Inside tag, skip + go (c : rest) False acc = go rest False (c : acc) -- Outside tag, keep + collapseWhitespace = Text.unwords <. Text.words + +-- | Maximum chars to send for summarization (keep it small for fast LLM response) +maxContentForSummary :: Int +maxContentForSummary = 15000 + +-- | Maximum summary length to return +maxSummaryLength :: Int +maxSummaryLength = 1000 + +-- | Timeout for trafilatura extraction in microseconds (10 seconds) +extractTimeoutMicros :: Int +extractTimeoutMicros = 10 * 1000000 + +-- | Extract article content using trafilatura (Python library) +-- Falls back to naive extractText if trafilatura fails +extractWithTrafilatura :: Text -> IO Text +extractWithTrafilatura html = do + let pythonScript = + "import sys; import trafilatura; " + <> "html = sys.stdin.read(); " + <> "result = trafilatura.extract(html, include_comments=False, include_tables=False); " + <> "print(result if result else '')" + proc = + (Process.proc "python3" ["-c", Text.unpack pythonScript]) + { Process.std_in = Process.CreatePipe, + Process.std_out = Process.CreatePipe, + Process.std_err = Process.CreatePipe + } + result <- + Timeout.timeout extractTimeoutMicros <| do + (exitCode, stdoutStr, _stderrStr) <- Process.readCreateProcessWithExitCode proc (Text.unpack html) + case exitCode of + Exit.ExitSuccess -> pure (Text.strip (Text.pack stdoutStr)) + Exit.ExitFailure _ -> pure "" + case result of + Just txt | not (Text.null txt) -> pure txt + _ -> do + dbg "[webreader] trafilatura failed, falling back to naive extraction" + pure (extractText (Text.take 100000 html)) + +summarizeContent :: Text -> Text -> Text -> IO (Either Text Text) +summarizeContent apiKey url content = do + let truncatedContent = Text.take maxContentForSummary content + haiku = Provider.defaultOpenRouter apiKey "anthropic/claude-haiku-4.5" + dbg ("[webreader] Summarizing: " <> url <> " (" <> tshow (Text.length truncatedContent) <> " chars)") + dbg "[webreader] Calling LLM for summarization..." + startTime <- getCurrentTime + result <- + Timeout.timeout summarizeTimeoutMicros + <| Provider.chat + haiku + [] + [ Provider.Message + Provider.System + ( "You are a webpage summarizer. Extract the key information in 3-5 bullet points. " + <> "Be extremely concise - max 500 characters total. No preamble, just bullets." + ) + Nothing + Nothing, + Provider.Message + Provider.User + ("Summarize: " <> url <> "\n\n" <> truncatedContent) + Nothing + Nothing + ] + endTime <- getCurrentTime + let elapsed = diffUTCTime endTime startTime + dbg ("[webreader] LLM call completed in " <> tshow elapsed) + case result of + Nothing -> do + dbg ("[webreader] Summarize timeout after " <> tshow elapsed <> ": " <> url) + pure (Left ("Timeout summarizing " <> url)) + Just (Left err) -> do + dbg ("[webreader] Summarize error: " <> url <> " - " <> err) + pure (Left ("Summarization failed: " <> err)) + Just (Right msg) -> do + let summary = Text.take maxSummaryLength (Provider.msgContent msg) + dbg ("[webreader] Summarized: " <> url <> " (" <> tshow (Text.length summary) <> " chars)") + pure (Right summary) + +-- | Fetch and summarize a single URL, returning a result object +-- This is the core function used by both single and batch tools +fetchAndSummarize :: Text -> Text -> IO Aeson.Value +fetchAndSummarize apiKey url = do + fetchResult <- fetchWebpage url + case fetchResult of + Left err -> + pure (Aeson.object ["url" .= url, "error" .= err]) + Right html -> do + dbg ("[webreader] Extracting article from: " <> url <> " (" <> tshow (Text.length html) <> " chars HTML)") + extractStart <- getCurrentTime + textContent <- extractWithTrafilatura html + extractEnd <- getCurrentTime + let extractElapsed = diffUTCTime extractEnd extractStart + dbg ("[webreader] Extracted: " <> url <> " (" <> tshow (Text.length textContent) <> " chars text) in " <> tshow extractElapsed) + if Text.null (Text.strip textContent) + then pure (Aeson.object ["url" .= url, "error" .= ("Page appears to be empty or JavaScript-only" :: Text)]) + else do + summaryResult <- summarizeContent apiKey url textContent + case summaryResult of + Left err -> + pure + ( Aeson.object + [ "url" .= url, + "error" .= err, + "raw_content" .= Text.take 2000 textContent + ] + ) + Right summary -> + pure + ( Aeson.object + [ "url" .= url, + "success" .= True, + "summary" .= summary + ] + ) + +-- | Web reader tool - fetches and summarizes webpages in parallel +webReaderTool :: Text -> Engine.Tool +webReaderTool apiKey = + Engine.Tool + { Engine.toolName = "read_webpages", + Engine.toolDescription = + "Fetch and summarize webpages in parallel. Each page is processed independently - " + <> "failures on one page won't affect others. Returns a list of summaries.", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "urls" + .= Aeson.object + [ "type" .= ("array" :: Text), + "items" .= Aeson.object ["type" .= ("string" :: Text)], + "description" .= ("List of URLs to read and summarize" :: Text) + ] + ], + "required" .= (["urls"] :: [Text]) + ], + Engine.toolExecute = executeWebReader apiKey + } + +executeWebReader :: Text -> Aeson.Value -> IO Aeson.Value +executeWebReader apiKey v = + case Aeson.fromJSON v of + Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e]) + Aeson.Success (args :: WebReaderArgs) -> do + let urls = wrUrls args + dbg ("[webreader] Starting batch: " <> tshow (length urls) <> " URLs") + results <- Sema.mapPool maxConcurrentFetches (fetchAndSummarize apiKey) urls + let succeeded = length (filter isSuccess results) + dbg ("[webreader] Batch complete: " <> tshow succeeded <> "/" <> tshow (length urls) <> " succeeded") + pure + ( Aeson.object + [ "results" .= results, + "total" .= length urls, + "succeeded" .= succeeded + ] + ) + where + isSuccess (Aeson.Object obj) = KeyMap.member "success" obj + isSuccess _ = False + +newtype WebReaderArgs = WebReaderArgs + { wrUrls :: [Text] + } + deriving (Generic) + +instance Aeson.FromJSON WebReaderArgs where + parseJSON = + Aeson.withObject "WebReaderArgs" <| \v -> + WebReaderArgs </ (v Aeson..: "urls") diff --git a/Omni/Agent/Tools/WebReaderTest.hs b/Omni/Agent/Tools/WebReaderTest.hs new file mode 100644 index 0000000..ca4c119 --- /dev/null +++ b/Omni/Agent/Tools/WebReaderTest.hs @@ -0,0 +1,53 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE NoImplicitPrelude #-} + +-- | Quick test for WebReader to debug hangs +-- +-- : out webreader-test +-- : dep http-conduit +-- : run trafilatura +module Omni.Agent.Tools.WebReaderTest where + +import Alpha +import qualified Data.Text as Text +import qualified Data.Text.IO as TIO +import Data.Time.Clock (diffUTCTime, getCurrentTime) +import qualified Omni.Agent.Tools.WebReader as WebReader + +main :: IO () +main = do + TIO.putStrLn "=== WebReader Debug Test ===" + + TIO.putStrLn "\n--- Test 1: Small page (httpbin) ---" + testUrl "https://httpbin.org/html" + + TIO.putStrLn "\n--- Test 2: Medium page (example.com) ---" + testUrl "https://example.com" + + TIO.putStrLn "\n--- Test 3: Large page (github) ---" + testUrl "https://github.com/anthropics/skills" + + TIO.putStrLn "\n=== Done ===" + +testUrl :: Text -> IO () +testUrl url = do + TIO.putStrLn ("Fetching: " <> url) + + startFetch <- getCurrentTime + result <- WebReader.fetchWebpage url + endFetch <- getCurrentTime + TIO.putStrLn ("Fetch took: " <> tshow (diffUTCTime endFetch startFetch)) + + case result of + Left err -> TIO.putStrLn ("Fetch error: " <> err) + Right html -> do + TIO.putStrLn ("HTML size: " <> tshow (Text.length html) <> " chars") + + TIO.putStrLn "Extracting text (naive, 100k truncated)..." + startExtract <- getCurrentTime + let !text = WebReader.extractText (Text.take 100000 html) + endExtract <- getCurrentTime + TIO.putStrLn ("Extract took: " <> tshow (diffUTCTime endExtract startExtract)) + TIO.putStrLn ("Text size: " <> tshow (Text.length text) <> " chars") + TIO.putStrLn ("Preview: " <> Text.take 200 text) diff --git a/Omni/Agent/Tools/WebSearch.hs b/Omni/Agent/Tools/WebSearch.hs new file mode 100644 index 0000000..58c945c --- /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: 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) |