summaryrefslogtreecommitdiff
path: root/Omni
diff options
context:
space:
mode:
authorBen Sima <ben@bensima.com>2025-12-14 22:45:09 -0500
committerBen Sima <ben@bensima.com>2025-12-14 22:45:09 -0500
commit8c07a16dd9a7a3ad1847d0c665265e98f7df5438 (patch)
tree983ec8406dc738315d6645dde6f48e898e11283b /Omni
parent89d9fc7449ab2e799742470c3294c6e062e6de0b (diff)
Add python_exec tool for agent Python execution
- Create Omni/Agent/Tools/Python.hs with python_exec tool - Execute Python snippets via subprocess with 30s default timeout - Return structured JSON with stdout, stderr, exit_code - Add 8 unit tests covering print, imports, errors, timeout - Wire tool into Telegram agent's tool list Completes t-265.1
Diffstat (limited to 'Omni')
-rw-r--r--Omni/Agent/Telegram.hs57
-rw-r--r--Omni/Agent/Tools/Python.hs217
2 files changed, 272 insertions, 2 deletions
diff --git a/Omni/Agent/Telegram.hs b/Omni/Agent/Telegram.hs
index 6da1484..2f0a029 100644
--- a/Omni/Agent/Telegram.hs
+++ b/Omni/Agent/Telegram.hs
@@ -92,6 +92,7 @@ import qualified Omni.Agent.Tools.Email as Email
import qualified Omni.Agent.Tools.Hledger as Hledger
import qualified Omni.Agent.Tools.Notes as Notes
import qualified Omni.Agent.Tools.Pdf as Pdf
+import qualified Omni.Agent.Tools.Python as Python
import qualified Omni.Agent.Tools.Todos as Todos
import qualified Omni.Agent.Tools.WebReader as WebReader
import qualified Omni.Agent.Tools.WebSearch as WebSearch
@@ -959,7 +960,8 @@ processEngagedMessage tgConfig provider engineCfg msg uid userName chatId userMe
if isEmailAuthorized userName
then Email.allEmailTools
else []
- tools = memoryTools <> searchTools <> webReaderTools <> pdfTools <> notesTools <> calendarTools <> todoTools <> messageTools <> hledgerTools <> emailTools
+ pythonTools = [Python.pythonExecTool]
+ tools = memoryTools <> searchTools <> webReaderTools <> pdfTools <> notesTools <> calendarTools <> todoTools <> messageTools <> hledgerTools <> emailTools <> pythonTools
let agentCfg =
Engine.defaultAgentConfig
@@ -1000,7 +1002,9 @@ processEngagedMessage tgConfig provider engineCfg msg uid userName chatId userMe
_ <- Messages.enqueueImmediate (Just uid) chatId threadId "hmm, i don't have a response for that" (Just "agent_response") Nothing
pure ()
else do
- _ <- Messages.enqueueImmediate (Just uid) chatId threadId response (Just "agent_response") Nothing
+ parts <- splitMessageForChat (Types.tgOpenRouterApiKey tgConfig) response
+ putText <| "Split response into " <> tshow (length parts) <> " parts"
+ enqueueMultipart (Just uid) chatId threadId parts (Just "agent_response")
unless isGroup <| checkAndSummarize (Types.tgOpenRouterApiKey tgConfig) uid chatId
let cost = Engine.resultTotalCost agentResult
costStr = Text.pack (printf "%.2f" cost)
@@ -1053,6 +1057,55 @@ checkAndSummarize openRouterKey uid chatId = do
_ <- Memory.summarizeAndArchive uid chatId summary
putText "Conversation summarized and archived (gemini)"
+splitMessageForChat :: Text -> Text -> IO [Text]
+splitMessageForChat openRouterKey message = do
+ if Text.length message < 200
+ then pure [message]
+ else do
+ let haiku = Provider.defaultOpenRouter openRouterKey "anthropic/claude-haiku-4.5"
+ result <-
+ Provider.chat
+ haiku
+ []
+ [ Provider.Message
+ Provider.System
+ ( Text.unlines
+ [ "Split this message into separate chat messages that feel natural in a messaging app.",
+ "Each part should be logically independent - a complete thought.",
+ "Separate parts with exactly '---' on its own line.",
+ "Keep the original text, just add separators. Don't add any commentary.",
+ "If the message is already short/simple, return it unchanged (no separators).",
+ "Aim for 2-4 parts maximum. Don't over-split.",
+ "",
+ "Good splits: between topics, after questions, between a statement and follow-up",
+ "Bad splits: mid-sentence, between closely related points"
+ ]
+ )
+ Nothing
+ Nothing,
+ Provider.Message Provider.User message Nothing Nothing
+ ]
+ case result of
+ Left err -> do
+ putText <| "Message split failed: " <> err
+ pure [message]
+ Right msg -> do
+ let parts = map Text.strip (Text.splitOn "---" (Provider.msgContent msg))
+ validParts = filter (not <. Text.null) parts
+ if null validParts
+ then pure [message]
+ else pure validParts
+
+enqueueMultipart :: Maybe Text -> Int -> Maybe Int -> [Text] -> Maybe Text -> IO ()
+enqueueMultipart _ _ _ [] _ = pure ()
+enqueueMultipart mUid chatId mThreadId parts msgType = do
+ forM_ (zip [0 ..] parts) <| \(i :: Int, part) -> do
+ if i == 0
+ then void <| Messages.enqueueImmediate mUid chatId mThreadId part msgType Nothing
+ else do
+ let delaySeconds = fromIntegral (i * 2)
+ void <| Messages.enqueueDelayed mUid chatId mThreadId part delaySeconds msgType Nothing
+
shouldEngageInGroup :: Text -> Text -> IO Bool
shouldEngageInGroup openRouterKey messageText = do
let gemini = Provider.defaultOpenRouter openRouterKey "google/gemini-2.0-flash-001"
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