Merge Ava deployment changes

1 files changed, 217 insertions, 0 deletions

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