From a6863d562a76eff5de36e0faa244e6ae2310bc22 Mon Sep 17 00:00:00 2001 From: Ben Sima Date: Fri, 12 Dec 2025 18:55:15 -0500 Subject: Add PDF and Notes tools to Telegram bot - Omni/Agent/Tools/Pdf.hs: Extract text from PDFs using pdftotext - Omni/Agent/Tools/Notes.hs: Quick notes CRUD with topics - Add notes table schema to Memory.hs initMemoryDb - Wire both tools into Telegram bot with logging callbacks --- Omni/Agent/Tools/Notes.hs | 357 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 357 insertions(+) create mode 100644 Omni/Agent/Tools/Notes.hs (limited to 'Omni/Agent/Tools/Notes.hs') 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 .: "user_id") + <*> (v .: "topic") + <*> (v .: "content") + <*> (v .: "created_at") + +instance SQL.FromRow Note where + fromRow = + (Note 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 .: "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 .:? "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