{-# LANGUAGE DeriveGeneric #-} {-# LANGUAGE LambdaCase #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE NoImplicitPrelude #-} -- | Cross-agent shared memory system with vector similarity search. -- -- Provides persistent memory that is: -- - Shared across all agents (Telegram, researcher, coder, etc.) -- - Private per user (users can't see each other's memories) -- - Searchable via semantic similarity using embeddings -- -- Uses sqlite-vss for vector similarity search and Ollama for embeddings. -- -- : out omni-agent-memory -- : dep aeson -- : dep http-conduit -- : dep sqlite-simple -- : dep uuid -- : dep vector -- : dep directory -- : dep bytestring module Omni.Agent.Memory ( -- * Types User (..), Memory (..), MemorySource (..), ConversationMessage (..), ConversationSummary (..), MessageRole (..), RelationType (..), MemoryLink (..), -- * User Management createUser, getUser, getUserByTelegramId, getOrCreateUserByTelegramId, -- * Memory Operations storeMemory, recallMemories, forgetMemory, getAllMemoriesForUser, updateMemoryAccess, -- * Knowledge Graph linkMemories, getMemoryLinks, getLinkedMemories, queryGraph, -- * Conversation History (DMs) saveMessage, getRecentMessages, getConversationContext, summarizeAndArchive, estimateTokens, -- * Group Conversation History saveGroupMessage, getGroupRecentMessages, getGroupConversationContext, -- * Group Memories storeGroupMemory, recallGroupMemories, -- * Embeddings embedText, -- * Agent Integration rememberTool, recallTool, linkMemoriesTool, queryGraphTool, formatMemoriesForPrompt, runAgentWithMemory, -- * Database withMemoryDb, initMemoryDb, getMemoryDbPath, -- * Testing main, test, ) where import Alpha import Data.Aeson ((.!=), (.:), (.:?), (.=)) import qualified Data.Aeson as Aeson import qualified Data.Aeson.KeyMap as KeyMap import qualified Data.ByteString as BS import qualified Data.ByteString.Lazy as BL import qualified Data.List as List import qualified Data.Text as Text import qualified Data.Text.Encoding as TE import Data.Time (UTCTime, getCurrentTime) import Data.Time.Format (defaultTimeLocale, formatTime) import qualified Data.UUID as UUID import qualified Data.UUID.V4 as UUID import qualified Data.Vector.Storable as VS import qualified Database.SQLite.Simple as SQL import Database.SQLite.Simple.FromField () import qualified Database.SQLite.Simple.ToField as SQL import Foreign.Storable () import qualified Network.HTTP.Simple as HTTP import qualified Omni.Agent.Engine as Engine import qualified Omni.Test as Test import System.Directory (createDirectoryIfMissing) import System.Environment (lookupEnv) import System.FilePath (takeDirectory, ()) main :: IO () main = Test.run test test :: Test.Tree test = Test.group "Omni.Agent.Memory" [ Test.unit "User JSON roundtrip" <| do now <- getCurrentTime let user = User { userId = "test-uuid", userTelegramId = Just 12345, userEmail = Nothing, userName = "Test User", userCreatedAt = now } case Aeson.decode (Aeson.encode user) of Nothing -> Test.assertFailure "Failed to decode User" Just decoded -> userName decoded Test.@=? "Test User", Test.unit "Memory JSON roundtrip" <| do now <- getCurrentTime let mem = Memory { memoryId = "mem-uuid", memoryUserId = "user-uuid", memoryContent = "User is an AI engineer", memoryEmbedding = Nothing, memorySource = MemorySource { sourceAgent = "telegram", sourceSession = Nothing, sourceContext = "User mentioned in chat" }, memoryConfidence = 0.9, memoryCreatedAt = now, memoryLastAccessedAt = now, memoryTags = ["profession", "ai"] } case Aeson.decode (Aeson.encode mem) of Nothing -> Test.assertFailure "Failed to decode Memory" Just decoded -> memoryContent decoded Test.@=? "User is an AI engineer", Test.unit "MemorySource JSON roundtrip" <| do let src = MemorySource { sourceAgent = "researcher", sourceSession = Just "session-123", sourceContext = "Extracted from conversation" } case Aeson.decode (Aeson.encode src) of Nothing -> Test.assertFailure "Failed to decode MemorySource" Just decoded -> sourceAgent decoded Test.@=? "researcher", Test.unit "formatMemoriesForPrompt formats correctly" <| do now <- getCurrentTime let mem1 = Memory { memoryId = "1", memoryUserId = "u", memoryContent = "User is an AI engineer", memoryEmbedding = Nothing, memorySource = MemorySource "telegram" Nothing "chat", memoryConfidence = 0.9, memoryCreatedAt = now, memoryLastAccessedAt = now, memoryTags = [] } mem2 = Memory { memoryId = "2", memoryUserId = "u", memoryContent = "User prefers Haskell", memoryEmbedding = Nothing, memorySource = MemorySource "coder" Nothing "code review", memoryConfidence = 0.8, memoryCreatedAt = now, memoryLastAccessedAt = now, memoryTags = [] } formatted = formatMemoriesForPrompt [mem1, mem2] ("AI engineer" `Text.isInfixOf` formatted) Test.@=? True ("Haskell" `Text.isInfixOf` formatted) Test.@=? True, Test.unit "cosineSimilarity identical vectors" <| do let v1 = VS.fromList [1.0, 0.0, 0.0 :: Float] v2 = VS.fromList [1.0, 0.0, 0.0 :: Float] abs (cosineSimilarity v1 v2 - 1.0) < 0.0001 Test.@=? True, Test.unit "cosineSimilarity orthogonal vectors" <| do let v1 = VS.fromList [1.0, 0.0, 0.0 :: Float] v2 = VS.fromList [0.0, 1.0, 0.0 :: Float] abs (cosineSimilarity v1 v2) < 0.0001 Test.@=? True, Test.unit "cosineSimilarity opposite vectors" <| do let v1 = VS.fromList [1.0, 0.0, 0.0 :: Float] v2 = VS.fromList [-1.0, 0.0, 0.0 :: Float] abs (cosineSimilarity v1 v2 + 1.0) < 0.0001 Test.@=? True, Test.unit "vectorToBlob and blobToVector roundtrip" <| do let v = VS.fromList [0.1, 0.2, 0.3, 0.4, 0.5 :: Float] blob = vectorToBlob v v' = blobToVector blob VS.length v Test.@=? VS.length v' VS.toList v Test.@=? VS.toList v', Test.unit "rememberTool has correct schema" <| do let tool = rememberTool "test-user-id" Engine.toolName tool Test.@=? "remember", Test.unit "recallTool has correct schema" <| do let tool = recallTool "test-user-id" Engine.toolName tool Test.@=? "recall", Test.unit "RelationType JSON roundtrip" <| do let types = [Contradicts, Supports, Elaborates, Supersedes, Related, ContingentOn] forM_ types <| \rt -> case Aeson.decode (Aeson.encode rt) of Nothing -> Test.assertFailure ("Failed to decode RelationType: " <> show rt) Just decoded -> decoded Test.@=? rt, Test.unit "MemoryLink JSON roundtrip" <| do now <- getCurrentTime let memLink = MemoryLink { linkFromMemoryId = "mem-1", linkToMemoryId = "mem-2", linkRelationType = Contradicts, linkCreatedAt = now } case Aeson.decode (Aeson.encode memLink) of Nothing -> Test.assertFailure "Failed to decode MemoryLink" Just decoded -> do linkFromMemoryId decoded Test.@=? "mem-1" linkToMemoryId decoded Test.@=? "mem-2" linkRelationType decoded Test.@=? Contradicts, Test.unit "relationTypeToText and textToRelationType roundtrip" <| do let types = [Contradicts, Supports, Elaborates, Supersedes, Related, ContingentOn] forM_ types <| \rt -> textToRelationType (relationTypeToText rt) Test.@=? Just rt, Test.unit "linkMemoriesTool has correct schema" <| do let tool = linkMemoriesTool "test-user-id" Engine.toolName tool Test.@=? "link_memories", Test.unit "queryGraphTool has correct schema" <| do let tool = queryGraphTool "test-user-id" Engine.toolName tool Test.@=? "query_graph" ] -- | User record for multi-user memory system. data User = User { userId :: Text, userTelegramId :: Maybe Int, userEmail :: Maybe Text, userName :: Text, userCreatedAt :: UTCTime } deriving (Show, Eq, Generic) instance Aeson.ToJSON User where toJSON u = Aeson.object [ "id" .= userId u, "telegram_id" .= userTelegramId u, "email" .= userEmail u, "name" .= userName u, "created_at" .= userCreatedAt u ] instance Aeson.FromJSON User where parseJSON = Aeson.withObject "User" <| \v -> (User (v .:? "telegram_id") <*> (v .:? "email") <*> (v .: "name") <*> (v .: "created_at") instance SQL.FromRow User where fromRow = User SQL.field <*> SQL.field <*> SQL.field <*> SQL.field instance SQL.ToRow User where toRow u = [ SQL.toField (userId u), SQL.toField (userTelegramId u), SQL.toField (userEmail u), SQL.toField (userName u), SQL.toField (userCreatedAt u) ] -- | Source information for a memory. data MemorySource = MemorySource { sourceAgent :: Text, sourceSession :: Maybe Text, sourceContext :: Text } deriving (Show, Eq, Generic) instance Aeson.ToJSON MemorySource where toJSON s = Aeson.object [ "agent" .= sourceAgent s, "session" .= sourceSession s, "context" .= sourceContext s ] instance Aeson.FromJSON MemorySource where parseJSON = Aeson.withObject "MemorySource" <| \v -> (MemorySource (v .:? "session") <*> (v .: "context") -- | A memory stored in the system. data Memory = Memory { memoryId :: Text, memoryUserId :: Text, memoryContent :: Text, memoryEmbedding :: Maybe (VS.Vector Float), memorySource :: MemorySource, memoryConfidence :: Double, memoryCreatedAt :: UTCTime, memoryLastAccessedAt :: UTCTime, memoryTags :: [Text] } deriving (Show, Eq, Generic) instance Aeson.ToJSON Memory where toJSON m = Aeson.object [ "id" .= memoryId m, "user_id" .= memoryUserId m, "content" .= memoryContent m, "source" .= memorySource m, "confidence" .= memoryConfidence m, "created_at" .= memoryCreatedAt m, "last_accessed_at" .= memoryLastAccessedAt m, "tags" .= memoryTags m ] instance Aeson.FromJSON Memory where parseJSON = Aeson.withObject "Memory" <| \v -> ( Memory (v .: "user_id") <*> (v .: "content") <*> pure Nothing <*> (v .: "source") <*> (v .:? "confidence" .!= 0.8) <*> (v .: "created_at") <*> (v .: "last_accessed_at") <*> (v .:? "tags" .!= []) -- SQLite instances for Memory (partial - embedding handled separately) instance SQL.FromRow Memory where fromRow = do mid <- SQL.field uid <- SQL.field content <- SQL.field embeddingBlob <- SQL.field agent <- SQL.field session <- SQL.field context <- SQL.field confidence <- SQL.field createdAt <- SQL.field lastAccessedAt <- SQL.field tagsJson <- SQL.field let embedding = blobToVector (Aeson.decode <. BL.fromStrict <. TE.encodeUtf8)) pure Memory { memoryId = mid, memoryUserId = uid, memoryContent = content, memoryEmbedding = embedding, memorySource = source, memoryConfidence = confidence, memoryCreatedAt = createdAt, memoryLastAccessedAt = lastAccessedAt, memoryTags = tags } -- | Role in a conversation message. data MessageRole = UserRole | AssistantRole deriving (Show, Eq, Generic) instance Aeson.ToJSON MessageRole where toJSON UserRole = Aeson.String "user" toJSON AssistantRole = Aeson.String "assistant" instance Aeson.FromJSON MessageRole where parseJSON = Aeson.withText "MessageRole" <| \case "user" -> pure UserRole "assistant" -> pure AssistantRole _ -> empty -- | A message in a conversation. data ConversationMessage = ConversationMessage { cmId :: Maybe Int, cmUserId :: Text, cmChatId :: Int, cmRole :: MessageRole, cmSenderName :: Maybe Text, cmContent :: Text, cmTokensEstimate :: Int, cmCreatedAt :: UTCTime } deriving (Show, Eq, Generic) instance Aeson.ToJSON ConversationMessage where toJSON m = Aeson.object [ "id" .= cmId m, "user_id" .= cmUserId m, "chat_id" .= cmChatId m, "role" .= cmRole m, "sender_name" .= cmSenderName m, "content" .= cmContent m, "tokens_estimate" .= cmTokensEstimate m, "created_at" .= cmCreatedAt m ] instance SQL.FromRow ConversationMessage where fromRow = (ConversationMessage SQL.field <*> SQL.field <*> (parseRole SQL.field <*> SQL.field <*> (fromMaybe 0 SQL.field where parseRole :: Text -> MessageRole parseRole "user" = UserRole parseRole _ = AssistantRole -- | A summary of older conversation messages. data ConversationSummary = ConversationSummary { csId :: Maybe Int, csUserId :: Text, csChatId :: Int, csSummary :: Text, csMessagesSummarized :: Int, csTokensSaved :: Maybe Int, csCreatedAt :: UTCTime } deriving (Show, Eq, Generic) instance Aeson.ToJSON ConversationSummary where toJSON s = Aeson.object [ "id" .= csId s, "user_id" .= csUserId s, "chat_id" .= csChatId s, "summary" .= csSummary s, "messages_summarized" .= csMessagesSummarized s, "tokens_saved" .= csTokensSaved s, "created_at" .= csCreatedAt s ] instance SQL.FromRow ConversationSummary where fromRow = (ConversationSummary SQL.field <*> SQL.field <*> SQL.field <*> SQL.field <*> SQL.field <*> SQL.field -- | Relation types for the knowledge graph. data RelationType = Contradicts | Supports | Elaborates | Supersedes | Related | ContingentOn deriving (Show, Eq, Generic, Ord) instance Aeson.ToJSON RelationType where toJSON Contradicts = Aeson.String "contradicts" toJSON Supports = Aeson.String "supports" toJSON Elaborates = Aeson.String "elaborates" toJSON Supersedes = Aeson.String "supersedes" toJSON Related = Aeson.String "related" toJSON ContingentOn = Aeson.String "contingent_on" instance Aeson.FromJSON RelationType where parseJSON = Aeson.withText "RelationType" <| \case "contradicts" -> pure Contradicts "supports" -> pure Supports "elaborates" -> pure Elaborates "supersedes" -> pure Supersedes "related" -> pure Related "contingent_on" -> pure ContingentOn _ -> empty relationTypeToText :: RelationType -> Text relationTypeToText Contradicts = "contradicts" relationTypeToText Supports = "supports" relationTypeToText Elaborates = "elaborates" relationTypeToText Supersedes = "supersedes" relationTypeToText Related = "related" relationTypeToText ContingentOn = "contingent_on" textToRelationType :: Text -> Maybe RelationType textToRelationType "contradicts" = Just Contradicts textToRelationType "supports" = Just Supports textToRelationType "elaborates" = Just Elaborates textToRelationType "supersedes" = Just Supersedes textToRelationType "related" = Just Related textToRelationType "contingent_on" = Just ContingentOn textToRelationType _ = Nothing -- | A link between two memories in the knowledge graph. data MemoryLink = MemoryLink { linkFromMemoryId :: Text, linkToMemoryId :: Text, linkRelationType :: RelationType, linkCreatedAt :: UTCTime } deriving (Show, Eq, Generic) instance Aeson.ToJSON MemoryLink where toJSON l = Aeson.object [ "from_memory_id" .= linkFromMemoryId l, "to_memory_id" .= linkToMemoryId l, "relation_type" .= linkRelationType l, "created_at" .= linkCreatedAt l ] instance Aeson.FromJSON MemoryLink where parseJSON = Aeson.withObject "MemoryLink" <| \v -> (MemoryLink (v .: "to_memory_id") <*> (v .: "relation_type") <*> (v .: "created_at") instance SQL.FromRow MemoryLink where fromRow = do fromId <- SQL.field toId <- SQL.field relTypeText <- SQL.field createdAt <- SQL.field let relType = fromMaybe Related (textToRelationType relTypeText) pure MemoryLink { linkFromMemoryId = fromId, linkToMemoryId = toId, linkRelationType = relType, linkCreatedAt = createdAt } -- | Get the path to memory.db getMemoryDbPath :: IO FilePath getMemoryDbPath = do maybeEnv <- lookupEnv "MEMORY_DB_PATH" case maybeEnv of Just p -> pure p Nothing -> do home <- lookupEnv "HOME" case home of Just h -> pure (h ".local/share/omni/memory.db") Nothing -> pure "_/memory.db" -- | Run an action with the memory database connection. withMemoryDb :: (SQL.Connection -> IO a) -> IO a withMemoryDb action = do dbPath <- getMemoryDbPath createDirectoryIfMissing True (takeDirectory dbPath) SQL.withConnection dbPath <| \conn -> do initMemoryDb conn action conn -- | Initialize the memory database schema. initMemoryDb :: SQL.Connection -> IO () initMemoryDb conn = do SQL.execute_ conn "PRAGMA busy_timeout = 10000" SQL.execute_ conn "PRAGMA foreign_keys = ON" _ <- SQL.query_ conn "PRAGMA journal_mode = WAL" :: IO [[Text]] SQL.execute_ conn "CREATE TABLE IF NOT EXISTS users (\ \ id TEXT PRIMARY KEY,\ \ telegram_id INTEGER UNIQUE,\ \ email TEXT UNIQUE,\ \ name TEXT NOT NULL,\ \ created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP\ \)" SQL.execute_ conn "CREATE TABLE IF NOT EXISTS memories (\ \ id TEXT PRIMARY KEY,\ \ user_id TEXT NOT NULL REFERENCES users(id),\ \ content TEXT NOT NULL,\ \ embedding BLOB,\ \ source_agent TEXT NOT NULL,\ \ source_session TEXT,\ \ source_context TEXT,\ \ confidence REAL DEFAULT 0.8,\ \ created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,\ \ last_accessed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,\ \ tags TEXT\ \)" SQL.execute_ conn "CREATE INDEX IF NOT EXISTS idx_memories_user ON memories(user_id)" SQL.execute_ conn "CREATE INDEX IF NOT EXISTS idx_memories_agent ON memories(source_agent)" SQL.execute_ conn "CREATE TABLE IF NOT EXISTS conversation_messages (\ \ id INTEGER PRIMARY KEY AUTOINCREMENT,\ \ user_id TEXT NOT NULL REFERENCES users(id),\ \ chat_id INTEGER NOT NULL,\ \ role TEXT NOT NULL,\ \ sender_name TEXT,\ \ content TEXT NOT NULL,\ \ tokens_estimate INTEGER,\ \ created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP\ \)" SQL.execute_ conn "CREATE INDEX IF NOT EXISTS idx_conv_user_chat ON conversation_messages(user_id, chat_id)" migrateConversationMessages conn SQL.execute_ conn "CREATE TABLE IF NOT EXISTS conversation_summaries (\ \ id INTEGER PRIMARY KEY AUTOINCREMENT,\ \ user_id TEXT NOT NULL REFERENCES users(id),\ \ chat_id INTEGER NOT NULL,\ \ summary TEXT NOT NULL,\ \ messages_summarized INTEGER NOT NULL,\ \ tokens_saved INTEGER,\ \ created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP\ \)" SQL.execute_ conn "CREATE INDEX IF NOT EXISTS idx_summary_user_chat ON conversation_summaries(user_id, chat_id)" SQL.execute_ conn "CREATE TABLE IF NOT EXISTS notes (\ \ id INTEGER PRIMARY KEY AUTOINCREMENT,\ \ user_id TEXT NOT NULL,\ \ topic TEXT NOT NULL,\ \ content TEXT NOT NULL,\ \ created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP\ \)" SQL.execute_ conn "CREATE INDEX IF NOT EXISTS idx_notes_user ON notes(user_id)" SQL.execute_ conn "CREATE INDEX IF NOT EXISTS idx_notes_topic ON notes(user_id, topic)" SQL.execute_ conn "CREATE TABLE IF NOT EXISTS todos (\ \ id INTEGER PRIMARY KEY AUTOINCREMENT,\ \ user_id TEXT NOT NULL,\ \ title TEXT NOT NULL,\ \ due_date TIMESTAMP,\ \ completed INTEGER NOT NULL DEFAULT 0,\ \ created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP\ \)" SQL.execute_ conn "CREATE INDEX IF NOT EXISTS idx_todos_user ON todos(user_id)" SQL.execute_ conn "CREATE INDEX IF NOT EXISTS idx_todos_due ON todos(user_id, due_date)" SQL.execute_ conn "CREATE TABLE IF NOT EXISTS memory_links (\ \ from_memory_id TEXT NOT NULL REFERENCES memories(id) ON DELETE CASCADE,\ \ to_memory_id TEXT NOT NULL REFERENCES memories(id) ON DELETE CASCADE,\ \ relation_type TEXT NOT NULL,\ \ created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,\ \ PRIMARY KEY (from_memory_id, to_memory_id, relation_type)\ \)" SQL.execute_ conn "CREATE INDEX IF NOT EXISTS idx_memory_links_from ON memory_links(from_memory_id)" SQL.execute_ conn "CREATE INDEX IF NOT EXISTS idx_memory_links_to ON memory_links(to_memory_id)" SQL.execute_ conn "CREATE INDEX IF NOT EXISTS idx_memory_links_type ON memory_links(relation_type)" -- | Migrate conversation_messages to add sender_name and thread_id columns. migrateConversationMessages :: SQL.Connection -> IO () migrateConversationMessages conn = do columns <- SQL.query_ conn "PRAGMA table_info(conversation_messages)" :: IO [(Int, Text, Text, Int, Maybe Text, Int)] let columnNames = map (\(_, name, _, _, _, _) -> name) columns unless ("sender_name" `elem` columnNames) <| do SQL.execute_ conn "ALTER TABLE conversation_messages ADD COLUMN sender_name TEXT" SQL.execute_ conn "UPDATE conversation_messages SET sender_name = 'bensima' WHERE role = 'user' AND sender_name IS NULL" unless ("thread_id" `elem` columnNames) <| do SQL.execute_ conn "ALTER TABLE conversation_messages ADD COLUMN thread_id INTEGER" SQL.execute_ conn "CREATE INDEX IF NOT EXISTS idx_conv_chat_thread ON conversation_messages(chat_id, thread_id)" -- | Create a new user. createUser :: Text -> Maybe Int -> IO User createUser name telegramId = do uuid <- UUID.nextRandom now <- getCurrentTime let user = User { userId = UUID.toText uuid, userTelegramId = telegramId, userEmail = Nothing, userName = name, userCreatedAt = now } withMemoryDb <| \conn -> SQL.execute conn "INSERT INTO users (id, telegram_id, email, name, created_at) VALUES (?, ?, ?, ?, ?)" user pure user -- | Get a user by ID. getUser :: Text -> IO (Maybe User) getUser uid = withMemoryDb <| \conn -> do results <- SQL.query conn "SELECT id, telegram_id, email, name, created_at FROM users WHERE id = ?" (SQL.Only uid) pure (listToMaybe results) -- | Get a user by Telegram ID. getUserByTelegramId :: Int -> IO (Maybe User) getUserByTelegramId tid = withMemoryDb <| \conn -> do results <- SQL.query conn "SELECT id, telegram_id, email, name, created_at FROM users WHERE telegram_id = ?" (SQL.Only tid) pure (listToMaybe results) -- | Get or create a user by Telegram ID. getOrCreateUserByTelegramId :: Int -> Text -> IO User getOrCreateUserByTelegramId tid name = do existing <- getUserByTelegramId tid case existing of Just user -> pure user Nothing -> createUser name (Just tid) -- | Store a memory for a user. storeMemory :: Text -> Text -> MemorySource -> IO Memory storeMemory uid content source = storeMemoryWithTags uid content source [] -- | Store a memory with tags. storeMemoryWithTags :: Text -> Text -> MemorySource -> [Text] -> IO Memory storeMemoryWithTags uid content source tags = do uuid <- UUID.nextRandom now <- getCurrentTime embedding <- embedText content let mem = Memory { memoryId = UUID.toText uuid, memoryUserId = uid, memoryContent = content, memoryEmbedding = either (const Nothing) Just embedding, memorySource = source, memoryConfidence = 0.8, memoryCreatedAt = now, memoryLastAccessedAt = now, memoryTags = tags } withMemoryDb <| \conn -> SQL.execute conn "INSERT INTO memories (id, user_id, content, embedding, source_agent, source_session, source_context, confidence, created_at, last_accessed_at, tags) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)" ( ( memoryId mem, memoryUserId mem, memoryContent mem, vectorToBlob Text -> Int -> IO [Memory] recallMemories uid query limit = do queryEmbedding <- embedText query case queryEmbedding of Left _ -> recallMemoriesByRecency uid limit Right qEmb -> do allMems <- getAllMemoriesForUser uid let scored = [ (m, cosineSimilarity qEmb emb) | m <- allMems, Just emb <- [memoryEmbedding m] ] sorted = List.sortBy (\(_, s1) (_, s2) -> compare s2 s1) scored topN = take limit sorted now <- getCurrentTime traverse_ (updateMemoryAccess now <. memoryId <. fst) topN pure (map fst topN) -- | Recall memories by recency (fallback when embedding fails). recallMemoriesByRecency :: Text -> Int -> IO [Memory] recallMemoriesByRecency uid limit = withMemoryDb <| \conn -> do SQL.query conn "SELECT id, user_id, content, embedding, source_agent, source_session, source_context, confidence, created_at, last_accessed_at, tags \ \FROM memories WHERE user_id = ? ORDER BY last_accessed_at DESC LIMIT ?" (uid, limit) -- | Get all memories for a user. getAllMemoriesForUser :: Text -> IO [Memory] getAllMemoriesForUser uid = withMemoryDb <| \conn -> SQL.query conn "SELECT id, user_id, content, embedding, source_agent, source_session, source_context, confidence, created_at, last_accessed_at, tags \ \FROM memories WHERE user_id = ?" (SQL.Only uid) -- | Delete a memory. forgetMemory :: Text -> IO () forgetMemory mid = withMemoryDb <| \conn -> SQL.execute conn "DELETE FROM memories WHERE id = ?" (SQL.Only mid) -- | Update memory's last accessed timestamp. updateMemoryAccess :: UTCTime -> Text -> IO () updateMemoryAccess now mid = withMemoryDb <| \conn -> SQL.execute conn "UPDATE memories SET last_accessed_at = ? WHERE id = ?" (now, mid) -- | Create a link between two memories. linkMemories :: Text -> Text -> RelationType -> IO MemoryLink linkMemories fromId toId relType = do now <- getCurrentTime withMemoryDb <| \conn -> SQL.execute conn "INSERT OR REPLACE INTO memory_links (from_memory_id, to_memory_id, relation_type, created_at) VALUES (?, ?, ?, ?)" (fromId, toId, relationTypeToText relType, now) pure MemoryLink { linkFromMemoryId = fromId, linkToMemoryId = toId, linkRelationType = relType, linkCreatedAt = now } -- | Get all links from a memory. getMemoryLinks :: Text -> IO [MemoryLink] getMemoryLinks memId = withMemoryDb <| \conn -> SQL.query conn "SELECT from_memory_id, to_memory_id, relation_type, created_at \ \FROM memory_links WHERE from_memory_id = ? OR to_memory_id = ?" (memId, memId) -- | Get memories linked to a given memory with their content. getLinkedMemories :: Text -> Maybe RelationType -> IO [(MemoryLink, Memory)] getLinkedMemories memId maybeRelType = do links <- getMemoryLinks memId let filteredLinks = case maybeRelType of Nothing -> links Just rt -> filter (\l -> linkRelationType l == rt) links mems <- traverse loadMemory filteredLinks pure [(l, m) | (l, Just m) <- zip filteredLinks mems] where loadMemory memLink = do let targetId = if linkFromMemoryId memLink == memId then linkToMemoryId memLink else linkFromMemoryId memLink withMemoryDb <| \conn -> do results <- SQL.query conn "SELECT id, user_id, content, embedding, source_agent, source_session, source_context, confidence, created_at, last_accessed_at, tags \ \FROM memories WHERE id = ?" (SQL.Only targetId) pure (listToMaybe results) -- | Query the knowledge graph by traversing links from a starting memory. -- Returns all memories reachable within the given depth. queryGraph :: Text -> Int -> Maybe RelationType -> IO [(Memory, [MemoryLink])] queryGraph startMemId maxDepth maybeRelType = do startMem <- getMemoryById startMemId case startMem of Nothing -> pure [] Just mem -> go [startMemId] [(mem, [])] 0 where go :: [Text] -> [(Memory, [MemoryLink])] -> Int -> IO [(Memory, [MemoryLink])] go _ acc depth | depth >= maxDepth = pure acc go visitedIds acc depth = do let currentIds = map (memoryId <. fst) acc newIds = filter (`notElem` visitedIds) currentIds if null newIds then pure acc else do newLinked <- concat map (memoryId <. fst) newMems go newVisited (acc <> newMems) (depth + 1) -- | Get a memory by ID. getMemoryById :: Text -> IO (Maybe Memory) getMemoryById memId = withMemoryDb <| \conn -> do results <- SQL.query conn "SELECT id, user_id, content, embedding, source_agent, source_session, source_context, confidence, created_at, last_accessed_at, tags \ \FROM memories WHERE id = ?" (SQL.Only memId) pure (listToMaybe results) -- | Embed text using Ollama's nomic-embed-text model. embedText :: Text -> IO (Either Text (VS.Vector Float)) embedText content = do ollamaUrl <- fromMaybe "http://localhost:11434" "/api/embeddings" req0 <- HTTP.parseRequest url let body = Aeson.object [ "model" .= ("nomic-embed-text" :: Text), "prompt" .= content ] req = HTTP.setRequestMethod "POST" <| HTTP.setRequestHeader "Content-Type" ["application/json"] <| HTTP.setRequestBodyLBS (Aeson.encode body) <| req0 result <- try (HTTP.httpLBS req) case result of Left (e :: SomeException) -> pure (Left ("Embedding request failed: " <> tshow e)) Right response -> do let status = HTTP.getResponseStatusCode response if status >= 200 && status < 300 then case Aeson.decode (HTTP.getResponseBody response) of Just (Aeson.Object obj) -> case KeyMap.lookup "embedding" obj of Just (Aeson.Array arr) -> let floats = [f | Aeson.Number n <- toList arr, let f = realToFrac n] in pure (Right (VS.fromList floats)) _ -> pure (Left "No embedding in response") _ -> pure (Left "Failed to parse embedding response") else pure (Left ("Embedding HTTP error: " <> tshow status)) -- | Convert a vector to a blob for storage. vectorToBlob :: VS.Vector Float -> BS.ByteString vectorToBlob v = let bytes = VS.unsafeCast v :: VS.Vector Word8 in BS.pack (VS.toList bytes) -- | Convert a blob back to a vector. blobToVector :: BS.ByteString -> VS.Vector Float blobToVector bs = let bytes = VS.fromList (BS.unpack bs) :: VS.Vector Word8 in VS.unsafeCast bytes -- | Calculate cosine similarity between two vectors. cosineSimilarity :: VS.Vector Float -> VS.Vector Float -> Float cosineSimilarity v1 v2 | VS.length v1 /= VS.length v2 = 0 | otherwise = let dot = VS.sum (VS.zipWith (*) v1 v2) mag1 = sqrt (VS.sum (VS.map (\x -> x * x) v1)) mag2 = sqrt (VS.sum (VS.map (\x -> x * x) v2)) in if mag1 == 0 || mag2 == 0 then 0 else dot / (mag1 * mag2) -- | Format memories for inclusion in a prompt. formatMemoriesForPrompt :: [Memory] -> Text formatMemoriesForPrompt [] = "No prior context available." formatMemoriesForPrompt mems = Text.unlines [ "Known context about this user:", "", Text.unlines (map formatMem mems) ] where formatMem m = "- " <> memoryContent m <> " (via " <> sourceAgent (memorySource m) <> ")" -- | Run an agent with memory context. -- Recalls relevant memories for the user and injects them into the system prompt. runAgentWithMemory :: User -> Engine.EngineConfig -> Engine.AgentConfig -> Text -> IO (Either Text Engine.AgentResult) runAgentWithMemory user engineCfg agentCfg userPrompt = do memories <- recallMemories (userId user) userPrompt 10 let memoryContext = formatMemoriesForPrompt memories enhancedPrompt = Engine.agentSystemPrompt agentCfg <> "\n\n## Known about this user\n" <> memoryContext enhancedConfig = agentCfg { Engine.agentSystemPrompt = enhancedPrompt, Engine.agentTools = Engine.agentTools agentCfg <> [ rememberTool (userId user), recallTool (userId user), linkMemoriesTool (userId user), queryGraphTool (userId user) ] } Engine.runAgent engineCfg enhancedConfig userPrompt -- | Tool for agents to store memories about users. rememberTool :: Text -> Engine.Tool rememberTool uid = Engine.Tool { Engine.toolName = "remember", Engine.toolDescription = "Store a piece of information about the user for future reference. " <> "Use this when the user shares personal facts, preferences, or context " <> "that would be useful to recall in future conversations.", Engine.toolJsonSchema = Aeson.object [ "type" .= ("object" :: Text), "properties" .= Aeson.object [ "content" .= Aeson.object [ "type" .= ("string" :: Text), "description" .= ("The information to remember about the user" :: Text) ], "context" .= Aeson.object [ "type" .= ("string" :: Text), "description" .= ("How/why this was learned (e.g., 'user mentioned in chat')" :: Text) ], "tags" .= Aeson.object [ "type" .= ("array" :: Text), "items" .= Aeson.object ["type" .= ("string" :: Text)], "description" .= ("Optional tags for categorization" :: Text) ] ], "required" .= (["content", "context"] :: [Text]) ], Engine.toolExecute = executeRemember uid } executeRemember :: Text -> Aeson.Value -> IO Aeson.Value executeRemember uid v = case Aeson.fromJSON v of Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e]) Aeson.Success (args :: RememberArgs) -> do let source = MemorySource { sourceAgent = "agent", sourceSession = Nothing, sourceContext = rememberContext args } mem <- storeMemoryWithTags uid (rememberContent args) source (rememberTags args) pure ( Aeson.object [ "success" .= True, "memory_id" .= memoryId mem, "message" .= ("Remembered: " <> rememberContent args) ] ) -- | Tool for agents to recall memories about users. recallTool :: Text -> Engine.Tool recallTool uid = Engine.Tool { Engine.toolName = "recall", Engine.toolDescription = "Search your memory for information about the user. " <> "Use this to retrieve previously stored facts, preferences, or context.", Engine.toolJsonSchema = Aeson.object [ "type" .= ("object" :: Text), "properties" .= Aeson.object [ "query" .= Aeson.object [ "type" .= ("string" :: Text), "description" .= ("What to search for in memory" :: Text) ], "limit" .= Aeson.object [ "type" .= ("integer" :: Text), "description" .= ("Maximum memories to return (default: 5)" :: Text) ] ], "required" .= (["query"] :: [Text]) ], Engine.toolExecute = executeRecall uid } executeRecall :: Text -> Aeson.Value -> IO Aeson.Value executeRecall uid v = case Aeson.fromJSON v of Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e]) Aeson.Success (args :: RecallArgs) -> do mems <- recallMemories uid (recallQuery args) (recallLimit args) pure ( Aeson.object [ "success" .= True, "count" .= length mems, "memories" .= map ( \m -> Aeson.object [ "id" .= memoryId m, "content" .= memoryContent m, "confidence" .= memoryConfidence m, "source" .= sourceAgent (memorySource m), "tags" .= memoryTags m ] ) mems ] ) -- Helper for parsing remember args data RememberArgs = RememberArgs { rememberContent :: Text, rememberContext :: Text, rememberTags :: [Text] } deriving (Generic) instance Aeson.FromJSON RememberArgs where parseJSON = Aeson.withObject "RememberArgs" <| \v -> (RememberArgs (v .:? "context" .!= "agent observation") <*> (v .:? "tags" .!= []) data RecallArgs = RecallArgs { recallQuery :: Text, recallLimit :: Int } deriving (Generic) instance Aeson.FromJSON RecallArgs where parseJSON = Aeson.withObject "RecallArgs" <| \v -> (RecallArgs (v .:? "limit" .!= 5) -- | Tool for agents to link memories in the knowledge graph. linkMemoriesTool :: Text -> Engine.Tool linkMemoriesTool _uid = Engine.Tool { Engine.toolName = "link_memories", Engine.toolDescription = "Create a typed relationship between two memories. " <> "Use this to connect related information. Relation types:\n" <> "- contradicts: conflicting information\n" <> "- supports: evidence that reinforces another memory\n" <> "- elaborates: adds detail to an existing memory\n" <> "- supersedes: newer info replaces older\n" <> "- related: general topical connection\n" <> "- contingent_on: depends on another fact being true", Engine.toolJsonSchema = Aeson.object [ "type" .= ("object" :: Text), "properties" .= Aeson.object [ "from_memory_id" .= Aeson.object [ "type" .= ("string" :: Text), "description" .= ("ID of the source memory" :: Text) ], "to_memory_id" .= Aeson.object [ "type" .= ("string" :: Text), "description" .= ("ID of the target memory" :: Text) ], "relation_type" .= Aeson.object [ "type" .= ("string" :: Text), "enum" .= (["contradicts", "supports", "elaborates", "supersedes", "related", "contingent_on"] :: [Text]), "description" .= ("Type of relationship between memories" :: Text) ] ], "required" .= (["from_memory_id", "to_memory_id", "relation_type"] :: [Text]) ], Engine.toolExecute = executeLinkMemories } executeLinkMemories :: Aeson.Value -> IO Aeson.Value executeLinkMemories v = case Aeson.fromJSON v of Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e]) Aeson.Success (args :: LinkMemoriesArgs) -> do case textToRelationType (linkArgsRelationType args) of Nothing -> pure ( Aeson.object [ "success" .= False, "error" .= ("Invalid relation type: " <> linkArgsRelationType args) ] ) Just relType -> do memLink <- linkMemories (linkArgsFromId args) (linkArgsToId args) relType pure ( Aeson.object [ "success" .= True, "message" .= ( "Linked memory " <> linkFromMemoryId memLink <> " -> " <> linkToMemoryId memLink <> " (" <> relationTypeToText (linkRelationType memLink) <> ")" ) ] ) data LinkMemoriesArgs = LinkMemoriesArgs { linkArgsFromId :: Text, linkArgsToId :: Text, linkArgsRelationType :: Text } deriving (Generic) instance Aeson.FromJSON LinkMemoriesArgs where parseJSON = Aeson.withObject "LinkMemoriesArgs" <| \v -> (LinkMemoriesArgs (v .: "to_memory_id") <*> (v .: "relation_type") -- | Tool for agents to query the memory knowledge graph. queryGraphTool :: Text -> Engine.Tool queryGraphTool _uid = Engine.Tool { Engine.toolName = "query_graph", Engine.toolDescription = "Explore the knowledge graph to find related memories. " <> "Given a starting memory, traverse links to find connected memories. " <> "Useful for understanding context and finding contradictions or supporting evidence.", Engine.toolJsonSchema = Aeson.object [ "type" .= ("object" :: Text), "properties" .= Aeson.object [ "memory_id" .= Aeson.object [ "type" .= ("string" :: Text), "description" .= ("ID of the memory to start from" :: Text) ], "depth" .= Aeson.object [ "type" .= ("integer" :: Text), "description" .= ("How many link hops to traverse (default: 2)" :: Text) ], "relation_type" .= Aeson.object [ "type" .= ("string" :: Text), "enum" .= (["contradicts", "supports", "elaborates", "supersedes", "related", "contingent_on"] :: [Text]), "description" .= ("Optional: filter by relation type" :: Text) ] ], "required" .= (["memory_id"] :: [Text]) ], Engine.toolExecute = executeQueryGraph } executeQueryGraph :: Aeson.Value -> IO Aeson.Value executeQueryGraph v = case Aeson.fromJSON v of Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e]) Aeson.Success (args :: QueryGraphArgs) -> do let maybeRelType = queryArgsRelationType args +> textToRelationType results <- queryGraph (queryArgsMemoryId args) (queryArgsDepth args) maybeRelType pure ( Aeson.object [ "success" .= True, "count" .= length results, "memories" .= map ( \(m, links) -> Aeson.object [ "id" .= memoryId m, "content" .= memoryContent m, "links" .= map ( \l -> Aeson.object [ "from" .= linkFromMemoryId l, "to" .= linkToMemoryId l, "relation" .= linkRelationType l ] ) links ] ) results ] ) data QueryGraphArgs = QueryGraphArgs { queryArgsMemoryId :: Text, queryArgsDepth :: Int, queryArgsRelationType :: Maybe Text } deriving (Generic) instance Aeson.FromJSON QueryGraphArgs where parseJSON = Aeson.withObject "QueryGraphArgs" <| \v -> (QueryGraphArgs (v .:? "depth" .!= 2) <*> (v .:? "relation_type") -- | Estimate token count for text (rough: ~4 chars per token). estimateTokens :: Text -> Int estimateTokens t = max 1 (Text.length t `div` 4) -- | Save a message to conversation history. saveMessage :: Text -> Int -> MessageRole -> Maybe Text -> Text -> IO ConversationMessage saveMessage uid chatId role senderName content = do now <- getCurrentTime let tokens = estimateTokens content withMemoryDb <| \conn -> do SQL.execute conn "INSERT INTO conversation_messages (user_id, chat_id, role, sender_name, content, tokens_estimate, created_at) VALUES (?, ?, ?, ?, ?, ?, ?)" (uid, chatId, roleToText role, senderName, content, tokens, now) rowId <- SQL.lastInsertRowId conn pure ConversationMessage { cmId = Just (fromIntegral rowId), cmUserId = uid, cmChatId = chatId, cmRole = role, cmSenderName = senderName, cmContent = content, cmTokensEstimate = tokens, cmCreatedAt = now } where roleToText UserRole = "user" :: Text roleToText AssistantRole = "assistant" -- | Get recent messages for a user/chat, newest first. getRecentMessages :: Text -> Int -> Int -> IO [ConversationMessage] getRecentMessages uid chatId limit = withMemoryDb <| \conn -> SQL.query conn "SELECT id, user_id, chat_id, role, sender_name, content, tokens_estimate, created_at \ \FROM conversation_messages \ \WHERE user_id = ? AND chat_id = ? \ \ORDER BY created_at DESC LIMIT ?" (uid, chatId, limit) -- | Get the most recent summary for a chat. getLatestSummary :: Text -> Int -> IO (Maybe ConversationSummary) getLatestSummary uid chatId = withMemoryDb <| \conn -> do rows <- SQL.query conn "SELECT id, user_id, chat_id, summary, messages_summarized, tokens_saved, created_at \ \FROM conversation_summaries \ \WHERE user_id = ? AND chat_id = ? \ \ORDER BY created_at DESC LIMIT 1" (uid, chatId) pure (listToMaybe rows) -- | Build conversation context for the LLM. -- Returns (context text, total token estimate). getConversationContext :: Text -> Int -> Int -> IO (Text, Int) getConversationContext uid chatId maxTokens = do maybeSummary <- getLatestSummary uid chatId recentMsgs <- getRecentMessages uid chatId 50 let summaryText = maybe "" (\s -> "## Previous conversation summary\n" <> csSummary s <> "\n\n") maybeSummary summaryTokens = maybe 0 (estimateTokens <. csSummary) maybeSummary msgsOldestFirst = reverse recentMsgs availableTokens = maxTokens - summaryTokens - 100 (selectedMsgs, usedTokens) = selectMessages msgsOldestFirst availableTokens formattedMsgs = if null selectedMsgs then "" else "## Recent conversation\n" <> Text.unlines (map formatMsg selectedMsgs) pure (summaryText <> formattedMsgs, summaryTokens + usedTokens) where selectMessages :: [ConversationMessage] -> Int -> ([ConversationMessage], Int) selectMessages msgs budget = go (reverse msgs) budget [] where go [] _ acc = (acc, sum (map cmTokensEstimate acc)) go (m : ms) remaining acc | cmTokensEstimate m <= remaining = go ms (remaining - cmTokensEstimate m) (m : acc) | otherwise = (acc, sum (map cmTokensEstimate acc)) formatMsg m = let timestamp = Text.pack (formatTime defaultTimeLocale "%Y-%m-%dT%H:%M:%SZ" (cmCreatedAt m)) prefix = case cmRole m of UserRole -> "[" <> timestamp <> "] " <> fromMaybe "User" (cmSenderName m) <> ": " AssistantRole -> "[" <> timestamp <> "] Assistant: " in prefix <> cmContent m -- | Summarize old messages and archive them. -- Returns the new summary text. summarizeAndArchive :: Text -> Int -> Text -> IO Text summarizeAndArchive uid chatId summaryText = do now <- getCurrentTime (oldMsgCount, tokensSaved) <- withMemoryDb <| \conn -> do rows <- SQL.query conn "SELECT COUNT(*), COALESCE(SUM(tokens_estimate), 0) FROM conversation_messages WHERE user_id = ? AND chat_id = ?" (uid, chatId) :: IO [(Int, Int)] let (count, tokens) = fromMaybe (0, 0) (listToMaybe rows) SQL.execute conn "INSERT INTO conversation_summaries (user_id, chat_id, summary, messages_summarized, tokens_saved, created_at) VALUES (?, ?, ?, ?, ?, ?)" (uid, chatId, summaryText, count, tokens, now) SQL.execute conn "DELETE FROM conversation_messages WHERE user_id = ? AND chat_id = ?" (uid, chatId) pure (count, tokens) putText <| "Archived " <> tshow oldMsgCount <> " messages (" <> tshow tokensSaved <> " tokens) for chat " <> tshow chatId pure summaryText -- ----------------------------------------------------------------------------- -- Group Conversation History -- ----------------------------------------------------------------------------- -- | Save a message to group conversation history. -- Unlike saveMessage, this is keyed by (chat_id, thread_id) not (user_id, chat_id). -- The sender_name is preserved for attribution. saveGroupMessage :: Int -> Maybe Int -> MessageRole -> Text -> Text -> IO ConversationMessage saveGroupMessage chatId mThreadId role senderName content = do now <- getCurrentTime let tokens = estimateTokens content withMemoryDb <| \conn -> do SQL.execute conn "INSERT INTO conversation_messages (user_id, chat_id, thread_id, role, sender_name, content, tokens_estimate, created_at) VALUES (NULL, ?, ?, ?, ?, ?, ?, ?)" (chatId, mThreadId, roleToText role, senderName, content, tokens, now) rowId <- SQL.lastInsertRowId conn pure ConversationMessage { cmId = Just (fromIntegral rowId), cmUserId = "", cmChatId = chatId, cmRole = role, cmSenderName = Just senderName, cmContent = content, cmTokensEstimate = tokens, cmCreatedAt = now } where roleToText UserRole = "user" :: Text roleToText AssistantRole = "assistant" -- | Get recent messages for a group chat/topic, newest first. getGroupRecentMessages :: Int -> Maybe Int -> Int -> IO [ConversationMessage] getGroupRecentMessages chatId mThreadId limit = withMemoryDb <| \conn -> case mThreadId of Just threadId -> SQL.query conn "SELECT id, COALESCE(user_id, ''), chat_id, role, sender_name, content, tokens_estimate, created_at \ \FROM conversation_messages \ \WHERE chat_id = ? AND thread_id = ? \ \ORDER BY created_at DESC LIMIT ?" (chatId, threadId, limit) Nothing -> SQL.query conn "SELECT id, COALESCE(user_id, ''), chat_id, role, sender_name, content, tokens_estimate, created_at \ \FROM conversation_messages \ \WHERE chat_id = ? AND thread_id IS NULL \ \ORDER BY created_at DESC LIMIT ?" (chatId, limit) -- | Build conversation context for a group chat. -- Returns (context text, total token estimate). getGroupConversationContext :: Int -> Maybe Int -> Int -> IO (Text, Int) getGroupConversationContext chatId mThreadId maxTokens = do recentMsgs <- getGroupRecentMessages chatId mThreadId 50 let msgsOldestFirst = reverse recentMsgs availableTokens = maxTokens - 100 (selectedMsgs, usedTokens) = selectMessages msgsOldestFirst availableTokens formattedMsgs = if null selectedMsgs then "" else "## Recent conversation\n" <> Text.unlines (map formatMsg selectedMsgs) pure (formattedMsgs, usedTokens) where selectMessages :: [ConversationMessage] -> Int -> ([ConversationMessage], Int) selectMessages msgs budget = go (reverse msgs) budget [] where go [] _ acc = (acc, sum (map cmTokensEstimate acc)) go (m : ms) remaining acc | cmTokensEstimate m <= remaining = go ms (remaining - cmTokensEstimate m) (m : acc) | otherwise = (acc, sum (map cmTokensEstimate acc)) formatMsg m = let timestamp = Text.pack (formatTime defaultTimeLocale "%Y-%m-%dT%H:%M:%SZ" (cmCreatedAt m)) prefix = case cmRole m of UserRole -> "[" <> timestamp <> "] " <> fromMaybe "User" (cmSenderName m) <> ": " AssistantRole -> "[" <> timestamp <> "] Assistant: " in prefix <> cmContent m -- ----------------------------------------------------------------------------- -- Group Memories -- ----------------------------------------------------------------------------- -- | Generate a synthetic user_id for group-level memories. groupUserId :: Int -> Text groupUserId chatId = "group:" <> tshow chatId -- | Store a memory associated with a group (not a user). -- These memories are shared across all users in the group. storeGroupMemory :: Int -> Text -> MemorySource -> IO Memory storeGroupMemory chatId = storeMemory (groupUserId chatId) -- | Recall memories for a group. recallGroupMemories :: Int -> Text -> Int -> IO [Memory] recallGroupMemories chatId = recallMemories (groupUserId chatId)