summaryrefslogtreecommitdiff
path: root/Omni/Agent/Memory.hs
diff options
context:
space:
mode:
Diffstat (limited to 'Omni/Agent/Memory.hs')
-rw-r--r--Omni/Agent/Memory.hs1575
1 files changed, 1575 insertions, 0 deletions
diff --git a/Omni/Agent/Memory.hs b/Omni/Agent/Memory.hs
new file mode 100644
index 0000000..4aaa438
--- /dev/null
+++ b/Omni/Agent/Memory.hs
@@ -0,0 +1,1575 @@
+{-# 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 .: "id"))
+ <*> (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
+ <*> 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 .: "agent"))
+ <*> (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 .: "id")
+ )
+ <*> (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 </ (embeddingBlob :: Maybe BS.ByteString)
+ source = MemorySource agent session context
+ tags = fromMaybe [] ((tagsJson :: Maybe Text) +> (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
+ <*> SQL.field
+ <*> (parseRole </ SQL.field)
+ <*> SQL.field
+ <*> SQL.field
+ <*> (fromMaybe 0 </ SQL.field)
+ <*> 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
+ <*> 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 .: "from_memory_id"))
+ <*> (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 </ memoryEmbedding mem,
+ sourceAgent (memorySource mem),
+ sourceSession (memorySource mem),
+ sourceContext (memorySource mem)
+ )
+ SQL.:. ( memoryConfidence mem,
+ memoryCreatedAt mem,
+ memoryLastAccessedAt mem,
+ TE.decodeUtf8 (BL.toStrict (Aeson.encode (memoryTags mem)))
+ )
+ )
+ pure mem
+
+-- | Recall memories for a user using semantic similarity.
+recallMemories :: Text -> 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 </ traverse (`getLinkedMemories` maybeRelType) newIds
+ let newMems = [(m, [l]) | (l, m) <- newLinked, memoryId m `notElem` visitedIds]
+ newVisited = visitedIds <> 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" </ lookupEnv "OLLAMA_URL"
+ let url = ollamaUrl <> "/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 .: "content"))
+ <*> (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 .: "query"))
+ <*> (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 .: "from_memory_id"))
+ <*> (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 .: "memory_id"))
+ <*> (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)
generated by cgit v1.2.3 (git 2.25.1) at 2025-12-19 00:25:23 +0000