A-gent-0.11.0.20: src/Internal/LLM/SystemOne.hs
{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE NoGeneralizedNewtypeDeriving #-}
{-# LANGUAGE Safe #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# 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
--------------------------------------------------------------------------------
module Internal.LLM.SystemOne
( --
RequestChoice(..)
, State(..)
, TypedHole(..)
, Question(..)
, Criteria(..)
--
, ResponseChoice(..)
, Answer(..)
, Probability(..)
, Usage(..)
--
, request
, response
--
, encode
, decode
--
, replaceHoles
)
where
--------------------------------------------------------------------------------
import Prelude hiding ( lines )
import Data.Bifunctor ( bimap )
import Data.Data ( Data )
import Data.Either ( partitionEithers )
import Data.List
( groupBy
, nub
, sort
, sortOn
)
import Data.Maybe ( fromMaybe, isJust )
import Text.Read ( readMaybe )
import Internal.GaloisInc.Text.JSON
( JSON
, readJSON
, showJSON
, showJSValue
)
import Internal.GaloisInc.Text.JSON.Generic
( Result (Error, Ok)
, fromJSON
, toJSON
)
import Internal.GaloisInc.Text.JSON.String
( readJSValue
, runGetJSON
)
import Internal.GaloisInc.Text.JSON.Types
( JSObject (..)
, JSString (..)
, JSValue (..)
, fromJSString
, toJSObject
, toJSString
)
import qualified Internal.LLM as LLM
--------------------------------------------------------------------------------
{- Sample:
{
"state": {
"filepath": "src/Foo.hs",
"lines": [
…
],
"typedholes": [
{
"line": 38,
"column": 15,
"hole": "_x",
"details": "_x :: Char Or perhaps ‘_x’ is mis-spelled, or not in scope",
"description": [
"In an equation for ‘foo’: foo (Foo x) = _x"
],
"bindings": [
"x :: Char",
"foo :: Foo -> Char"
],
"constraints": [],
"mostLikely": [
"x"
]
}
]
},
"questions": {
"0": {
"type": "choice",
"instructions": "Substitute the 'typed-hole' from src/Foo.hs at line 38 and column 15 with? (hint: 'mostLikely' value, if any)",
"criteria": {
"0": {
"sig": "x :: Char",
"val": "x"
},
"1": {
"sig": "foo :: Foo -> Char",
"val": "(foo (_ :: Foo))"
}
}
}
}
}
-}
data RequestChoice = RequestChoice
{ state :: !State
, questions :: ![(Int, Question)]
} deriving Show
data State = State
{ filepath :: !FilePath
, lines :: ![String]
, typedholes :: ![TypedHole]
} deriving (Data, Show)
data TypedHole = TypedHole
{ line :: !Int
, column :: !Int
, hole :: !String
, details :: !String
, description :: ![String]
, bindings :: ![String]
, constraints :: ![String]
, mostLikely :: ![String]
} deriving (Data, Show)
data Question = Question
{ type_ :: !String
, instructions :: !String
, criteria :: ![(Int, Criteria)]
} deriving Show
data Criteria = Criteria
{ sig :: !String
, val :: !String
} deriving Show
--------------------------------------------------------------------------------
{- Sample:
{
"model": "sm-mockup",
"answers": {
"0": {
"type": "choice",
"choice": "0",
"confidence": 0.9651,
"probabilities": {
"0": 0.9651,
"1": 0.0182,
"2": 0.0112,
"3": 0.0055
}
}
},
"usage": {
"input_tokens": 0,
"output_tokens": 0,
"latency_ms": 1000
}
}
-}
data ResponseChoice = ResponseChoice
{ model :: !String
, answers :: ![(Int, Answer)]
, usage :: !Usage
} deriving Show
data Answer = Answer
{ type_ :: !String
, choice :: !String
, confidence :: !Double
, probabilities :: ![(Int, Probability)]
} deriving Show
newtype Probability = Probability Double
deriving Show
data Usage = Usage
{ input_tokens :: !Int
, output_tokens :: !Int
, latency_ms :: !Double
} deriving (Data, Show)
--------------------------------------------------------------------------------
instance JSON RequestChoice where
readJSON (JSObject o) =
case sortOn fst (fromJsobject o) of
[ ( "questions", JSObject q )
, ( "state", s ) ] ->
case (res2either . fromJSON) s of
Right sv ->
if null ls && all isJust os then
Ok $ RequestChoice sv (zip is rs)
else
Error "Unable to read RequestChoice"
________ ->
Error "Unable to read RequestChoice"
where
ms = fromJsobject q
ks = map fst ms
os = map (readMaybe :: String -> Maybe Int) ks
is = map (fromMaybe 0) os
vs = map snd ms
(ls,rs) = partitionEithers $ map (res2either . readJSON) vs
_____________________________________ ->
Error "Unable to read RequestChoice"
readJSON _________________ =
Error "Unable to read RequestChoice"
showJSON (RequestChoice s qs) =
toJsobject
[ ( "state", toJSON s )
, ( "questions", (toJsobject . map (bimap show showJSON)) qs )
]
instance JSON Question where
readJSON (JSObject o) =
case sortOn fst (fromJsobject o) of
[ ( "criteria", JSObject c )
, ( "instructions", JSString i )
, ( "type", JSString t ) ] ->
if null ls && all isJust os then
Ok
$ Question
(fromJsstring t)
(fromJsstring i)
(zip is rs)
else
Error "Unable to read Question"
where
ms = fromJsobject c
ks = map fst ms
os = map (readMaybe :: String -> Maybe Int) ks
is = map (fromMaybe 0) os
vs = map snd ms
(ls,rs) = partitionEithers $ map (res2either . readJSON) vs
_____________________________________ ->
Error "Unable to read Question"
readJSON _________________ =
Error "Unable to read Question"
showJSON (Question t i cs) =
toJsobject
[ ( "type", toJsstring t )
, ( "instructions", toJsstring i )
, ( "criteria", (toJsobject . map (bimap show showJSON)) cs )
]
instance JSON Criteria where
readJSON (JSObject o) =
case sortOn fst (fromJsobject o) of
[ ( "sig", JSString s )
, ( "val", JSString v ) ] ->
Ok $ Criteria { sig = fromJsstring s, val = fromJsstring v }
_____________________________________ ->
Error "Unable to read Criteria"
readJSON _________________ =
Error "Unable to read Criteria"
showJSON (Criteria { sig = s, val = v }) =
toJsobject
[ ( "sig", toJsstring s )
, ( "val", toJsstring v )
]
--------------------------------------------------------------------------------
instance JSON ResponseChoice where
readJSON (JSObject o) =
case sortOn fst (fromJsobject o) of
[ ( "answers", JSObject a )
, ( "model", JSString m )
, ( "usage", u ) ] ->
case (res2either . fromJSON) u of
Right uv ->
if null ls && all isJust os then
Ok
$ ResponseChoice
(fromJsstring m)
(zip is rs) uv
else
Error "Unable to read ResponseChoice"
________ ->
Error "Unable to read ResponseChoice"
where
ms = fromJsobject a
ks = map fst ms
os = map (readMaybe :: String -> Maybe Int) ks
is = map (fromMaybe 0) os
vs = map snd ms
(ls,rs) = partitionEithers $ map (res2either . readJSON) vs
_____________________________________ ->
Error "Unable to read ResponseChoice"
readJSON _________________ =
Error "Unable to read ResponseChoice"
showJSON (ResponseChoice m as u) =
toJsobject
[ ( "model", toJsstring m )
, ( "answers", (toJsobject . map (bimap show showJSON)) as )
, ( "usage", toJSON u )
]
instance JSON Answer where
readJSON (JSObject o) =
case sortOn fst (fromJsobject o) of
[ ( "choice", JSString ch )
, ( "confidence", JSRational _ co )
, ( "probabilities", JSObject ps )
, ( "type", JSString t ) ] ->
if null ls && all isJust os then
Ok
$ Answer
(fromJsstring t)
(fromJsstring ch)
(rat2double co)
(zip is rs)
else
Error "Unable to read Answer"
where
ms = fromJsobject ps
ks = map fst ms
os = map (readMaybe :: String -> Maybe Int) ks
is = map (fromMaybe 0) os
vs = map snd ms
(ls,rs) = partitionEithers $ map (res2either . readJSON) vs
_____________________________________ ->
Error "Unable to read Answer"
readJSON _________________ =
Error "Unable to read Answer"
showJSON (Answer t ch co ps) =
toJsobject
[ ( "type", toJsstring t )
, ( "choice", toJsstring ch )
, ( "confidence", (toJsrational . double2rat) co )
, ( "probabilities", (toJsobject . map (bimap show showJSON)) ps )
]
instance JSON Probability where
readJSON (JSRational _ r) =
(Ok . Probability . rat2double) r
readJSON _____________________ =
Error "Unable to read Probability"
showJSON (Probability r) =
(toJsrational . double2rat) r
--------------------------------------------------------------------------------
toJsobject :: [(String, JSValue)] -> JSValue
toJsobject = JSObject . toJSObject
toJsrational :: Rational -> JSValue
toJsrational = JSRational True
toJsstring :: String -> JSValue
toJsstring = JSString . toJSString
fromJsobject :: JSObject e -> [(String, e)]
fromJsobject = fromJSObject
fromJsstring :: JSString -> String
fromJsstring = fromJSString
--------------------------------------------------------------------------------
rat2double :: Rational -> Double
rat2double = fromRational
double2rat :: Double -> Rational
double2rat = toRational
--------------------------------------------------------------------------------
res2either :: Result a -> Either String a
res2either (Error e) = Left e
res2either (Ok a) = Right a
--------------------------------------------------------------------------------
encodeHlp :: JSON a => a -> String
encodeHlp = flip showJSValue [] . showJSON
decodeHlp :: JSON a => String -> Result a
decodeHlp s =
case runGetJSON readJSValue s of
Right a -> readJSON a
Left err -> Error err
--------------------------------------------------------------------------------
request :: LLM.File -> [LLM.Hole] -> RequestChoice
request (LLM.File (fp, ls)) hs =
RequestChoice s (zip [0..] qs)
where
s =
State
{ filepath = fp
, lines = ls
, typedholes =
sortOn
( \ TypedHole { line = l, column = c} -> (l, c)
)
$ map
( \ ( LLM.Hole (LLM.Location _ l c) (LLM.Signature hn hs_) d bs cs fs rs
) ->
let
h = hn ++ " :: " ++ hs_
u =
case hn of
[ ] -> []
['_'] -> []
__:xs -> xs
in
TypedHole
{ line = l
, column = c
, hole = hn
, details = h
, description = d
, mostLikely =
( nub
. sort
. filter (u ==)
)
( map (\ (n,_,_) -> n) fs ++ map (\ (n,_,_) -> n) rs
)
, bindings =
map
( \ (n, s_) ->
n ++ " :: " ++ s_
)
$ filter ((u ==) . fst) bs ++ filter ((u /=) . fst) bs
, constraints =
map
( \ (n, s_) ->
n ++ " :: " ++ s_
)
cs
}
)
hs
}
qs =
map
( \ ( LLM.Hole (LLM.Location fp_ l c) _ _ _ _ fs rs
) ->
let
is =
-- "With what do we replace the typed-hole from " ++ fp_ ++
-- "What should go instead of the typed-hole from " ++ fp_ ++
-- "How do we substitute the typed-hole from " ++ fp_ ++
"Substitute the 'typed-hole' from " ++ fp_ ++
" at line " ++ show l ++
" and column " ++ show c ++
-- " with? (have in mind the value from 'mostLikely', if any)"
" with? (hint: 'mostLikely' value, if any)"
in
Question
{ type_ = "choice"
, instructions = is
, criteria =
let
-- NOTE: Sometimes a typed-hole can show up both as a fit
-- and a refinement. It's the same logic, however, with
-- point-free-style, functions can be called without
-- (apparent) arguments
fs_ =
map
( \ (c_, s_, _) ->
Criteria
{ sig = c_ ++ " :: " ++ s_
, val = c_
}
)
fs
rs_ =
map
( \ (n, c_, s_) ->
Criteria
{ sig = s_
, val = "(" ++ n ++ c_ ++ ")"
}
)
rs
in
zip [0..] (fs_ ++ rs_)
}
)
hs
response :: ResponseChoice -> [(Int, Int, Double)]
response res =
map
( \ (i, a) ->
let
j = (read . choice) a :: Int
p = confidence a
in
(i, j, p)
)
as
where
as = answers res
--------------------------------------------------------------------------------
encode :: RequestChoice -> String
encode =
encodeHlp
decode :: String -> Either String ResponseChoice
decode =
( \ case
Ok a -> Right a
Error e -> Left e
)
<$> decodeHlp
--------------------------------------------------------------------------------
replaceHolesHlp
:: [String]
-> [(String, Int, Int)]
-> [[String]]
-> [Int]
-> [String]
replaceHolesHlp lines_ holes_ questions_ answers_ =
lin 1 lines_ gbl
where
hqa = zip3 holes_ questions_ answers_
-- NOTE: Assuming holes_ are sorted ascending by lines
gbl = groupBy (\ ((_,i,_),_,_) ((_,j,_),_,_) -> i == j) hqa
--
lin _ [] ______ = []
lin _ ls [ ] = ls
lin m ls (b:bs) =
case b of
[ ] -> ls
(((_,i,_),_,_):_) ->
case ds of
[ ] -> ts
(l:rs) -> ts ++ col 1 0 b l : lin j rs bs
where
j = i + 1
--
k = i - m
--
ts = take k ls
ds = drop k ls
--
col _ _ [ ] cs = cs
col _ _ ___________________ [] = []
col n o (((h,_,i),qs,a):rs) cs = t ++ v ++ col i l rs d
where
v = qs !! a
--
k = i - n - o
--
l = length h
--
t = take k cs
d = drop (k + l) cs
replaceHoles
:: RequestChoice
-> ResponseChoice
-> LLM.File
replaceHoles req@(RequestChoice s@(State fp ls _) _) res =
LLM.File
( fp
, replaceHolesHlp ls hs qs as
)
where
hs =
( map
( \ (TypedHole l c h _ _ _ _ _) ->
(h, l, c)
)
. typedholes
) s
qs =
( map
( \ (_, Question _ _ cs) ->
map
( \ (_, Criteria { val = v }) ->
v
)
cs
)
. questions
) req
as =
( map
( \ (_, Answer _ c _ _) ->
read c :: Int
)
. answers
) res