path: root/Omni/Agent/Memory.hs

{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE NoImplicitPrelude #-}

-- | Cross-agent shared memory system with vector similarity search.
--
-- Provides persistent memory that is:
-- - Shared across all agents (Telegram, researcher, coder, etc.)
-- - Private per user (users can't see each other's memories)
-- - Searchable via semantic similarity using embeddings
--
-- Uses sqlite-vss for vector similarity search and Ollama for embeddings.
--
-- : out omni-agent-memory
-- : dep aeson
-- : dep http-conduit
-- : dep sqlite-simple
-- : dep uuid
-- : dep vector
-- : dep directory
-- : dep bytestring
module Omni.Agent.Memory
  ( -- * Types
    User (..),
    Memory (..),
    MemorySource (..),
    ConversationMessage (..),
    ConversationSummary (..),
    MessageRole (..),

    -- * User Management
    createUser,
    getUser,
    getUserByTelegramId,
    getOrCreateUserByTelegramId,

    -- * Memory Operations
    storeMemory,
    recallMemories,
    forgetMemory,
    getAllMemoriesForUser,
    updateMemoryAccess,

    -- * Conversation History
    saveMessage,
    getRecentMessages,
    getConversationContext,
    summarizeAndArchive,
    estimateTokens,

    -- * Embeddings
    embedText,

    -- * Agent Integration
    rememberTool,
    recallTool,
    formatMemoriesForPrompt,
    runAgentWithMemory,

    -- * Database
    withMemoryDb,
    initMemoryDb,
    getMemoryDbPath,

    -- * Testing
    main,
    test,
  )
where

import Alpha
import Data.Aeson ((.!=), (.:), (.:?), (.=))
import qualified Data.Aeson as Aeson
import qualified Data.Aeson.KeyMap as KeyMap
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BL
import qualified Data.List as List
import qualified Data.Text as Text
import qualified Data.Text.Encoding as TE
import Data.Time (UTCTime, getCurrentTime)
import qualified Data.UUID as UUID
import qualified Data.UUID.V4 as UUID
import qualified Data.Vector.Storable as VS
import qualified Database.SQLite.Simple as SQL
import Database.SQLite.Simple.FromField ()
import qualified Database.SQLite.Simple.ToField as SQL
import Foreign.Storable ()
import qualified Network.HTTP.Simple as HTTP
import qualified Omni.Agent.Engine as Engine
import qualified Omni.Test as Test
import System.Directory (createDirectoryIfMissing)
import System.Environment (lookupEnv)
import System.FilePath (takeDirectory, (</>))

main :: IO ()
main = Test.run test

test :: Test.Tree
test =
  Test.group
    "Omni.Agent.Memory"
    [ Test.unit "User JSON roundtrip" <| do
        now <- getCurrentTime
        let user =
              User
                { userId = "test-uuid",
                  userTelegramId = Just 12345,
                  userEmail = Nothing,
                  userName = "Test User",
                  userCreatedAt = now
                }
        case Aeson.decode (Aeson.encode user) of
          Nothing -> Test.assertFailure "Failed to decode User"
          Just decoded -> userName decoded Test.@=? "Test User",
      Test.unit "Memory JSON roundtrip" <| do
        now <- getCurrentTime
        let mem =
              Memory
                { memoryId = "mem-uuid",
                  memoryUserId = "user-uuid",
                  memoryContent = "User is an AI engineer",
                  memoryEmbedding = Nothing,
                  memorySource =
                    MemorySource
                      { sourceAgent = "telegram",
                        sourceSession = Nothing,
                        sourceContext = "User mentioned in chat"
                      },
                  memoryConfidence = 0.9,
                  memoryCreatedAt = now,
                  memoryLastAccessedAt = now,
                  memoryTags = ["profession", "ai"]
                }
        case Aeson.decode (Aeson.encode mem) of
          Nothing -> Test.assertFailure "Failed to decode Memory"
          Just decoded -> memoryContent decoded Test.@=? "User is an AI engineer",
      Test.unit "MemorySource JSON roundtrip" <| do
        let src =
              MemorySource
                { sourceAgent = "researcher",
                  sourceSession = Just "session-123",
                  sourceContext = "Extracted from conversation"
                }
        case Aeson.decode (Aeson.encode src) of
          Nothing -> Test.assertFailure "Failed to decode MemorySource"
          Just decoded -> sourceAgent decoded Test.@=? "researcher",
      Test.unit "formatMemoriesForPrompt formats correctly" <| do
        now <- getCurrentTime
        let mem1 =
              Memory
                { memoryId = "1",
                  memoryUserId = "u",
                  memoryContent = "User is an AI engineer",
                  memoryEmbedding = Nothing,
                  memorySource = MemorySource "telegram" Nothing "chat",
                  memoryConfidence = 0.9,
                  memoryCreatedAt = now,
                  memoryLastAccessedAt = now,
                  memoryTags = []
                }
            mem2 =
              Memory
                { memoryId = "2",
                  memoryUserId = "u",
                  memoryContent = "User prefers Haskell",
                  memoryEmbedding = Nothing,
                  memorySource = MemorySource "coder" Nothing "code review",
                  memoryConfidence = 0.8,
                  memoryCreatedAt = now,
                  memoryLastAccessedAt = now,
                  memoryTags = []
                }
            formatted = formatMemoriesForPrompt [mem1, mem2]
        ("AI engineer" `Text.isInfixOf` formatted) Test.@=? True
        ("Haskell" `Text.isInfixOf` formatted) Test.@=? True,
      Test.unit "cosineSimilarity identical vectors" <| do
        let v1 = VS.fromList [1.0, 0.0, 0.0 :: Float]
            v2 = VS.fromList [1.0, 0.0, 0.0 :: Float]
        abs (cosineSimilarity v1 v2 - 1.0) < 0.0001 Test.@=? True,
      Test.unit "cosineSimilarity orthogonal vectors" <| do
        let v1 = VS.fromList [1.0, 0.0, 0.0 :: Float]
            v2 = VS.fromList [0.0, 1.0, 0.0 :: Float]
        abs (cosineSimilarity v1 v2) < 0.0001 Test.@=? True,
      Test.unit "cosineSimilarity opposite vectors" <| do
        let v1 = VS.fromList [1.0, 0.0, 0.0 :: Float]
            v2 = VS.fromList [-1.0, 0.0, 0.0 :: Float]
        abs (cosineSimilarity v1 v2 + 1.0) < 0.0001 Test.@=? True,
      Test.unit "vectorToBlob and blobToVector roundtrip" <| do
        let v = VS.fromList [0.1, 0.2, 0.3, 0.4, 0.5 :: Float]
            blob = vectorToBlob v
            v' = blobToVector blob
        VS.length v Test.@=? VS.length v'
        VS.toList v Test.@=? VS.toList v',
      Test.unit "rememberTool has correct schema" <| do
        let tool = rememberTool "test-user-id"
        Engine.toolName tool Test.@=? "remember",
      Test.unit "recallTool has correct schema" <| do
        let tool = recallTool "test-user-id"
        Engine.toolName tool Test.@=? "recall"
    ]

-- | User record for multi-user memory system.
data User = User
  { userId :: Text,
    userTelegramId :: Maybe Int,
    userEmail :: Maybe Text,
    userName :: Text,
    userCreatedAt :: UTCTime
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON User where
  toJSON u =
    Aeson.object
      [ "id" .= userId u,
        "telegram_id" .= userTelegramId u,
        "email" .= userEmail u,
        "name" .= userName u,
        "created_at" .= userCreatedAt u
      ]

instance Aeson.FromJSON User where
  parseJSON =
    Aeson.withObject "User" <| \v ->
      (User </ (v .: "id"))
        <*> (v .:? "telegram_id")
        <*> (v .:? "email")
        <*> (v .: "name")
        <*> (v .: "created_at")

instance SQL.FromRow User where
  fromRow =
    User
      </ SQL.field
      <*> SQL.field
      <*> SQL.field
      <*> SQL.field
      <*> SQL.field

instance SQL.ToRow User where
  toRow u =
    [ SQL.toField (userId u),
      SQL.toField (userTelegramId u),
      SQL.toField (userEmail u),
      SQL.toField (userName u),
      SQL.toField (userCreatedAt u)
    ]

-- | Source information for a memory.
data MemorySource = MemorySource
  { sourceAgent :: Text,
    sourceSession :: Maybe Text,
    sourceContext :: Text
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON MemorySource where
  toJSON s =
    Aeson.object
      [ "agent" .= sourceAgent s,
        "session" .= sourceSession s,
        "context" .= sourceContext s
      ]

instance Aeson.FromJSON MemorySource where
  parseJSON =
    Aeson.withObject "MemorySource" <| \v ->
      (MemorySource </ (v .: "agent"))
        <*> (v .:? "session")
        <*> (v .: "context")

-- | A memory stored in the system.
data Memory = Memory
  { memoryId :: Text,
    memoryUserId :: Text,
    memoryContent :: Text,
    memoryEmbedding :: Maybe (VS.Vector Float),
    memorySource :: MemorySource,
    memoryConfidence :: Double,
    memoryCreatedAt :: UTCTime,
    memoryLastAccessedAt :: UTCTime,
    memoryTags :: [Text]
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON Memory where
  toJSON m =
    Aeson.object
      [ "id" .= memoryId m,
        "user_id" .= memoryUserId m,
        "content" .= memoryContent m,
        "source" .= memorySource m,
        "confidence" .= memoryConfidence m,
        "created_at" .= memoryCreatedAt m,
        "last_accessed_at" .= memoryLastAccessedAt m,
        "tags" .= memoryTags m
      ]

instance Aeson.FromJSON Memory where
  parseJSON =
    Aeson.withObject "Memory" <| \v ->
      ( Memory
          </ (v .: "id")
      )
        <*> (v .: "user_id")
        <*> (v .: "content")
        <*> pure Nothing
        <*> (v .: "source")
        <*> (v .:? "confidence" .!= 0.8)
        <*> (v .: "created_at")
        <*> (v .: "last_accessed_at")
        <*> (v .:? "tags" .!= [])

-- SQLite instances for Memory (partial - embedding handled separately)
instance SQL.FromRow Memory where
  fromRow = do
    mid <- SQL.field
    uid <- SQL.field
    content <- SQL.field
    embeddingBlob <- SQL.field
    agent <- SQL.field
    session <- SQL.field
    context <- SQL.field
    confidence <- SQL.field
    createdAt <- SQL.field
    lastAccessedAt <- SQL.field
    tagsJson <- SQL.field
    let embedding = blobToVector </ (embeddingBlob :: Maybe BS.ByteString)
        source = MemorySource agent session context
        tags = fromMaybe [] ((tagsJson :: Maybe Text) +> (Aeson.decode <. BL.fromStrict <. TE.encodeUtf8))
    pure
      Memory
        { memoryId = mid,
          memoryUserId = uid,
          memoryContent = content,
          memoryEmbedding = embedding,
          memorySource = source,
          memoryConfidence = confidence,
          memoryCreatedAt = createdAt,
          memoryLastAccessedAt = lastAccessedAt,
          memoryTags = tags
        }

-- | Role in a conversation message.
data MessageRole = UserRole | AssistantRole
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON MessageRole where
  toJSON UserRole = Aeson.String "user"
  toJSON AssistantRole = Aeson.String "assistant"

instance Aeson.FromJSON MessageRole where
  parseJSON =
    Aeson.withText "MessageRole" <| \case
      "user" -> pure UserRole
      "assistant" -> pure AssistantRole
      _ -> empty

-- | A message in a conversation.
data ConversationMessage = ConversationMessage
  { cmId :: Maybe Int,
    cmUserId :: Text,
    cmChatId :: Int,
    cmRole :: MessageRole,
    cmSenderName :: Maybe Text,
    cmContent :: Text,
    cmTokensEstimate :: Int,
    cmCreatedAt :: UTCTime
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON ConversationMessage where
  toJSON m =
    Aeson.object
      [ "id" .= cmId m,
        "user_id" .= cmUserId m,
        "chat_id" .= cmChatId m,
        "role" .= cmRole m,
        "sender_name" .= cmSenderName m,
        "content" .= cmContent m,
        "tokens_estimate" .= cmTokensEstimate m,
        "created_at" .= cmCreatedAt m
      ]

instance SQL.FromRow ConversationMessage where
  fromRow =
    (ConversationMessage </ SQL.field)
      <*> SQL.field
      <*> SQL.field
      <*> (parseRole </ SQL.field)
      <*> SQL.field
      <*> SQL.field
      <*> (fromMaybe 0 </ SQL.field)
      <*> SQL.field
    where
      parseRole :: Text -> MessageRole
      parseRole "user" = UserRole
      parseRole _ = AssistantRole

-- | A summary of older conversation messages.
data ConversationSummary = ConversationSummary
  { csId :: Maybe Int,
    csUserId :: Text,
    csChatId :: Int,
    csSummary :: Text,
    csMessagesSummarized :: Int,
    csTokensSaved :: Maybe Int,
    csCreatedAt :: UTCTime
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON ConversationSummary where
  toJSON s =
    Aeson.object
      [ "id" .= csId s,
        "user_id" .= csUserId s,
        "chat_id" .= csChatId s,
        "summary" .= csSummary s,
        "messages_summarized" .= csMessagesSummarized s,
        "tokens_saved" .= csTokensSaved s,
        "created_at" .= csCreatedAt s
      ]

instance SQL.FromRow ConversationSummary where
  fromRow =
    (ConversationSummary </ SQL.field)
      <*> SQL.field
      <*> SQL.field
      <*> SQL.field
      <*> SQL.field
      <*> SQL.field
      <*> SQL.field

-- | Get the path to memory.db
getMemoryDbPath :: IO FilePath
getMemoryDbPath = do
  maybeEnv <- lookupEnv "MEMORY_DB_PATH"
  case maybeEnv of
    Just p -> pure p
    Nothing -> do
      home <- lookupEnv "HOME"
      case home of
        Just h -> pure (h </> ".local/share/omni/memory.db")
        Nothing -> pure "_/memory.db"

-- | Run an action with the memory database connection.
withMemoryDb :: (SQL.Connection -> IO a) -> IO a
withMemoryDb action = do
  dbPath <- getMemoryDbPath
  createDirectoryIfMissing True (takeDirectory dbPath)
  SQL.withConnection dbPath <| \conn -> do
    initMemoryDb conn
    action conn

-- | Initialize the memory database schema.
initMemoryDb :: SQL.Connection -> IO ()
initMemoryDb conn = do
  SQL.execute_ conn "PRAGMA foreign_keys = ON"
  SQL.execute_
    conn
    "CREATE TABLE IF NOT EXISTS users (\
    \  id TEXT PRIMARY KEY,\
    \  telegram_id INTEGER UNIQUE,\
    \  email TEXT UNIQUE,\
    \  name TEXT NOT NULL,\
    \  created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP\
    \)"
  SQL.execute_
    conn
    "CREATE TABLE IF NOT EXISTS memories (\
    \  id TEXT PRIMARY KEY,\
    \  user_id TEXT NOT NULL REFERENCES users(id),\
    \  content TEXT NOT NULL,\
    \  embedding BLOB,\
    \  source_agent TEXT NOT NULL,\
    \  source_session TEXT,\
    \  source_context TEXT,\
    \  confidence REAL DEFAULT 0.8,\
    \  created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,\
    \  last_accessed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,\
    \  tags TEXT\
    \)"
  SQL.execute_
    conn
    "CREATE INDEX IF NOT EXISTS idx_memories_user ON memories(user_id)"
  SQL.execute_
    conn
    "CREATE INDEX IF NOT EXISTS idx_memories_agent ON memories(source_agent)"
  SQL.execute_
    conn
    "CREATE TABLE IF NOT EXISTS conversation_messages (\
    \  id INTEGER PRIMARY KEY AUTOINCREMENT,\
    \  user_id TEXT NOT NULL REFERENCES users(id),\
    \  chat_id INTEGER NOT NULL,\
    \  role TEXT NOT NULL,\
    \  sender_name TEXT,\
    \  content TEXT NOT NULL,\
    \  tokens_estimate INTEGER,\
    \  created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP\
    \)"
  SQL.execute_
    conn
    "CREATE INDEX IF NOT EXISTS idx_conv_user_chat ON conversation_messages(user_id, chat_id)"
  migrateConversationMessages conn
  SQL.execute_
    conn
    "CREATE TABLE IF NOT EXISTS conversation_summaries (\
    \  id INTEGER PRIMARY KEY AUTOINCREMENT,\
    \  user_id TEXT NOT NULL REFERENCES users(id),\
    \  chat_id INTEGER NOT NULL,\
    \  summary TEXT NOT NULL,\
    \  messages_summarized INTEGER NOT NULL,\
    \  tokens_saved INTEGER,\
    \  created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP\
    \)"
  SQL.execute_
    conn
    "CREATE INDEX IF NOT EXISTS idx_summary_user_chat ON conversation_summaries(user_id, chat_id)"
  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)"
  SQL.execute_
    conn
    "CREATE TABLE IF NOT EXISTS todos (\
    \  id INTEGER PRIMARY KEY AUTOINCREMENT,\
    \  user_id TEXT NOT NULL,\
    \  title TEXT NOT NULL,\
    \  due_date TIMESTAMP,\
    \  completed INTEGER NOT NULL DEFAULT 0,\
    \  created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP\
    \)"
  SQL.execute_
    conn
    "CREATE INDEX IF NOT EXISTS idx_todos_user ON todos(user_id)"
  SQL.execute_
    conn
    "CREATE INDEX IF NOT EXISTS idx_todos_due ON todos(user_id, due_date)"

-- | Migrate conversation_messages to add sender_name column.
migrateConversationMessages :: SQL.Connection -> IO ()
migrateConversationMessages conn = do
  columns <- SQL.query_ conn "PRAGMA table_info(conversation_messages)" :: IO [(Int, Text, Text, Int, Maybe Text, Int)]
  let columnNames = map (\(_, name, _, _, _, _) -> name) columns
  unless ("sender_name" `elem` columnNames) <| do
    SQL.execute_ conn "ALTER TABLE conversation_messages ADD COLUMN sender_name TEXT"
    SQL.execute_ conn "UPDATE conversation_messages SET sender_name = 'bensima' WHERE role = 'user' AND sender_name IS NULL"

-- | Create a new user.
createUser :: Text -> Maybe Int -> IO User
createUser name telegramId = do
  uuid <- UUID.nextRandom
  now <- getCurrentTime
  let user =
        User
          { userId = UUID.toText uuid,
            userTelegramId = telegramId,
            userEmail = Nothing,
            userName = name,
            userCreatedAt = now
          }
  withMemoryDb <| \conn ->
    SQL.execute
      conn
      "INSERT INTO users (id, telegram_id, email, name, created_at) VALUES (?, ?, ?, ?, ?)"
      user
  pure user

-- | Get a user by ID.
getUser :: Text -> IO (Maybe User)
getUser uid =
  withMemoryDb <| \conn -> do
    results <- SQL.query conn "SELECT id, telegram_id, email, name, created_at FROM users WHERE id = ?" (SQL.Only uid)
    pure (listToMaybe results)

-- | Get a user by Telegram ID.
getUserByTelegramId :: Int -> IO (Maybe User)
getUserByTelegramId tid =
  withMemoryDb <| \conn -> do
    results <- SQL.query conn "SELECT id, telegram_id, email, name, created_at FROM users WHERE telegram_id = ?" (SQL.Only tid)
    pure (listToMaybe results)

-- | Get or create a user by Telegram ID.
getOrCreateUserByTelegramId :: Int -> Text -> IO User
getOrCreateUserByTelegramId tid name = do
  existing <- getUserByTelegramId tid
  case existing of
    Just user -> pure user
    Nothing -> createUser name (Just tid)

-- | Store a memory for a user.
storeMemory :: Text -> Text -> MemorySource -> IO Memory
storeMemory uid content source = storeMemoryWithTags uid content source []

-- | Store a memory with tags.
storeMemoryWithTags :: Text -> Text -> MemorySource -> [Text] -> IO Memory
storeMemoryWithTags uid content source tags = do
  uuid <- UUID.nextRandom
  now <- getCurrentTime
  embedding <- embedText content
  let mem =
        Memory
          { memoryId = UUID.toText uuid,
            memoryUserId = uid,
            memoryContent = content,
            memoryEmbedding = either (const Nothing) Just embedding,
            memorySource = source,
            memoryConfidence = 0.8,
            memoryCreatedAt = now,
            memoryLastAccessedAt = now,
            memoryTags = tags
          }
  withMemoryDb <| \conn ->
    SQL.execute
      conn
      "INSERT INTO memories (id, user_id, content, embedding, source_agent, source_session, source_context, confidence, created_at, last_accessed_at, tags) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"
      ( ( memoryId mem,
          memoryUserId mem,
          memoryContent mem,
          vectorToBlob </ memoryEmbedding mem,
          sourceAgent (memorySource mem),
          sourceSession (memorySource mem),
          sourceContext (memorySource mem)
        )
          SQL.:. ( memoryConfidence mem,
                   memoryCreatedAt mem,
                   memoryLastAccessedAt mem,
                   TE.decodeUtf8 (BL.toStrict (Aeson.encode (memoryTags mem)))
                 )
      )
  pure mem

-- | Recall memories for a user using semantic similarity.
recallMemories :: Text -> Text -> Int -> IO [Memory]
recallMemories uid query limit = do
  queryEmbedding <- embedText query
  case queryEmbedding of
    Left _ -> recallMemoriesByRecency uid limit
    Right qEmb -> do
      allMems <- getAllMemoriesForUser uid
      let scored =
            [ (m, cosineSimilarity qEmb emb)
              | m <- allMems,
                Just emb <- [memoryEmbedding m]
            ]
          sorted = List.sortBy (\(_, s1) (_, s2) -> compare s2 s1) scored
          topN = take limit sorted
      now <- getCurrentTime
      traverse_ (updateMemoryAccess now <. memoryId <. fst) topN
      pure (map fst topN)

-- | Recall memories by recency (fallback when embedding fails).
recallMemoriesByRecency :: Text -> Int -> IO [Memory]
recallMemoriesByRecency uid limit =
  withMemoryDb <| \conn -> do
    SQL.query
      conn
      "SELECT id, user_id, content, embedding, source_agent, source_session, source_context, confidence, created_at, last_accessed_at, tags \
      \FROM memories WHERE user_id = ? ORDER BY last_accessed_at DESC LIMIT ?"
      (uid, limit)

-- | Get all memories for a user.
getAllMemoriesForUser :: Text -> IO [Memory]
getAllMemoriesForUser uid =
  withMemoryDb <| \conn ->
    SQL.query
      conn
      "SELECT id, user_id, content, embedding, source_agent, source_session, source_context, confidence, created_at, last_accessed_at, tags \
      \FROM memories WHERE user_id = ?"
      (SQL.Only uid)

-- | Delete a memory.
forgetMemory :: Text -> IO ()
forgetMemory mid =
  withMemoryDb <| \conn ->
    SQL.execute conn "DELETE FROM memories WHERE id = ?" (SQL.Only mid)

-- | Update memory's last accessed timestamp.
updateMemoryAccess :: UTCTime -> Text -> IO ()
updateMemoryAccess now mid =
  withMemoryDb <| \conn ->
    SQL.execute conn "UPDATE memories SET last_accessed_at = ? WHERE id = ?" (now, mid)

-- | Embed text using Ollama's nomic-embed-text model.
embedText :: Text -> IO (Either Text (VS.Vector Float))
embedText content = do
  ollamaUrl <- fromMaybe "http://localhost:11434" </ lookupEnv "OLLAMA_URL"
  let url = ollamaUrl <> "/api/embeddings"
  req0 <- HTTP.parseRequest url
  let body =
        Aeson.object
          [ "model" .= ("nomic-embed-text" :: Text),
            "prompt" .= content
          ]
      req =
        HTTP.setRequestMethod "POST"
          <| HTTP.setRequestHeader "Content-Type" ["application/json"]
          <| HTTP.setRequestBodyLBS (Aeson.encode body)
          <| req0
  result <- try (HTTP.httpLBS req)
  case result of
    Left (e :: SomeException) ->
      pure (Left ("Embedding request failed: " <> tshow e))
    Right response -> do
      let status = HTTP.getResponseStatusCode response
      if status >= 200 && status < 300
        then case Aeson.decode (HTTP.getResponseBody response) of
          Just (Aeson.Object obj) -> case KeyMap.lookup "embedding" obj of
            Just (Aeson.Array arr) ->
              let floats = [f | Aeson.Number n <- toList arr, let f = realToFrac n]
               in pure (Right (VS.fromList floats))
            _ -> pure (Left "No embedding in response")
          _ -> pure (Left "Failed to parse embedding response")
        else pure (Left ("Embedding HTTP error: " <> tshow status))

-- | Convert a vector to a blob for storage.
vectorToBlob :: VS.Vector Float -> BS.ByteString
vectorToBlob v =
  let bytes = VS.unsafeCast v :: VS.Vector Word8
   in BS.pack (VS.toList bytes)

-- | Convert a blob back to a vector.
blobToVector :: BS.ByteString -> VS.Vector Float
blobToVector bs =
  let bytes = VS.fromList (BS.unpack bs) :: VS.Vector Word8
   in VS.unsafeCast bytes

-- | Calculate cosine similarity between two vectors.
cosineSimilarity :: VS.Vector Float -> VS.Vector Float -> Float
cosineSimilarity v1 v2
  | VS.length v1 /= VS.length v2 = 0
  | otherwise =
      let dot = VS.sum (VS.zipWith (*) v1 v2)
          mag1 = sqrt (VS.sum (VS.map (\x -> x * x) v1))
          mag2 = sqrt (VS.sum (VS.map (\x -> x * x) v2))
       in if mag1 == 0 || mag2 == 0 then 0 else dot / (mag1 * mag2)

-- | Format memories for inclusion in a prompt.
formatMemoriesForPrompt :: [Memory] -> Text
formatMemoriesForPrompt [] = "No prior context available."
formatMemoriesForPrompt mems =
  Text.unlines
    [ "Known context about this user:",
      "",
      Text.unlines (map formatMem mems)
    ]
  where
    formatMem m =
      "- " <> memoryContent m <> " (via " <> sourceAgent (memorySource m) <> ")"

-- | Run an agent with memory context.
-- Recalls relevant memories for the user and injects them into the system prompt.
runAgentWithMemory ::
  User ->
  Engine.EngineConfig ->
  Engine.AgentConfig ->
  Text ->
  IO (Either Text Engine.AgentResult)
runAgentWithMemory user engineCfg agentCfg userPrompt = do
  memories <- recallMemories (userId user) userPrompt 10
  let memoryContext = formatMemoriesForPrompt memories
      enhancedPrompt =
        Engine.agentSystemPrompt agentCfg
          <> "\n\n## Known about this user\n"
          <> memoryContext
      enhancedConfig =
        agentCfg
          { Engine.agentSystemPrompt = enhancedPrompt,
            Engine.agentTools =
              Engine.agentTools agentCfg
                <> [rememberTool (userId user), recallTool (userId user)]
          }
  Engine.runAgent engineCfg enhancedConfig userPrompt

-- | Tool for agents to store memories about users.
rememberTool :: Text -> Engine.Tool
rememberTool uid =
  Engine.Tool
    { Engine.toolName = "remember",
      Engine.toolDescription =
        "Store a piece of information about the user for future reference. "
          <> "Use this when the user shares personal facts, preferences, or context "
          <> "that would be useful to recall in future conversations.",
      Engine.toolJsonSchema =
        Aeson.object
          [ "type" .= ("object" :: Text),
            "properties"
              .= Aeson.object
                [ "content"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("The information to remember about the user" :: Text)
                      ],
                  "context"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("How/why this was learned (e.g., 'user mentioned in chat')" :: Text)
                      ],
                  "tags"
                    .= Aeson.object
                      [ "type" .= ("array" :: Text),
                        "items" .= Aeson.object ["type" .= ("string" :: Text)],
                        "description" .= ("Optional tags for categorization" :: Text)
                      ]
                ],
            "required" .= (["content", "context"] :: [Text])
          ],
      Engine.toolExecute = executeRemember uid
    }

executeRemember :: Text -> Aeson.Value -> IO Aeson.Value
executeRemember uid v =
  case Aeson.fromJSON v of
    Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e])
    Aeson.Success (args :: RememberArgs) -> do
      let source =
            MemorySource
              { sourceAgent = "agent",
                sourceSession = Nothing,
                sourceContext = rememberContext args
              }
      mem <- storeMemoryWithTags uid (rememberContent args) source (rememberTags args)
      pure
        ( Aeson.object
            [ "success" .= True,
              "memory_id" .= memoryId mem,
              "message" .= ("Remembered: " <> rememberContent args)
            ]
        )

-- | Tool for agents to recall memories about users.
recallTool :: Text -> Engine.Tool
recallTool uid =
  Engine.Tool
    { Engine.toolName = "recall",
      Engine.toolDescription =
        "Search your memory for information about the user. "
          <> "Use this to retrieve previously stored facts, preferences, or context.",
      Engine.toolJsonSchema =
        Aeson.object
          [ "type" .= ("object" :: Text),
            "properties"
              .= Aeson.object
                [ "query"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("What to search for in memory" :: Text)
                      ],
                  "limit"
                    .= Aeson.object
                      [ "type" .= ("integer" :: Text),
                        "description" .= ("Maximum memories to return (default: 5)" :: Text)
                      ]
                ],
            "required" .= (["query"] :: [Text])
          ],
      Engine.toolExecute = executeRecall uid
    }

