From 8c07a16dd9a7a3ad1847d0c665265e98f7df5438 Mon Sep 17 00:00:00 2001 From: Ben Sima Date: Sun, 14 Dec 2025 22:45:09 -0500 Subject: 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 --- Omni/Agent/Tools/Python.hs | 217 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 217 insertions(+) create mode 100644 Omni/Agent/Tools/Python.hs (limited to 'Omni/Agent/Tools/Python.hs') 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 .:? "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 .: "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 -- cgit v1.2.3