{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE OverloadedStrings #-}

{- |
Module      : Data.Ollama.Conversation
Copyright   : (c) 2025 Tushar Adhatrao
License     : MIT
Maintainer  : Tushar Adhatrao <tusharadhatrao@gmail.com>
Stability   : experimental
Description : Conversation management for the Ollama client, including storage and retrieval of chat sessions.

This module provides types and functions for managing conversations in the Ollama client. It defines
a 'Conversation' type to represent a chat session, a 'ConversationStore' typeclass for storage operations,
and an in-memory implementation using 'InMemoryStore' and 'ConvoM'. The module supports saving, loading,
listing, and deleting conversations, with thread-safe operations using STM (Software Transactional Memory).

The 'Conversation' type includes metadata such as a unique ID, messages, model name, and timestamps.
The 'ConversationStore' typeclass defines a generic interface for conversation storage, while 'InMemoryStore'
provides a concrete in-memory implementation. The 'ConvoM' monad integrates with 'InMemoryStore' for
monadic operations.

Example:

>>> store <- initInMemoryStore
>>> let conv = Conversation "conv1" [userMessage "Hello!"] "gemma3" <$> getCurrentTime <*> getCurrentTime
>>> runInMemoryConvo store $ saveConversation conv
>>> runInMemoryConvo store $ loadConversation "conv1"
Just (Conversation ...)
-}
module Data.Ollama.Conversation
  ( -- * Conversation Types
    Conversation (..)
  , ConversationStore (..)

    -- * In-Memory Store
  , InMemoryStore (..)
  , ConvoM (..)
  , initInMemoryStore
  , runInMemoryConvo

    -- * Validation
  , validateConversation
  ) where

import Control.Concurrent.STM
import Control.Monad.Reader
import Data.Aeson (FromJSON, ToJSON)
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Ollama.Common.Types
import Data.Text (Text)
import Data.Text qualified as T
import Data.Time (UTCTime, getCurrentTime)
import GHC.Generics (Generic)

{- | Represents a chat session with metadata and messages.

Stores a conversation's unique identifier, list of messages, model name, creation time, and last updated time.
-}
data Conversation = Conversation
  { Conversation -> Text
conversationId :: !Text
  -- ^ Unique identifier for the conversation.
  , Conversation -> [Message]
messages :: ![Message]
  -- ^ List of messages in the conversation.
  , Conversation -> Text
model :: !Text
  -- ^ Name of the model used in the conversation (e.g., "gemma3").
  , Conversation -> UTCTime
createdAt :: !UTCTime
  -- ^ Timestamp when the conversation was created.
  , Conversation -> UTCTime
lastUpdated :: !UTCTime
  -- ^ Timestamp when the conversation was last updated.
  }
  deriving (Int -> Conversation -> ShowS
[Conversation] -> ShowS
Conversation -> String
(Int -> Conversation -> ShowS)
-> (Conversation -> String)
-> ([Conversation] -> ShowS)
-> Show Conversation
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Conversation -> ShowS
showsPrec :: Int -> Conversation -> ShowS
$cshow :: Conversation -> String
show :: Conversation -> String
$cshowList :: [Conversation] -> ShowS
showList :: [Conversation] -> ShowS
Show, Conversation -> Conversation -> Bool
(Conversation -> Conversation -> Bool)
-> (Conversation -> Conversation -> Bool) -> Eq Conversation
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Conversation -> Conversation -> Bool
== :: Conversation -> Conversation -> Bool
$c/= :: Conversation -> Conversation -> Bool
/= :: Conversation -> Conversation -> Bool
Eq, (forall x. Conversation -> Rep Conversation x)
-> (forall x. Rep Conversation x -> Conversation)
-> Generic Conversation
forall x. Rep Conversation x -> Conversation
forall x. Conversation -> Rep Conversation x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. Conversation -> Rep Conversation x
from :: forall x. Conversation -> Rep Conversation x
$cto :: forall x. Rep Conversation x -> Conversation
to :: forall x. Rep Conversation x -> Conversation
Generic)

instance ToJSON Conversation
instance FromJSON Conversation

{- | Typeclass defining operations for storing and managing conversations.

Provides methods for saving, loading, listing, and deleting conversations in a monadic context.
--
-- @since 0.2.0.0
-}
class Monad m => ConversationStore m where
  -- | Saves a conversation to the store.
  --
  -- Validates the conversation and updates its 'lastUpdated' timestamp before saving.
  saveConversation :: Conversation -> m ()

  -- | Loads a conversation by its ID.
  --
  -- Returns 'Just' the conversation if found, or 'Nothing' if not.
  loadConversation :: Text -> m (Maybe Conversation)

  -- | Lists all conversations in the store.
  listConversations :: m [Conversation]

  -- | Deletes a conversation by its ID.
  --
  -- Returns 'True' if the conversation was found and deleted, 'False' otherwise.
  deleteConversation :: Text -> m Bool

{- | In-memory conversation store using a 'TVar' for thread-safe operations.

Stores conversations in a 'Map' keyed by conversation IDs, wrapped in a 'TVar' for concurrent access.
-}
newtype InMemoryStore = InMemoryStore (TVar (Map Text Conversation))

{- | Monad for operations with 'InMemoryStore'.

A wrapper around 'ReaderT' that provides access to an 'InMemoryStore' in a monadic context.
-}
newtype ConvoM a = ConvoM {forall a. ConvoM a -> ReaderT InMemoryStore IO a
runConvoM :: ReaderT InMemoryStore IO a}
  deriving ((forall a b. (a -> b) -> ConvoM a -> ConvoM b)
-> (forall a b. a -> ConvoM b -> ConvoM a) -> Functor ConvoM
forall a b. a -> ConvoM b -> ConvoM a
forall a b. (a -> b) -> ConvoM a -> ConvoM b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> ConvoM a -> ConvoM b
fmap :: forall a b. (a -> b) -> ConvoM a -> ConvoM b
$c<$ :: forall a b. a -> ConvoM b -> ConvoM a
<$ :: forall a b. a -> ConvoM b -> ConvoM a
Functor, Functor ConvoM
Functor ConvoM =>
(forall a. a -> ConvoM a)
-> (forall a b. ConvoM (a -> b) -> ConvoM a -> ConvoM b)
-> (forall a b c.
    (a -> b -> c) -> ConvoM a -> ConvoM b -> ConvoM c)
-> (forall a b. ConvoM a -> ConvoM b -> ConvoM b)
-> (forall a b. ConvoM a -> ConvoM b -> ConvoM a)
-> Applicative ConvoM
forall a. a -> ConvoM a
forall a b. ConvoM a -> ConvoM b -> ConvoM a
forall a b. ConvoM a -> ConvoM b -> ConvoM b
forall a b. ConvoM (a -> b) -> ConvoM a -> ConvoM b
forall a b c. (a -> b -> c) -> ConvoM a -> ConvoM b -> ConvoM c
forall (f :: * -> *).
Functor f =>
(forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
$cpure :: forall a. a -> ConvoM a
pure :: forall a. a -> ConvoM a
$c<*> :: forall a b. ConvoM (a -> b) -> ConvoM a -> ConvoM b
<*> :: forall a b. ConvoM (a -> b) -> ConvoM a -> ConvoM b
$cliftA2 :: forall a b c. (a -> b -> c) -> ConvoM a -> ConvoM b -> ConvoM c
liftA2 :: forall a b c. (a -> b -> c) -> ConvoM a -> ConvoM b -> ConvoM c
$c*> :: forall a b. ConvoM a -> ConvoM b -> ConvoM b
*> :: forall a b. ConvoM a -> ConvoM b -> ConvoM b
$c<* :: forall a b. ConvoM a -> ConvoM b -> ConvoM a
<* :: forall a b. ConvoM a -> ConvoM b -> ConvoM a
Applicative, Applicative ConvoM
Applicative ConvoM =>
(forall a b. ConvoM a -> (a -> ConvoM b) -> ConvoM b)
-> (forall a b. ConvoM a -> ConvoM b -> ConvoM b)
-> (forall a. a -> ConvoM a)
-> Monad ConvoM
forall a. a -> ConvoM a
forall a b. ConvoM a -> ConvoM b -> ConvoM b
forall a b. ConvoM a -> (a -> ConvoM b) -> ConvoM b
forall (m :: * -> *).
Applicative m =>
(forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
$c>>= :: forall a b. ConvoM a -> (a -> ConvoM b) -> ConvoM b
>>= :: forall a b. ConvoM a -> (a -> ConvoM b) -> ConvoM b
$c>> :: forall a b. ConvoM a -> ConvoM b -> ConvoM b
>> :: forall a b. ConvoM a -> ConvoM b -> ConvoM b
$creturn :: forall a. a -> ConvoM a
return :: forall a. a -> ConvoM a
Monad, Monad ConvoM
Monad ConvoM => (forall a. IO a -> ConvoM a) -> MonadIO ConvoM
forall a. IO a -> ConvoM a
forall (m :: * -> *).
Monad m =>
(forall a. IO a -> m a) -> MonadIO m
$cliftIO :: forall a. IO a -> ConvoM a
liftIO :: forall a. IO a -> ConvoM a
MonadIO, MonadReader InMemoryStore)

{- | Runs a 'ConvoM' action with the given 'InMemoryStore'.

Executes a monadic computation in the context of an in-memory store.

Example:

>>> store <- initInMemoryStore
>>> runInMemoryConvo store $ saveConversation conv
-}
runInMemoryConvo :: InMemoryStore -> ConvoM a -> IO a
runInMemoryConvo :: forall a. InMemoryStore -> ConvoM a -> IO a
runInMemoryConvo InMemoryStore
store = (ReaderT InMemoryStore IO a -> InMemoryStore -> IO a)
-> InMemoryStore -> ReaderT InMemoryStore IO a -> IO a
forall a b c. (a -> b -> c) -> b -> a -> c
flip ReaderT InMemoryStore IO a -> InMemoryStore -> IO a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT InMemoryStore
store (ReaderT InMemoryStore IO a -> IO a)
-> (ConvoM a -> ReaderT InMemoryStore IO a) -> ConvoM a -> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ConvoM a -> ReaderT InMemoryStore IO a
forall a. ConvoM a -> ReaderT InMemoryStore IO a
runConvoM

instance ConversationStore ConvoM where
  saveConversation :: Conversation -> ConvoM ()
saveConversation Conversation
conv = do
    case Conversation -> Either Text Conversation
validateConversation Conversation
conv of
      Left Text
err -> IO () -> ConvoM ()
forall a. IO a -> ConvoM a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> ConvoM ()) -> IO () -> ConvoM ()
forall a b. (a -> b) -> a -> b
$ String -> IO ()
putStrLn (String
"Validation error: " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Text -> String
T.unpack Text
err)
      Right Conversation
validConv -> do
        UTCTime
now <- IO UTCTime -> ConvoM UTCTime
forall a. IO a -> ConvoM a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO UTCTime
getCurrentTime
        let updatedConv :: Conversation
updatedConv = Conversation
validConv {lastUpdated = now}
        InMemoryStore TVar (Map Text Conversation)
ref <- ConvoM InMemoryStore
forall r (m :: * -> *). MonadReader r m => m r
ask
        IO () -> ConvoM ()
forall a. IO a -> ConvoM a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> ConvoM ()) -> (STM () -> IO ()) -> STM () -> ConvoM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> ConvoM ()) -> STM () -> ConvoM ()
forall a b. (a -> b) -> a -> b
$ TVar (Map Text Conversation)
-> (Map Text Conversation -> Map Text Conversation) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar (Map Text Conversation)
ref (Text
-> Conversation -> Map Text Conversation -> Map Text Conversation
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert (Conversation -> Text
conversationId Conversation
updatedConv) Conversation
updatedConv)

  loadConversation :: Text -> ConvoM (Maybe Conversation)