executeRecall :: Text -> Aeson.Value -> IO Aeson.Value
executeRecall uid v =
  case Aeson.fromJSON v of
    Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e])
    Aeson.Success (args :: RecallArgs) -> do
      mems <- recallMemories uid (recallQuery args) (recallLimit args)
      pure
        ( Aeson.object
            [ "success" .= True,
              "count" .= length mems,
              "memories"
                .= map
                  ( \m ->
                      Aeson.object
                        [ "content" .= memoryContent m,
                          "confidence" .= memoryConfidence m,
                          "source" .= sourceAgent (memorySource m),
                          "tags" .= memoryTags m
                        ]
                  )
                  mems
            ]
        )

-- Helper for parsing remember args
data RememberArgs = RememberArgs
  { rememberContent :: Text,
    rememberContext :: Text,
    rememberTags :: [Text]
  }
  deriving (Generic)

instance Aeson.FromJSON RememberArgs where
  parseJSON =
    Aeson.withObject "RememberArgs" <| \v ->
      (RememberArgs </ (v .: "content"))
        <*> (v .:? "context" .!= "agent observation")
        <*> (v .:? "tags" .!= [])

data RecallArgs = RecallArgs
  { recallQuery :: Text,
    recallLimit :: Int
  }
  deriving (Generic)

