packages feed

A-gent-0.11.0.13: src/Agent/LLM.hs

{-# OPTIONS_GHC -Wall -Werror #-}

{-# LANGUAGE NoGeneralizedNewtypeDeriving #-}
{-# LANGUAGE Safe                         #-}

{-# LANGUAGE RankNTypes                   #-}

--------------------------------------------------------------------------------

-- |
-- Copyright  : (c) 2026 SPISE MISU ApS
-- License    : SSPL-1.0 OR AGPL-3.0-only
-- Maintainer : SPISE MISU <mail+hackage@spisemisu.com>
-- Stability  : experimental
--
-- Polite and well educated LLM agent with excellent manners that always behaves
-- well.

--------------------------------------------------------------------------------

module Agent.LLM
  ( -- * Modes
    Mode(..)
    -- * Context
  , Context(..)
  , FilePaths
  , Files
  , Filter
  , History
  , Load
  , UUID
    -- * Paramenters
  , Eval
    -- * Methods
  , repl
  , replWithMode
  )
where

--------------------------------------------------------------------------------

import           Prelude                    hiding ( head, mod, print, read )


import           Data.Char                  ( toLower, toUpper )
import           Data.Maybe                 ( fromMaybe )
import           GHC.IO.Encoding            ( setLocaleEncoding )
import           System.Environment         ( getProgName )
import qualified System.Environment.Blank   as ENV
import           System.Exit
  ( ExitCode (ExitFailure, ExitSuccess)
  )
import           System.IO
  ( BufferMode (LineBuffering, NoBuffering)
  , hFlush
  , hSetBuffering
  , hSetEcho
  , stderr
  , stdin
  , stdout
  , utf8
  )
import           System.Process             ( readProcessWithExitCode )
import           Text.Read                  ( readMaybe )

import           Internal.LLM
  ( Chit (..)
  , FilePaths
  , Files
  , Filter
  , History (..)
  , Mode
  , UUID
  , modes
  )
import qualified Internal.LLM               as INT
import           Internal.LLM.Action        ( Action (..) )
import           Internal.RIO               ( RIO (..), input, output )
import qualified Internal.Utils             as UTL


import qualified Agent.Data.ANSI.EscapeCode as AEC
import qualified Agent.LLM.Message          as MSG

import           Agent.Data.JSON            ( Data )
import           Agent.LLM.Message          ( Message )


--------------------------------------------------------------------------------

type Load a =
  ( Data a
  , Show a
  )
  => Maybe a

data Context a =
  Context
    { uuid :: !UUID
    , mode :: !Mode
    , load :: !(Load a)
    , hist :: !History
    , list :: !(Maybe Filter)
    , pile :: !(Maybe FilePaths)
    , ruck :: !(Maybe Files)
    }

type Eval a =
  Context a
  -> Message
  -> RIO (Context a, Action)

--------------------------------------------------------------------------------

repl
  :: Eval a
  -> IO ()
repl =
  replWithMode INT.Chat

replWithMode
  :: INT.Mode
  -> Eval a
  -> IO ()
replWithMode mod proc =
  do
    constraint <- integrity
    case constraint of
      Right __ ->
        do
          setLocaleEncoding utf8
          hSetBuffering stderr NoBuffering
          hSetBuffering stdin  NoBuffering
          -- NOTE: logic based on `hFlush stdout`
          hSetBuffering stdout LineBuffering
          -- NOTE: No default output when typing
          hSetEcho      stdin  False
          putStrLn head >> hFlush stdout
          run $ loop ctx proc
      Left err ->
        putStrLn err >> hFlush stdout
  where
    ctx =
      Context
        { uuid = INT.UUID []
        , mode = mod
        , load = Nothing
        , hist =
          History
            { chit =
                Chit
                  { prev = []
                  , next = []
                  }
            , chat = []
            }
        , list = Nothing
        , pile = Just $ INT.FilePaths []
        , ruck = Just $ INT.Files []
        }

--------------------------------------------------------------------------------

-- HELPERS (private)

{- NOTE: Doesn't seem to work
ctrlC :: IO () -> IO ()
ctrlC task =
  -- NOTE: Catch CTRL+C
  -- https://neilmitchell.blogspot.com/2015/05/handling-control-c-in-haskell.html
  sync . (:[]) <$> fork task
-}

head :: String
head =
  "# Exit Λ-gent with /e or /exit. For more commands, type /? or /help."

help :: String
help =
  unlines
    [ "# Supported commands:"
    , ds
    , "/?   or /help   | This message"
    , "/e   or /exit   | Exit Λ-gent"
    , "/w   or /wipe   | Clear screen"
    , ds
    , "/m m or /mode m | Change to {" ++ ms ++ "}. Ex: /m plan"
    , ds
    , "/h   or /hist   | View history (input and output)"
    , "        /chit   | View history (only input) "
    , "        /chat   | View history (only output)"
    , ds
    , "/l f or /list f | List of files, limited by mode, file masks and filter"
    , "/p   or /pile   | Show stored list of files. See /list above"
    , "/f i or /file i | Show file with the given index. See /pile above"
    , "/n i or /nums i | Show file with line numbers. See /file above"
    , ds
    , "/s m or /send m | Sending /pile to LLM as context to the message"
    , "/r o or /ruck o | View files (opt index) from LLM response. See /send"
    , "/a d or /atom d | Atomically save /ruck files to a GIT branch (+ description)"
    , "        /repo   | List GIT branches (+ description). See /atom above"
    , ds
    ]
  where
    ds = replicate 80 '-'
    ms =
      foldl1 (\ x y -> x ++ "|" ++ y) $ map (map toLower . show) modes

loop
  :: Context a
  -> Eval a
  -> RIO ()
loop ctx eval =
  print prompt >>
  read hps hns >>= \ txt ->
  printLn [ ]  >>
  case txt2act txt of
    -- TODO: Refactoring needed

    -- Nested -- Start
    None                      ->
      loop (nxt txt ctx Nothing) eval
    File _                    ->
      loop (nxt txt ctx Nothing) eval
    Template _ _ _ _         ->
      loop (nxt txt ctx Nothing) eval
    -- Nested -- Stop

    Exit                      -> caseExit
    Help                      -> caseHelp       txt
    Hist (chits, chats)       -> caseHist       txt chits chats
    Mode c cs                 -> caseMode       txt c cs
    Paths (INT.FilePaths fps) -> casePaths      txt fps
    Pile                      -> casePile       txt
    Ruck mdix                 -> caseRuck       txt mdix
    Wipe                      -> caseWipe       txt
    UnknownCmd cmd            -> caseUnknownCmd txt cmd

    Branch msg fs             ->
      eval ctx (MSG.Atom msg fs) >>= \ (upd, action) ->
      case action of
        Message (MSG.Text cs)         ->
          printLn cs >>
          loop (nxt txt upd (Just cs)) eval
        None                      ->
          loop (nxt txt upd Nothing) eval
        _________________________ ->
          -- NOTE: Unexpected error
          printLn "Unexpected error" >>
          loop (nxt txt upd (Just [])) eval

    Prompt msg                ->

      eval ctx (MSG.Tmpl afps) >>= \ (upd0, action0) ->
      case action0 of
        Template f is es fs ->

          eval ctx (MSG.Send xml) >>= \ (upd1, action1) ->
          case action1 of
            Message (MSG.Text cs)         ->
              printLn cs >>
              loop
                ( nxt txt
                    ( upd1 { ruck = Just $ pfs }
                    )
                    ( Just cs
                    )
                ) eval
              where
                pfs =
                  case (map INT.File . f) cs of
                    [] -> INT.Files []
                    xs -> INT.Files xs
            None                      ->
              loop (nxt txt upd1 Nothing) eval
            _________________________ ->
              -- NOTE: Unexpected error
              printLn "Unexpected error" >>
              loop (nxt txt upd1 (Just [])) eval
          where
            ils = map ( \ (INT.File (_, ls)) -> unlines ls) is
            els = map ( \ (INT.File (_, ls)) -> unlines ls) es
            tpl = INT.Template msg ils els fs
            xml = INT.tpl2xmls tpl

        Message (MSG.Text cs)         ->
          printLn cs >>
          loop (nxt txt upd0 (Just cs)) eval
        None                      ->
          loop (nxt txt upd0 Nothing) eval
        _________________________ ->
          -- NOTE: Unexpected error
          printLn "Unexpected error" >>
          loop (nxt txt upd0 (Just [])) eval
      where
        afps =
          ( \ (INT.FilePaths fps) -> map INT.AbsoluteFilePath fps
          )
          <$> pile ctx

    Message (MSG.Path ns afp)     ->
      eval ctx (MSG.Path ns afp) >>= \ (upd, action) ->
      case action of
        File (INT.File (_,ls)) ->
          caseFile txt ns ls
        Message (MSG.Text cs)         ->
          printLn cs >>
          loop (nxt txt upd (Just cs)) eval
        None                      ->
          loop (nxt txt upd Nothing) eval
        _________________________ ->
          -- NOTE: Unexpected error
          printLn "Unexpected error" >>
          loop (nxt txt upd (Just [])) eval

    Message (MSG.List fil)        ->
      eval ctx (MSG.List fil) >>= \ (upd, action) ->
      case action of
        Paths (INT.FilePaths fps) ->
          casePaths txt fps
        Message (MSG.Text cs)         ->
          printLn cs >>
          loop (nxt txt upd (Just cs)) eval
        None                      ->
          loop (nxt txt upd Nothing) eval
        _________________________ ->
          -- NOTE: Unexpected error
          printLn "Unexpected error" >>
          loop (nxt txt upd (Just [])) eval

    Message  MSG.Repo             ->
      eval ctx MSG.Repo >>= \ (upd, action) ->
      case action of
        Message (MSG.Text cs)         ->
          printLn cs >>
          loop (nxt txt upd (Just cs)) eval
        None                      ->
          loop (nxt txt upd Nothing) eval
        _________________________ ->
          -- NOTE: Unexpected error
          printLn "Unexpected error" >>
          loop (nxt txt upd (Just [])) eval

    Message       msg  ->
      eval ctx msg >>= \ (upd, action) ->
      case action of
        Message (MSG.Text cs)         ->
          printLn cs >>
          loop (nxt txt upd (Just cs)) eval
        None                      ->
          loop (nxt txt upd Nothing) eval
        _________________________ ->
          -- NOTE: Unexpected error
          printLn "Unexpected error" >>
          loop (nxt txt upd (Just [])) eval
  where

    -- TODO: Refactoring needed
    txt2act txt =
      -- TODO: Refactoring needed
      case txt of

        '/':'m':'o':'d':'e':' ':c:cs -> Mode c cs
        '/':'m'            :' ':c:cs -> Mode c cs

        '/':'l':'i':'s':'t'    :mfil ->
          -- TODO: DRY
          Message (MSG.List fil)
          where
            fil =
              case mfil of
                [    ] -> Just $ INT.Filter mfil
                ' ':cs -> Just $ INT.Filter cs
                ______ -> Nothing
        '/':'l'                :mfil ->
          -- TODO: DRY
          Message (MSG.List fil)
          where
            fil =
              case mfil of
                [    ] -> Just $ INT.Filter mfil
                ' ':cs -> Just $ INT.Filter cs
                ______ -> Nothing

        '/':'f':'i':'l':'e':' ':midx ->
          -- TODO: DRY
          case oidx of
            Nothing  ->
              Message (MSG.Path False Nothing)
            Just idx ->
              if 0 <= idx && idx < length pfs then
                Message (MSG.Path False (Just $ INT.AbsoluteFilePath $ pfs !! idx))
              else
                Message (MSG.Path False Nothing)
              where
                pfs = concatMap INT.filePaths $ pile ctx
          where
            oidx =
              case midx of
                [] -> Nothing
                cs -> readMaybe cs :: Maybe Int
        '/':'f'            :' ':midx ->
          -- TODO: DRY
          case oidx of
            Nothing  ->
              Message (MSG.Path False Nothing)
            Just idx ->
              if 0 <= idx && idx < length pfs then
                Message (MSG.Path False (Just $ INT.AbsoluteFilePath $ pfs !! idx))
              else
                Message (MSG.Path False Nothing)
              where
                pfs = concatMap INT.filePaths $ pile ctx
          where
            oidx =
              case midx of
                [] -> Nothing
                cs -> readMaybe cs :: Maybe Int
        '/':'n':'u':'m':'s':' ':midx ->
          -- TODO: DRY
          case oidx of
            Nothing  ->
              Message (MSG.Path True Nothing)
            Just idx ->
              if 0 <= idx && idx < length pfs then
                Message (MSG.Path True (Just $ INT.AbsoluteFilePath $ pfs !! idx))
              else
                Message (MSG.Path True Nothing)
              where
                pfs = concatMap INT.filePaths $ pile ctx
          where
            oidx =
              case midx of
                [] -> Nothing
                cs -> readMaybe cs :: Maybe Int
        '/':'n'            :' ':midx ->
          -- TODO: DRY
          case oidx of
            Nothing  ->
              Message (MSG.Path True Nothing)
            Just idx ->
              if 0 <= idx && idx < length pfs then
                Message (MSG.Path True (Just $ INT.AbsoluteFilePath $ pfs !! idx))
              else
                Message (MSG.Path True Nothing)
              where
                pfs = concatMap INT.filePaths $ pile ctx
          where
            oidx =
              case midx of
                [] -> Nothing
                cs -> readMaybe cs :: Maybe Int

        '/':'s':'e':'n':'d':' ':mesg -> Prompt mesg
        '/':'s'            :' ':mesg -> Prompt mesg

        '/':'a':'t':'o':'m':' ':mesg ->
          -- TODO: DRY
          Branch mesg fs
          where
            fs = fromMaybe (INT.Files []) (ruck ctx)
        '/':'a'            :' ':mesg ->
          -- TODO: DRY
          Branch mesg fs
          where
            fs = fromMaybe (INT.Files []) (ruck ctx)

        '/':'r':'e':'p':'o'    :[  ] -> Message MSG.Repo

        '/':'r':'u':'c':'k'    :midx -> Ruck   midx
        '/':'r'                :midx -> Ruck   midx

        "/exit"                      -> Exit
        "/e"                         -> Exit
        "/help"                      -> Help
        "/?"                         -> Help
        "/hist"                      -> Hist ( True,  True  )
        "/h"                         -> Hist ( True,  True  )
        "/chit"                      -> Hist ( True,  False )
        "/chat"                      -> Hist ( False, True  )
        "/pile"                      -> Pile
        "/p"                         -> Pile
        "/wipe"                      -> Wipe
        "/w"                         -> Wipe

        '/':cmd                      -> UnknownCmd cmd
        ____________________________ -> Message (MSG.Text txt)

    -- TODO: Refactoring needed
    his       = hist ctx
    chi       = chit his
    hps       = prev chi
    hns       = next chi
    nxt p c o =
      c
        { hist =
          (hist c)
            { chit =
              (chit $ hist c)
                { prev = p : hps
                }
            , chat =
              case o of
                Nothing ->        chat $ hist c
                Just  r -> (r:) $ chat $ hist c
            }
        }

    -- TODO: Refactoring needed
    caseMode txt c cs =
      case mmod of
        Just mod ->
          printLn ("Changed to " ++ low ++ "-mode") >>
          loop
            ( (nxt txt ctx Nothing)
                { mode = mod
                , pile = Just $ INT.FilePaths []
                , ruck = Just $ INT.Files     []
                }
            ) eval
          where
            low = map toLower $ show mod
        Nothing ->
          printLn ("Invalid mode: " ++ c:cs) >>
          loop ctx eval
      where
        mmod = readMaybe (toUpper c : map toLower cs) :: Maybe INT.Mode

    caseHist txt chits chats =
      ( if chits then
          printLn "* Chits:" >>
          ( mapM_ printLn . UTL.combine . UTL.chits True . prev . chit . hist
          ) ctx
        else
          pure ()
      ) >>
      ( if chats then
          printLn "* Chats:" >>
          ( mapM_ printLn . UTL.combine . UTL.chats True .        chat . hist
          ) ctx
        else
          pure ()
      ) >>
      loop (nxt txt ctx Nothing) eval

    caseExit =
      printLn "Λ-gent will shutdown" >>
      return ()
    caseFile txt ns ls =
      ( mapM_ printLn .
        UTL.combine .
        UTL.file ns
      ) (unlines ls) >>
      loop (nxt txt ctx Nothing) eval
    caseHelp txt =
      printLn help >>
      loop (nxt txt ctx Nothing) eval

    casePaths txt fps =
      eval ctx MSG.Root >>= \ (upd, action) ->
      case action of
        Paths (INT.FilePaths rfps) ->
          ( case rfps of
              [cwd] ->
                mapM_ printLn .
                UTL.combine .
                UTL.files [] True .
                -- NOTE: Show relative paths
                map (drop (length cwd)) $
                fps
              _____ ->
                printLn "Missing root path"
          ) >>
          loop
            ( (nxt txt upd Nothing)
                { pile = Just $ INT.FilePaths fps
                }
            ) eval
        Message (MSG.Text cs)         ->
          printLn cs >>
          loop (nxt txt upd (Just cs)) eval
        None                      ->
          loop (nxt txt upd Nothing) eval
        _________________________ ->
          -- NOTE: Unexpected error
          printLn "Unexpected error" >>
          loop (nxt txt upd (Just [])) eval

    casePile txt =
      eval ctx MSG.Root >>= \ (upd, action) ->
      case action of
        Paths (INT.FilePaths fps) ->
          ( case fps of
              [cwd] ->
                mapM_ printLn .
                UTL.combine .
                UTL.files [] True .
                -- NOTE: Show relative paths
                map (drop (length cwd)) .
                concatMap INT.filePaths .
                pile $
                upd
              _____ ->
                printLn "Missing root path"
          ) >>
          loop (nxt txt upd Nothing) eval
        Message (MSG.Text cs)         ->
          printLn cs >>
          loop (nxt txt upd (Just cs)) eval
        None                      ->
          loop (nxt txt upd Nothing) eval
        _________________________ ->
          -- NOTE: Unexpected error
          printLn "Unexpected error" >>
          loop (nxt txt upd (Just [])) eval

    caseRuck txt midx =
      ( if null midx then
          -- NOTE: If no optional index provided, show only paths
          ( mapM_ printLn .
            UTL.combine .
            UTL.files [] True .
            concatMap fps .
            ruck
          ) ctx
        else
          -- NOTE: Otherwise, show content or index error
          ( case idx of
              Just i ->
                if 0 <= i && i < len then
                  ( mapM_ printLn .
                    UTL.combine .
                    UTL.file True .
                    concatMap f .
                    ruck
                  ) ctx
                else
                  let
                    m = "Index (" ++ show i ++ ") is out of bounds."
                  in
                    printLn m
                where
                  f (INT.Files fs) =
                    case fs !! i of
                      INT.File (_,ls) -> unlines ls
              Nothing ->
                printLn "Invalid use of /ruck. Use space between cmd and index"
          )
      ) >>
      loop (nxt txt ctx Nothing) eval
      where
        idx =
          case midx of
            [    ] -> Nothing
            ' ':cs -> readMaybe cs :: Maybe Int
            ______ -> Nothing
        len =
          case ruck ctx of
            Just (INT.Files fs) -> length fs
            Nothing             -> 0
        fps (INT.Files fs) =
          map ( \ (INT.File (p,_)) -> p) fs

    caseWipe txt =
      printLn clear >>
      loop (nxt txt ctx Nothing) eval
    caseUnknownCmd txt cmd =
      printLn msg >>
      loop (nxt txt ctx Nothing) eval
      where
        msg = "Command not recognized: " ++ cmd

    -- NOTE: Clear screen & Move top-left
    clear     = "\^[[H\^[[2J" -- NOTE: See `infocmp -x`
    prompt    =
      (show . AEC.bold . AEC.sgr)
      ("Λ-" ++ (map toLower . show . mode) ctx ++ "> ")
    read      = input
    print     = output . show . AEC.faint . AEC.sgr
    printLn x = print $ x ++ "\n"

integrity :: IO (Either String ())
integrity =
  do
    ohi <- ENV.getEnv "LLM_SIGN_SHA"
    case ohi of
      Nothing -> pure $ Right ()
      Just hi ->
        do
          op <- cwd
          pn <- getProgName
          case op of
            Left  e -> pure $ Left e
            Right p ->
              do
                ehf <- sha ap
                case ehf of
                  Left  ef -> pure $ Left ef
                  Right hf ->
                    if hi == h64 then
                      pure $ Right ()
                    else
                      pure $
                      Left $
                      "# Λ-gent integrity failed:\n" ++
                      "> " ++ hi  ++ " (expected)\n" ++
                      "> " ++ h64 ++ " (observed)"
                    where
                      h64 = take 64 hf
              where
                ap =
                  concatMap id
                  $ lines p ++ [ "/" ] ++ lines pn
  where
    -- TODO: Move "HELPERS (private)" private out from Internal/RIO.hs to
    -- Internal/IO.hs and then just wrap them with `RestrictedIO`. This will
    -- allow for logic like this to be simplified
    cwd =
      do
        (exitcode, out, err) <- readProcessWithExitCode "pwd" [] []
        case exitcode of
          ExitSuccess ->
            pure $ Right out
          ExitFailure _ ->
            pure $ Left err
    sha p =
      do
        (exitcode, out, err) <- readProcessWithExitCode "sha256sum" [ p ] []
        case exitcode of
          ExitSuccess ->
            pure $ Right out
          ExitFailure _ ->
            pure $ Left err