loadConversation Text
cid = do
    InMemoryStore TVar (Map Text Conversation)
ref <- ConvoM InMemoryStore
forall r (m :: * -> *). MonadReader r m => m r
ask
    Map Text Conversation
convs <- IO (Map Text Conversation) -> ConvoM (Map Text Conversation)
forall a. IO a -> ConvoM a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Map Text Conversation) -> ConvoM (Map Text Conversation))
-> IO (Map Text Conversation) -> ConvoM (Map Text Conversation)
forall a b. (a -> b) -> a -> b
$ TVar (Map Text Conversation) -> IO (Map Text Conversation)
forall a. TVar a -> IO a
readTVarIO TVar (Map Text Conversation)
ref
    Maybe Conversation -> ConvoM (Maybe Conversation)
forall a. a -> ConvoM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Conversation -> ConvoM (Maybe Conversation))
-> Maybe Conversation -> ConvoM (Maybe Conversation)
forall a b. (a -> b) -> a -> b
$ Text -> Map Text Conversation -> Maybe Conversation
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Text
cid Map Text Conversation
convs

  listConversations :: ConvoM [Conversation]
listConversations = do
    InMemoryStore TVar (Map Text Conversation)
ref <- ConvoM InMemoryStore
forall r (m :: * -> *). MonadReader r m => m r
ask
    Map Text Conversation