instance Aeson.FromJSON RecallArgs where
  parseJSON =
    Aeson.withObject "RecallArgs" <| \v ->
      (RecallArgs </ (v .: "query"))
        <*> (v .:? "limit" .!= 5)

-- | Estimate token count for text (rough: ~4 chars per token).
estimateTokens :: Text -> Int
estimateTokens t = max 1 (Text.length t `div` 4)

-- | Save a message to conversation history.
saveMessage :: Text -> Int -> MessageRole -> Maybe Text -> Text -> IO ConversationMessage
saveMessage uid chatId role senderName content = do
  now <- getCurrentTime
  let tokens = estimateTokens content
  withMemoryDb <| \conn -> do
    SQL.execute
      conn
      "INSERT INTO conversation_messages (user_id, chat_id, role, sender_name, content, tokens_estimate, created_at) VALUES (?, ?, ?, ?, ?, ?, ?)"
      (uid, chatId, roleToText role, senderName, content, tokens, now)
    rowId <- SQL.lastInsertRowId conn
    pure
      ConversationMessage
        { cmId = Just (fromIntegral rowId),
          cmUserId = uid,
          cmChatId = chatId,
          cmRole = role,
          cmSenderName = senderName,
          cmContent = content,
          cmTokensEstimate = tokens,
          cmCreatedAt = now
        }
  where
    roleToText UserRole = "user" :: Text
    roleToText AssistantRole = "assistant"

