EventStream.hs

module Kioku.Memory.EventStream
  ( MemoryEventStream,
    memoryEventStream,
    memoryCodec,
    memoryStream,
    parseMemoryEvent,
  )
where

import Data.Aeson (Value)
import Data.Aeson.Types (Parser, parseEither, withObject, (.:), (.:?))
import Data.Text qualified as Text
import Keiki.Core (HsPred)
import Keiki.Generics (emptyRegFile)
import Keiro.Codec (Codec (..), EventType (..))
import Keiro.EventStream (EventStream (..), SnapshotPolicy (..))
import Keiro.EventStream.Validate (ValidatedEventStream, mkEventStreamOrThrow)
import Keiro.Stream (Stream)
import Keiro.Stream qualified as Stream
import Kioku.Api.Access (RecordedPrincipal (..), legacyMemorySpaceId, legacyPrincipalRef)
import Kioku.Api.Scope (MemoryScope (..), Namespace (..), ScopeKind (..))
import Kioku.Id (MemoryId, SessionId, idText, parseIdLenient)
import Kioku.Memory.Domain
import Kioku.Prelude

type MemoryEventStream =
  EventStream (HsPred MemoryRegs MemoryCommand) MemoryRegs MemoryVertex MemoryCommand MemoryEvent

memoryStream :: MemoryId -> Stream MemoryEventStream
memoryStream mid = Stream.entityStream (Stream.categoryUnsafe "kioku_memory") (idText mid)

memoryEventStream :: ValidatedEventStream (HsPred MemoryRegs MemoryCommand) MemoryRegs MemoryVertex MemoryCommand MemoryEvent
memoryEventStream =
  mkEventStreamOrThrow "kioku-memory" memoryEventStreamDefinition

memoryEventStreamDefinition :: MemoryEventStream
memoryEventStreamDefinition =
  EventStream
    { transducer = memoryTransducer,
      initialState = NotCreated,
      initialRegisters = emptyRegFile,
      eventCodec = memoryCodec,
      resolveStreamName = Stream.streamName,
      snapshotPolicy = Never,
      stateCodec = Nothing
    }

memoryCodec :: Codec MemoryEvent
memoryCodec =
  Codec
    { eventTypes =
        EventType
          <$> "MemoryRecorded"
            :| [ "MemorySuperseded",
                 "MemoryArchived",
                 "MemoryTagsUpdated",
                 "MemoryConfidenceUpdated",
                 "MemoryMerged"
               ],
      eventType =
        EventType . \case
          MemoryRecorded {} -> "MemoryRecorded"
          MemorySuperseded {} -> "MemorySuperseded"
          MemoryArchived {} -> "MemoryArchived"
          MemoryTagsUpdated {} -> "MemoryTagsUpdated"
          MemoryConfidenceUpdated {} -> "MemoryConfidenceUpdated"
          MemoryMerged {} -> "MemoryMerged",
      schemaVersion = 1,
      encode = toJSON,
      decode = const parseMemoryEvent,
      upcasters = []
    }

parseMemoryEvent :: Value -> Either Text MemoryEvent
parseMemoryEvent value =
  case parseEither parseJSON value of
    Right event -> Right event
    Left nativeErr ->
      case parseEither parseLegacyMemoryEvent value of
        Right event -> Right event
        Left legacyErr -> Left (Text.pack nativeErr <> "; legacy decode failed: " <> Text.pack legacyErr)

parseLegacyMemoryEvent :: Value -> Parser MemoryEvent
parseLegacyMemoryEvent =
  withObject "Rei AgentMemoryEvent" $ \o -> do
    tag <- o .: "type"
    payload <- o .: "data"
    case tag of
      "agent_memory_recorded" -> MemoryRecorded <$> parseLegacyMemoryRecorded payload
      "agent_memory_superseded" -> MemorySuperseded <$> parseLegacyMemorySuperseded payload
      "agent_memory_archived" -> MemoryArchived <$> parseLegacyMemoryArchived payload
      "agent_memory_tags_updated" -> MemoryTagsUpdated <$> parseLegacyMemoryTagsUpdated payload
      "agent_memory_confidence_updated" -> MemoryConfidenceUpdated <$> parseLegacyMemoryConfidenceUpdated payload
      other -> fail ("Unknown Rei AgentMemoryEvent tag: " <> Text.unpack other)