convs <- IO (Map Text Conversation) -> ConvoM (Map Text Conversation)
forall a. IO a -> ConvoM a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Map Text Conversation) -> ConvoM (Map Text Conversation))
-> IO (Map Text Conversation) -> ConvoM (Map Text Conversation)
forall a b. (a -> b) -> a -> b
$ TVar (Map Text Conversation) -> IO (Map Text Conversation)
forall a. TVar a -> IO a
readTVarIO TVar (Map Text Conversation)
ref
    [Conversation] -> ConvoM [Conversation]
forall a. a -> ConvoM a
forall (m :: * -> *) a. Monad m => a -> m a
return ([Conversation] -> ConvoM [Conversation])
-> [Conversation] -> ConvoM [Conversation]
forall a b. (a -> b) -> a -> b
$ Map Text Conversation -> [Conversation]
forall k a. Map k a -> [a]
Map.elems Map Text Conversation
convs

  deleteConversation :: Text -> ConvoM Bool
deleteConversation Text
cid = do
    InMemoryStore TVar (Map Text Conversation)
ref <- ConvoM InMemoryStore
forall r (m :: * -> *). MonadReader r m => m r
ask
    IO Bool -> ConvoM Bool
forall a. IO a -> ConvoM a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Bool -> ConvoM Bool)
-> (STM Bool -> IO Bool) -> STM Bool -> ConvoM Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. STM Bool -> IO Bool
forall a. STM a -> IO a
atomically (STM Bool -> ConvoM Bool) -> STM Bool -> ConvoM Bool
forall a b. (a -> b) -> a -> b
$ do
      Map Text Conversation
