MLMemorySpec.hs

{-# LANGUAGE OverloadedStrings #-}

module Test.MLMemorySpec (spec) where

import Data.ByteString.Lazy.Char8 qualified as LBS
import Database.Bloodhound.OpenSearch3.Types
import TestsUtils.Import
import Prelude

-- | Create memory response: @{memory_id}@.
sampleCreateMemoryResponse :: LBS.ByteString
sampleCreateMemoryResponse = "{\"memory_id\":\"gW8Aa40BfUsSoeNTvOKI\"}"

-- | Get memory response. Note timestamps are RFC 3339 strings.
sampleGetMemory :: LBS.ByteString
sampleGetMemory =
  "{\
  \  \"memory_id\": \"ZnQ0JY0BfUsSoeNTqfN3\",\
  \  \"create_time\": \"2024-02-02T18:07:06.887061463Z\",\
  \  \"updated_time\": \"2024-02-02T18:07:06.887061463Z\",\
  \  \"name\": \"Conversation for a RAG pipeline\",\
  \  \"user\": \"admin\"\
  \}"

-- | List memories response: @{"memories": [...]}@.
sampleListMemories :: LBS.ByteString
sampleListMemories = "{\"memories\": []}"

-- | Delete memory response: @{success: <bool>}@.
sampleDeleteMemory :: LBS.ByteString
sampleDeleteMemory = "{\"success\":true}"

-- | Create message response: @{message_id}@.
sampleCreateMessageResponse :: LBS.ByteString
sampleCreateMessageResponse = "{\"message_id\":\"sxa2cY0BfUsSoeNTz-8m\"}"

-- | A message with trace fields (@parent_message_id@, @trace_number@).
sampleTraceMessage :: LBS.ByteString
sampleTraceMessage =
  "{\
  \  \"memory_id\": \"m1\",\
  \  \"message_id\": \"t1\",\
  \  \"parent_message_id\": \"p1\",\
  \  \"trace_number\": 0,\
  \  \"response\": \"intermediate tool output\"\
  \}"

-- | Get traces response: @{"traces": [...]}@.
sampleTraces :: LBS.ByteString
sampleTraces = "{\"traces\": []}"

spec :: Spec
spec = describe "ML Commons memory and message APIs" $ do
  describe "Memory" $ do
    it "decodes the doc get sample (RFC 3339 timestamps as opaque text)" $
      decode sampleGetMemory
        `shouldBe` Just
          Memory
            { memoryId = Just "ZnQ0JY0BfUsSoeNTqfN3",
              memoryCreateTime = Just "2024-02-02T18:07:06.887061463Z",
              memoryUpdatedTime = Just "2024-02-02T18:07:06.887061463Z",
              memoryName = Just "Conversation for a RAG pipeline",
              memoryUser = Just "admin",
              memoryApplicationType = Nothing
            }
    it "round-trips a sparse memory" $ do
      let m = Memory {memoryId = Just "x", memoryCreateTime = Nothing, memoryUpdatedTime = Nothing, memoryName = Just "n", memoryUser = Nothing, memoryApplicationType = Nothing}
      decode (encode m) `shouldBe` Just m

  describe "CreateMemoryResponse" $
    it "decodes the doc sample" $
      decode sampleCreateMemoryResponse
        `shouldBe` Just (CreateMemoryResponse {createMemoryResponseId = Just "gW8Aa40BfUsSoeNTvOKI"})

  describe "ListMemoriesResponse" $
    it "decodes an empty list envelope" $
      decode sampleListMemories `shouldBe` Just (ListMemoriesResponse {listMemoriesResponseMemories = []})

  describe "DeleteMemoryResponse" $
    it "decodes {success: true}" $
      decode sampleDeleteMemory `shouldBe` Just (DeleteMemoryResponse {deleteMemoryResponseSuccess = Just True})

  describe "MemoryMessage" $ do
    it "decodes a trace message (parent_message_id + trace_number present)" $
      decode sampleTraceMessage
        `shouldBe` Just
          MemoryMessage
            { memoryMessageMemoryId = Just "m1",
              memoryMessageId = Just "t1",
              memoryMessageCreateTime = Nothing,
              memoryMessageUpdatedTime = Nothing,
              memoryMessageInput = Nothing,
              memoryMessagePromptTemplate = Nothing,
              memoryMessageResponse = Just "intermediate tool output",
              memoryMessageOrigin = Nothing,
              memoryMessageAdditionalInfo = Nothing,
              memoryMessageParentMessageId = Just "p1",
              memoryMessageTraceNumber = Just 0
            }
    it "round-trips a message with additional_info" $ do
      let msg =
            MemoryMessage
              { memoryMessageMemoryId = Just "m",
                memoryMessageId = Just "x",
                memoryMessageCreateTime = Nothing,
                memoryMessageUpdatedTime = Nothing,
                memoryMessageInput = Just "hello",
                memoryMessagePromptTemplate = Nothing,
                memoryMessageResponse = Just "world",
                memoryMessageOrigin = Just "bot",
                memoryMessageAdditionalInfo = Just (object ["k" .= (1 :: Int)]),
                memoryMessageParentMessageId = Nothing,
                memoryMessageTraceNumber = Nothing
              }
      decode (encode msg) `shouldBe` Just msg

  describe "CreateMemoryMessageResponse" $
    it "decodes the doc sample" $
      decode sampleCreateMessageResponse
        `shouldBe` Just (CreateMemoryMessageResponse {createMemoryMessageResponseId = Just "sxa2cY0BfUsSoeNTz-8m"})

  describe "MemoryMessageTraces" $
    it "decodes an empty traces envelope" $
      decode sampleTraces `shouldBe` Just (MemoryMessageTraces {memoryMessageTracesTraces = []})

  describe "CreateMemoryMessageRequest" $
    it "encodes only the supplied fields" $ do
      let req =
            CreateMemoryMessageRequest
              { createMemoryMessageInput = Just "How do I make an interaction?",
                createMemoryMessagePromptTemplate = Nothing,
                createMemoryMessageResponse = Nothing,
                createMemoryMessageOrigin = Nothing,
                createMemoryMessageAdditionalInfo = Nothing
              }
      encode req `shouldBe` "{\"input\":\"How do I make an interaction?\"}"