-- | Get recent messages for a user/chat, newest first.
getRecentMessages :: Text -> Int -> Int -> IO [ConversationMessage]
getRecentMessages uid chatId limit =
  withMemoryDb <| \conn ->
    SQL.query
      conn
      "SELECT id, user_id, chat_id, role, sender_name, content, tokens_estimate, created_at \
      \FROM conversation_messages \
      \WHERE user_id = ? AND chat_id = ? \
      \ORDER BY created_at DESC LIMIT ?"
      (uid, chatId, limit)

-- | Get the most recent summary for a chat.
getLatestSummary :: Text -> Int -> IO (Maybe ConversationSummary)
getLatestSummary uid chatId =
  withMemoryDb <| \conn -> do
    rows <-
      SQL.query
        conn
        "SELECT id, user_id, chat_id, summary, messages_summarized, tokens_saved, created_at \
        \FROM conversation_summaries \
        \WHERE user_id = ? AND chat_id = ? \
        \ORDER BY created_at DESC LIMIT 1"
        (uid, chatId)
    pure (listToMaybe rows)

-- | Build conversation context for the LLM.
-- Returns (context text, total token estimate).
getConversationContext :: Text -> Int -> Int -> IO (Text, Int)
getConversationContext uid chatId maxTokens = do
  maybeSummary <- getLatestSummary uid chatId
  recentMsgs <- getRecentMessages uid chatId 50

  let summaryText = maybe "" (\s -> "## Previous conversation summary\n" <> csSummary s <> "\n\n") maybeSummary
      summaryTokens = maybe 0 (estimateTokens <. csSummary) maybeSummary

      msgsOldestFirst = reverse recentMsgs
      availableTokens = maxTokens - summaryTokens - 100

      (selectedMsgs, usedTokens) = selectMessages msgsOldestFirst availableTokens

      formattedMsgs =
        if null selectedMsgs
          then ""
          else
            "## Recent conversation\n"
              <> Text.unlines (map formatMsg selectedMsgs)

  pure (summaryText <> formattedMsgs, summaryTokens + usedTokens)
  where
    selectMessages :: [ConversationMessage] -> Int -> ([ConversationMessage], Int)
    selectMessages msgs budget = go (reverse msgs) budget []
      where
        go [] _ acc = (acc, sum (map cmTokensEstimate acc))
        go (m : ms) remaining acc
          | cmTokensEstimate m <= remaining =
              go ms (remaining - cmTokensEstimate m) (m : acc)
          | otherwise = (acc, sum (map cmTokensEstimate acc))

    formatMsg m =
      let prefix = case cmRole m of
            UserRole -> fromMaybe "User" (cmSenderName m) <> ": "
            AssistantRole -> "Assistant: "
       in prefix <> cmContent m

-- | Summarize old messages and archive them.
-- Returns the new summary text.
summarizeAndArchive :: Text -> Int -> Text -> IO Text
summarizeAndArchive uid chatId summaryText = do
  now <- getCurrentTime

  (oldMsgCount, tokensSaved) <-
    withMemoryDb <| \conn -> do
      rows <-
        SQL.query
          conn
          "SELECT COUNT(*), COALESCE(SUM(tokens_estimate), 0) FROM conversation_messages WHERE user_id = ? AND chat_id = ?"
          (uid, chatId) ::
          IO [(Int, Int)]
      let (count, tokens) = fromMaybe (0, 0) (listToMaybe rows)

      SQL.execute
        conn
        "INSERT INTO conversation_summaries (user_id, chat_id, summary, messages_summarized, tokens_saved, created_at) VALUES (?, ?, ?, ?, ?, ?)"
        (uid, chatId, summaryText, count, tokens, now)

      SQL.execute
        conn
        "DELETE FROM conversation_messages WHERE user_id = ? AND chat_id = ?"
        (uid, chatId)

      pure (count, tokens)

  putText <| "Archived " <> tshow oldMsgCount <> " messages (" <> tshow tokensSaved <> " tokens) for chat " <> tshow chatId
  pure summaryText