convs <- TVar (Map Text Conversation) -> STM (Map Text Conversation)
forall a. TVar a -> STM a
readTVar TVar (Map Text Conversation)
ref
      if Text -> Map Text Conversation -> Bool
forall k a. Ord k => k -> Map k a -> Bool
Map.member Text
cid Map Text Conversation
convs
        then do
          TVar (Map Text Conversation) -> Map Text Conversation -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (Map Text Conversation)
ref (Text -> Map Text Conversation -> Map Text Conversation
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete Text
cid Map Text Conversation
convs)
          Bool -> STM Bool
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
        else Bool -> STM Bool
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False

{- | Validates a 'Conversation' to ensure required fields are non-empty.

Checks that the 'conversationId' is not empty and that the 'messages' list contains at least one message.
Returns 'Right' with the validated conversation or 'Left' with an error message.

Example:

>>> let conv = Conversation "" [] "gemma3" time time
>>> validateConversation conv
Left "Conversation ID cannot be empty"
-}
validateConversation :: Conversation -> Either Text Conversation
validateConversation :: Conversation -> Either Text Conversation
validateConversation Conversation
conv
  | Text -> Bool
T.null (Conversation -> Text
conversationId Conversation
conv) = Text -> Either Text Conversation
forall a b. a -> Either a b
Left Text
"Conversation ID cannot be empty"
  | [Message] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (Conversation -> [Message]
messages Conversation
conv) = Text -> Either Text Conversation
forall a b. a -> Either a b
Left Text
"Conversation must have at least one message"
  | Bool
otherwise = Conversation -> Either Text Conversation
forall a b. b -> Either a b
Right Conversation
conv

{- | Creates a new empty in-memory conversation store.

Initializes an 'InMemoryStore' with an empty 'Map' wrapped in a 'TVar' for thread-safe operations.

Example:

>>> store <- initInMemoryStore
>>> runInMemoryConvo store $ listConversations
[]
-}
initInMemoryStore :: IO InMemoryStore
initInMemoryStore :: IO InMemoryStore
initInMemoryStore = TVar (Map Text Conversation) -> InMemoryStore
InMemoryStore (TVar (Map Text Conversation) -> InMemoryStore)
-> IO (TVar (Map Text Conversation)) -> IO InMemoryStore
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map Text Conversation -> IO (TVar (Map Text Conversation))
forall a. a -> IO (TVar a)
newTVarIO Map Text Conversation
forall k a. Map k a
Map.empty