-- | Rei payloads predate memory spaces by even more than Kioku's own older events do, so they
-- take the same two rules: the legacy space, and the legacy agent label kept verbatim. Rei's
-- @agentId@ is exactly the free-text label 'LegacyPrincipal' exists to mark.
parseLegacyMemoryRecorded :: Value -> Parser MemoryRecordedData
parseLegacyMemoryRecorded =
  withObject "Rei AgentMemoryRecordedData" $ \o -> do
    memoryId <- parseLegacyMemoryId =<< o .: "memoryId"
    agentId <- o .: "agentId"
    sessionId <- traverse parseLegacySessionId =<< o .:? "sessionId"
    scope <- parseLegacyAnchor =<< o .: "anchor"
    supersedes <- traverse parseLegacyMemoryId =<< o .:? "supersedes"
    MemoryRecordedData
      memoryId
      legacyMemorySpaceId
      (LegacyPrincipal (legacyPrincipalRef agentId))
      Nothing
      agentId
      sessionId
      scope
      <$> o .: "memoryType"
      <*> o .: "content"
      <*> pure 100
      <*> o .: "confidence"
      <*> o .: "tags"
      <*> pure supersedes
      <*> o .: "recordedAt"

parseLegacyMemorySuperseded :: Value -> Parser MemorySupersededData
parseLegacyMemorySuperseded =
  withObject "Rei AgentMemorySupersededData" $ \o ->
    MemorySupersededData
      <$> (parseLegacyMemoryId =<< o .: "memoryId")
      <*> pure legacyMemorySpaceId
      <*> pure UnattributedPrincipal
      <*> (parseLegacyMemoryId =<< o .: "supersededBy")
      <*> o .: "supersededAt"

parseLegacyMemoryArchived :: Value -> Parser MemoryArchivedData
parseLegacyMemoryArchived =
  withObject "Rei AgentMemoryArchivedData" $ \o ->
    MemoryArchivedData
      <$> (parseLegacyMemoryId =<< o .: "memoryId")
      <*> pure legacyMemorySpaceId
      <*> pure UnattributedPrincipal
      <*> o .: "archivedAt"

parseLegacyMemoryTagsUpdated :: Value -> Parser MemoryTagsUpdatedData
parseLegacyMemoryTagsUpdated =
  withObject "Rei AgentMemoryTagsUpdatedData" $ \o ->
    MemoryTagsUpdatedData
      <$> (parseLegacyMemoryId =<< o .: "memoryId")
      <*> pure legacyMemorySpaceId
      <*> pure UnattributedPrincipal
      <*> o .: "tags"
      <*> o .: "updatedAt"

parseLegacyMemoryConfidenceUpdated :: Value -> Parser MemoryConfidenceUpdatedData
parseLegacyMemoryConfidenceUpdated =
  withObject "Rei AgentMemoryConfidenceUpdatedData" $ \o ->
    MemoryConfidenceUpdatedData
      <$> (parseLegacyMemoryId =<< o .: "memoryId")
      <*> pure legacyMemorySpaceId
      <*> pure UnattributedPrincipal
      <*> o .: "confidence"
      <*> o .: "updatedAt"

parseLegacyAnchor :: Value -> Parser MemoryScope
parseLegacyAnchor =
  withObject "Rei MemoryAnchor" $ \o -> do
    anchorType <- o .: "type"
    case anchorType of
      "intention" -> ScopeEntity reiNamespace (ScopeKind "intention") <$> o .: "id"
      "habit" -> ScopeEntity reiNamespace (ScopeKind "habit") <$> o .: "id"
      "workspace" -> pure (ScopeGlobal reiNamespace)
      other -> fail ("Unknown Rei MemoryAnchor type: " <> Text.unpack other)

parseLegacyMemoryId :: Text -> Parser MemoryId
parseLegacyMemoryId = either (fail . Text.unpack) pure . parseIdLenient

parseLegacySessionId :: Text -> Parser SessionId
parseLegacySessionId = either (fail . Text.unpack) pure . parseIdLenient

reiNamespace :: Namespace
reiNamespace = Namespace "rei"