packages feed

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

{-# OPTIONS_GHC -Wall -Werror #-}

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

{-# LANGUAGE RankNTypes                   #-}

{-# LANGUAGE LambdaCase                   #-}

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

-- |
-- 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
    INT.Mode(..)
    -- * Context
  , Load
  , History(..)
  , Context(..)
  , INT.files
    -- * Paramenters
  , Eval
    -- * Methods
  , repl
  , replWithMode
  , rootPath
    -- * Helpers
  , combineFiles
  , getFile
  )
where

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

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

import           Data.Char                            ( toLower, toUpper )
import           GHC.IO.Encoding                      ( setLocaleEncoding )
import           System.IO
  ( BufferMode (LineBuffering, NoBuffering)
  , hFlush
  , hSetBuffering
  , hSetEcho
  , stderr
  , stdin
  , stdout
  , utf8
  )
import           Text.Read                            ( readMaybe )

import           Internal.GaloisInc.Text.JSON.Generic ( Data )
import           Internal.LLM
  ( Chit (..)
  , History (..)
  , Template (..)
  , View (..)
  , modes
  )
import qualified Internal.LLM                         as INT
import           Internal.RIO
  ( RIO (..)
  , findFiles
  , input
  , output
  , readFileStrict
  )

import qualified Agent.Data.ANSI.EscapeCode           as AEC
import qualified Agent.Data.JSON                      as JSON
import qualified Agent.Utils.Common                   as COM

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

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

data Context a =
  Context
    { exit :: !INT.Exit
    , root :: !(Maybe INT.Root)
    , mode :: !INT.Mode
    , load :: !(Load a)
    , hist :: !INT.History
    , list :: !(Maybe INT.Filter)
    , pile :: !(Maybe INT.Files)
    , file :: !(Maybe INT.File)
    }

type Eval a =
  Context a
  -> String
  -> RIO (Context a, String)

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

rootPath
  :: RIO (Maybe INT.Root)
  -> IO  (Maybe INT.Root)
rootPath =
  run

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

replWithMode
  :: INT.Mode
  -> Eval a
  -> Maybe INT.Root
  -> IO ()
replWithMode mod proc mroot =
  do
    setLocaleEncoding utf8
    hSetBuffering stderr NoBuffering
    hSetBuffering stdin  NoBuffering
    hSetBuffering stdout LineBuffering -- NOTE: logic based on `hFlush stdout`
    -- NOTE: No default output when typing
    hSetEcho      stdin  False
    putStrLn head >> hFlush stdout
    run $ loop ctx proc
    where
      ctx =
        Context
          { exit = INT.Exit False
          , root = mroot
          , mode = mod
          , load = Nothing
          , hist =
            History
              { view = []
              , chit =
                Chit
                  { prev = []
                  , next = []
                  }
              , chat = []
              }
          , list = Nothing
          , pile = Just $ INT.Files [] -- Nothing
          , file = Nothing
          }

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

-- HELPERS (public)

combineFiles
  :: INT.Files
  -> INT.Files
  -> INT.Files
combineFiles (INT.Files xs) (INT.Files ys) =
  INT.Files $ xs ++ ys

getFile
  :: INT.File
  -> (FilePath, [INT.FileLine])
getFile (INT.File f) =
  f

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

-- 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
-}

loop
  :: Context a
  -> Eval a
  -> RIO ()
loop ctx eval =
  case ctx of
    Context { exit = INT.Exit True }                    ->
      return ()
    Context { hist = History { view = [ Chits, Chats ] } } ->
      printLn "* Chits:" >>
      ( mapM_ printLn . COM.combine . COM.chits True . prev . chit . hist
      ) ctx >>
      printLn "* Chats:" >>
      ( mapM_ printLn . COM.combine . COM.chats True .        chat . hist
      ) ctx >>
      loop (ctx { hist = ((hist ctx) { view = [] }) }) eval
    Context { hist = History { view = [ Chits ] } } ->
      printLn "* Chits:" >>
      ( mapM_ printLn . COM.combine . COM.chits True . prev . chit . hist
      ) ctx >>
      loop (ctx { hist = ((hist ctx) { view = [] }) }) eval
    Context { hist = History { view = [ Chats ] } } ->
      printLn "* Chats:" >>
      ( mapM_ printLn . COM.combine . COM.chats True .        chat . hist
      ) ctx >>
      loop (ctx { hist = ((hist ctx) { view = [] }) }) eval
    Context { list = Just _, pile = Just _ } ->
      eval ctx [ ] >>= \ (upd, _) ->
      ( mapM_ printLn .
        COM.combine .
        COM.files [] True .
        -- NOTE: Show relative paths
        map (drop (length cwd)) .
        concatMap INT.files .
        pile
      ) upd >>
      loop (upd { list = Nothing }) eval
      where
        cwd =
          case root ctx of
            Just (INT.Root cs) -> cs
            Nothing            -> []
    Context { file = Just _ }   ->
      eval    ctx [ ] >>= \ (upd, res) ->
      printLn res     >>
      loop (upd { file = Nothing }) eval
    __________________________________________ ->
      print prompt >>
      read hps hns >>= \ txt ->
      printLn [ ]  >>
      case txt of
        '/':'m':'o':'d':'e':' ':c:cs -> caseMode txt c cs
        '/':'m'            :' ':c:cs -> caseMode txt c cs
        '/':'s':'e':'n':'d':' ':mesg -> caseSend txt mesg
        '/':'s'            :' ':mesg -> caseSend txt mesg
        '/':'f':'i':'l':'e'    :midx -> caseFile txt midx False
        '/':'f'                :midx -> caseFile txt midx False
        '/':'n':'u':'m':'s'    :midx -> caseFile txt midx True
        '/':'n'                :midx -> caseFile txt midx True
        '/':'l':'i':'s':'t'    :mfil -> caseList txt mfil
        '/':'l'                :mfil -> caseList txt mfil
        "/pile"                      -> casePile txt
        "/p"                         -> casePile txt
        "/hist"                      -> caseHist txt [ Chits, Chats ]
        "/h"                         -> caseHist txt [ Chits, Chats ]
        "/chit"                      -> caseHist txt [ Chits        ]
        "/chat"                      -> caseHist txt [        Chats ]
        "/wipe"                      -> caseWipe txt
        "/w"                         -> caseWipe txt
        "/help"                      -> caseHelp txt
        "/?"                         -> caseHelp txt
        "/exit"                      -> caseExit txt
        "/e"                         -> caseExit txt
        '/':cmd                      ->
          printLn msg >>
          loop (nxt txt ctx Nothing) eval
          where
            msg = "Command not recognized: " ++ cmd
        ____________________________ ->
          eval    ctx xml >>= \ (upd, res) ->
          printLn res     >>
          loop (nxt txt upd (Just res)) eval
          where
            tpl =
              Template
                { purpose      = txt
                , instructions = []
                , examples     = []
                , documents    = []
                }
            xml = INT.tpl2xml tpl
      where
        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
                }
            }
        caseMode txt c cs =
          case mmod of
            Just mod ->
              printLn ("Changed to " ++ low ++ "-mode") >>
              loop (nxt txt (ctx { mode = mod, pile = Nothing }) Nothing) eval
              where
                low = map toLower $ show mod
            Nothing ->
              printLn ("Invalid mode: " ++ c:cs) >>
              loop (nxt txt ctx Nothing) eval
          where
            mmod = readMaybe (toUpper c : map toLower cs) :: Maybe INT.Mode
        caseSend txt msg =
          pds                           >>= \  ds ->
          tis                           >>= \  is ->
          eval ctx (xml ds (concat is)) >>= \ (upd, res) ->
          printLn res                   >>
          loop (nxt txt upd (Just res)) eval
          where
            cwd =
              case root ctx of
                Just (INT.Root cs) -> cs
                Nothing            -> []
            pds =
              ( mapM
                  ( \ fap ->
                      INT.Document .
                      INT.File .
                      (,) (drop (length cwd) fap) .
                      lines
                      <$> readFileStrict fap
                  ) .
                concatMap INT.files .
                pile
              ) ctx

            tis =
              -- TODO: This need to be refactored
              findFiles cwd ("*/tpl/" ++ mod ++ "_instructions_*.json") >>= \ efs ->
              case efs of
                Right aps ->
                  mapM (\ ap -> dec <$> readFileStrict ap) aps
                Left  ___ ->
                  pure []
              where
                mod   = (map toLower . show . mode) ctx
                dec x =
                  case JSON.decode x :: Either JSON.DecodeError [String] of
                    Right xs -> xs
                    Left  __ -> []
            tpl ds is =
              Template
                { purpose      = msg
                , instructions = is
                , examples     =
                  -- TODO: Look in {root}/tpl/{mode}_examples_*.md
                  []
                , documents    = ds
                }
            xml xs ys = INT.tpl2xml $ tpl xs ys

        caseFile txt midx nums =
          ( case idx of
              Just i ->
                -- TODO: Test with empty list/pile
                if 0 <= i && i < length pfs then
                  let
                    cwd =
                      case root ctx of
                        Just (INT.Root cs) -> cs
                        Nothing            -> []
                    fap = pfs !! i
                  in
                    readFileStrict fap >>= \ f ->
                    ( mapM_ printLn .
                      COM.combine .
                      COM.file nums
                    ) f >>
                    loop
                      ( nxt txt
                        ( ctx
                          { file =
                            Just
                              ( INT.File
                                -- NOTE: Store relative path
                                ( drop (length cwd) fap
                                , lines f
                                )
                              )
                          }
                        ) Nothing
                      ) eval
                else
                  let
                    m = "Index (" ++ show i ++ ") is out of bounds."
                  in
                    printLn m >>
                    loop (nxt txt ctx Nothing) eval
              Nothing ->
                printLn "Invalid index" >>
                loop (nxt txt ctx Nothing) eval
          )
          where
            pfs = concatMap INT.files $ pile ctx
            idx =
              case midx of
                [    ] -> Nothing
                ' ':cs -> readMaybe cs :: Maybe Int
                ______ -> Nothing
        caseList txt mfil =
          ( case fil of
              Nothing ->
                printLn "Invalid use of list. Use space between cmd and filter"
              _______ ->
                pure ()
          ) >>
          loop (nxt txt (ctx { list = fil }) Nothing) eval
          where
            fil =
              case mfil of
                [    ] -> Just $ INT.Filter mfil
                ' ':cs -> Just $ INT.Filter cs
                ______ -> Nothing
        casePile txt =
          ( mapM_ printLn .
            COM.combine .
            COM.files [] True .
            -- NOTE: Show relative paths
            map (drop (length cwd)) .
            concatMap INT.files .
            pile
          ) ctx >>
          loop (nxt txt ctx Nothing) eval
          where
            cwd =
              case root ctx of
                Just (INT.Root cs) -> cs
                Nothing            -> []
        caseHist txt v =
          loop (nxt txt (ctx { hist = his { view = v } }) Nothing) eval
        caseWipe txt =
          printLn clear >>
          loop (nxt txt ctx Nothing) eval
        caseHelp txt =
          printLn help >>
          loop (nxt txt ctx Nothing) eval
        caseExit txt =
          printLn "Λ-gent will shutdown" >>
          loop (nxt txt (ctx { exit = INT.Exit True }) Nothing) eval
    where
      -- 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"

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

help :: String
help =
  "# Supported commands:\n" ++
  "/?   or /help   | This message\n" ++
  "/e   or /exit   | Exit Λ-gent\n" ++
  "/w   or /wipe   | Clear screen\n" ++
  "/m m or /mode m | Change to {" ++ ms ++ "}. Ex: /m code\n" ++
  "/h   or /hist   | View chit-chat (input and output) history\n" ++
  "        /chit   | View chit (input) history\n" ++
  "        /chat   | View chat (output) history\n" ++
  "/l f or /list f | List of files, limited by mode, file masks and filter\n" ++
  "/p   or /pile   | Show stored list of files. See /list above\n" ++
  "/f i or /file i | Show file with the given index. See, /pile above\n" ++
  "/n i or /nums i | Show file with line numbers. See /file above\n" ++
  "/s m or /send m | Sending /pile to LLM as context to the message"
  where
    ms =
      foldl1 (\ x y -> x ++ "|" ++ y) $ map (map toLower . show) modes