1 files changed, 148 insertions, 0 deletions
diff --git a/Omni/Ava/Trace.hs b/Omni/Ava/Trace.hs
new file mode 100644
index 0000000..6dbdf51
--- /dev/null
+++ b/Omni/Ava/Trace.hs
@@ -0,0 +1,148 @@
+{-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE NoImplicitPrelude #-}
+
+-- | Tool trace storage for Ava.
+--
+-- Records tool execution traces for debugging and analytics.
+--
+-- : out omni-ava-trace
+-- : dep aeson
+-- : dep sqlite-simple
+-- : dep uuid
+module Omni.Ava.Trace
+  ( TraceRecord (..),
+    insertTrace,
+    getTrace,
+    cleanupOldTraces,
+    main,
+    test,
+  )
+where
+
+import Alpha
+import Data.Aeson ((.=))
+import qualified Data.Aeson as Aeson
+import qualified Data.Text as Text
+import qualified Data.UUID as UUID
+import qualified Data.UUID.V4 as UUID
+import qualified Database.SQLite.Simple as SQL
+import qualified Database.SQLite.Simple.ToField as SQL
+import qualified Omni.Test as Test
+
+main :: IO ()
+main = Test.run test
+
+test :: Test.Tree
+test =
+  Test.group
+    "Omni.Ava.Trace"
+    [ Test.unit "TraceRecord JSON roundtrip" <| do
+        let tr =
+              TraceRecord
+                { trcId = "trace-123",
+                  trcCreatedAt = "2024-01-15T10:30:00Z",
+                  trcToolName = "web_search",
+                  trcInput = "{\"query\":\"test\"}",
+                  trcOutput = "{\"results\":[]}",
+                  trcDurationMs = 150,
+                  trcUserId = Just "user-456",
+                  trcChatId = Just "chat-789"
+                }
+        case Aeson.decode (Aeson.encode tr) of
+          Nothing -> Test.assertFailure "Failed to decode TraceRecord"
+          Just decoded -> do
+            trcToolName decoded Test.@=? "web_search"
+            trcDurationMs decoded Test.@=? 150
+    ]
+
+data TraceRecord = TraceRecord
+  { trcId :: Text,
+    trcCreatedAt :: Text,
+    trcToolName :: Text,
+    trcInput :: Text,
+    trcOutput :: Text,
+    trcDurationMs :: Int,
+    trcUserId :: Maybe Text,
+    trcChatId :: Maybe Text
+  }
+  deriving (Show, Eq, Generic)
+
+instance Aeson.ToJSON TraceRecord where
+  toJSON tr =
+    Aeson.object
+      [ "id" .= trcId tr,
+        "created_at" .= trcCreatedAt tr,
+        "tool_name" .= trcToolName tr,
+        "input" .= trcInput tr,
+        "output" .= trcOutput tr,
+        "duration_ms" .= trcDurationMs tr,
+        "user_id" .= trcUserId tr,
+        "chat_id" .= trcChatId tr
+      ]
+
+instance Aeson.FromJSON TraceRecord where
+  parseJSON =
+    Aeson.withObject "TraceRecord" <| \v ->
+      (TraceRecord </ (v Aeson..: "id"))
+        <*> (v Aeson..: "created_at")
+        <*> (v Aeson..: "tool_name")
+        <*> (v Aeson..: "input")
+        <*> (v Aeson..: "output")
+        <*> (v Aeson..: "duration_ms")
+        <*> (v Aeson..:? "user_id")
+        <*> (v Aeson..:? "chat_id")
+
+instance SQL.FromRow TraceRecord where
+  fromRow =
+    (TraceRecord </ SQL.field)
+      <*> SQL.field
+      <*> SQL.field
+      <*> SQL.field
+      <*> SQL.field
+      <*> SQL.field
+      <*> SQL.field
+      <*> SQL.field
+
+instance SQL.ToRow TraceRecord where
+  toRow tr =
+    [ SQL.toField (trcId tr),
+      SQL.toField (trcCreatedAt tr),
+      SQL.toField (trcToolName tr),
+      SQL.toField (trcInput tr),
+      SQL.toField (trcOutput tr),
+      SQL.toField (trcDurationMs tr),
+      SQL.toField (trcUserId tr),
+      SQL.toField (trcChatId tr)
+    ]
+
+insertTrace :: SQL.Connection -> TraceRecord -> IO Text
+insertTrace conn tr = do
+  tid <-
+    if Text.null (trcId tr)
+      then (Text.pack <. UUID.toString) </ UUID.nextRandom
+      else pure (trcId tr)
+  let trWithId = tr {trcId = tid}
+  SQL.execute
+    conn
+    "INSERT INTO tool_traces (id, created_at, tool_name, input, output, duration_ms, user_id, chat_id) \
+    \VALUES (?, ?, ?, ?, ?, ?, ?, ?)"
+    trWithId
+  pure tid
+
+getTrace :: SQL.Connection -> Text -> IO (Maybe TraceRecord)
+getTrace conn tid = do
+  results <-
+    SQL.query
+      conn
+      "SELECT id, created_at, tool_name, input, output, duration_ms, user_id, chat_id \
+      \FROM tool_traces WHERE id = ?"
+      (SQL.Only tid)
+  pure (listToMaybe results)
+
+cleanupOldTraces :: SQL.Connection -> IO Int
+cleanupOldTraces conn = do
+  SQL.execute_
+    conn
+    "DELETE FROM tool_traces WHERE created_at < datetime('now', '-7 days')"
+  SQL.changes conn