commit 8bce3f84468c481b29b1f714044d99287d82d944
Author: Ben Sima <ben@bensima.com>
Date: Mon Feb 9 17:53:59 2026
Unify passive/active context hydration
Unify the context builder so all channels (Telegram, CLI) share the same passive/active hydration paths, and remove bespoke Telegram context assembly/compaction heuristics that drifted from core behavior.
Design decisions:
- Introduce ConversationScope for DM vs group scoping, keeping semantic search chat-scoped only.
- Map passive-temporal/semantic/knowledge sources into hydration with explicit section IDs.
- Add context-confidence warnings to passive-temporal sections to nudge tool backfills.
- Use group memory owner IDs for group-scoped long-term memories.
Task-Id: t-574
diff --git a/Omni/Agent.hs b/Omni/Agent.hs
index afdc04e4..02e9934b 100755
--- a/Omni/Agent.hs
+++ b/Omni/Agent.hs
@@ -49,12 +49,12 @@ import qualified Data.Time as Time
import qualified Data.Time.Format as TimeFormat
import qualified Data.Yaml as Yaml
import qualified Omni.Agent.Auth as Auth
+import qualified Omni.Agent.Context as Context
import qualified Omni.Agent.Engine as Engine
import qualified Omni.Agent.Interpreter.Sequential as Seq
import qualified Omni.Agent.Memory as Memory
import qualified Omni.Agent.Op as Op
import qualified Omni.Agent.Programs.Agent as OpAgent
-import qualified Omni.Agent.Prompt.MemorySources as MemorySources
import qualified Omni.Agent.Provider as Provider
import qualified Omni.Agent.Tools as Tools
import qualified Omni.Agent.Trace as Trace
@@ -798,8 +798,8 @@ runAgent opts = do
-- Resolve memory config from environment
envMemoryUser <- Environment.lookupEnv "MEMORY_USER_ID"
envMemoryContext <- Environment.lookupEnv "MEMORY_CONTEXT"
- let memoryUser = Text.pack <$> envMemoryUser
- memoryContext = Text.pack <$> envMemoryContext
+ let memoryUser = Text.pack </ envMemoryUser
+ memoryContext = Text.pack </ envMemoryContext
-- Add memory tools if memory-user is provided
let memoryTools = case memoryUser of
@@ -895,12 +895,15 @@ runAgent opts = do
mHydrationConfig = case memoryUser of
Just uid ->
let memCtx = fromMaybe ("cli:" <> fromMaybe "session" (optRunId opts')) memoryContext
- memIdentity = Memory.MemoryIdentity uid memCtx
+ conversationScope =
+ case Context.parseChatIdFromContext memCtx of
+ Just chatId -> Context.DirectConversation uid chatId
+ Nothing -> Context.OwnerOnly uid
staticSections =
- [ MemorySources.mkTimeSection now tz,
- MemorySources.mkProjectSection "agent" (Text.pack cwd)
+ [ Context.mkTimeSection now tz,
+ Context.mkProjectSection "agent" (Text.pack cwd)
]
- in Just <| MemorySources.buildHydrationConfig systemPrompt tools staticSections memIdentity
+ in Just <| Context.buildContextHydrationConfig systemPrompt tools staticSections conversationScope
Nothing -> Nothing
seqConfig =
diff --git a/Omni/Agent/Context.hs b/Omni/Agent/Context.hs
new file mode 100644
index 00000000..25431f24
--- /dev/null
+++ b/Omni/Agent/Context.hs
@@ -0,0 +1,523 @@
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE NoImplicitPrelude #-}
+
+-- | Core context hydration sources and tools.
+--
+-- This module defines the four context hydration modes:
+-- - passive-temporal: recent messages window
+-- - passive-semantic: related messages + long-term memories
+-- - active-temporal: time-range lookup tools
+-- - active-semantic: semantic search tools
+module Omni.Agent.Context
+ ( ConversationScope (..),
+ scopeChatId,
+ scopeThreadId,
+ scopeMemoryOwner,
+ parseChatIdFromContext,
+
+ -- * Passive Context Sources
+ passiveTemporalSource,
+ passiveSemanticSource,
+ passiveSemanticKnowledgeSource,
+
+ -- * Hydration Config Builder
+ buildContextHydrationConfig,
+
+ -- * Active Context Tools
+ activeTemporalTools,
+ activeSemanticTools,
+ memoryWriteTools,
+
+ -- * Static Sections
+ mkProjectSection,
+ mkTimeSection,
+
+ -- * Helpers
+ estimateTokensSimple,
+
+ -- * Testing
+ main,
+ test,
+ )
+where
+
+import Alpha
+import qualified Data.Aeson as Aeson
+import qualified Data.Text as Text
+import qualified Data.Time as Time
+import qualified Database.SQLite.Simple as SQL
+import qualified Omni.Agent.Engine as Engine
+import qualified Omni.Agent.Memory as Memory
+import Omni.Agent.Prompt.Hydrate
+ ( ContextSource (..),
+ HydrationConfig (..),
+ KnowledgeResult (..),
+ SemanticResult (..),
+ SourceParams (..),
+ TemporalResult (..),
+ defaultHydrationConfig,
+ )
+import Omni.Agent.Prompt.IR
+ ( CompositionMode (..),
+ Priority (..),
+ Section (..),
+ SectionSource (..),
+ ToolDef (..),
+ )
+import qualified Omni.Test as Test
+import qualified System.Directory as Dir
+import qualified System.Environment as Env
+import System.FilePath ((</>))
+
+-- | Conversation scope for context hydration.
+data ConversationScope
+ = DirectConversation
+ { csOwner :: Text,
+ csChatId :: Int
+ }
+ | GroupConversation
+ { csChatId :: Int,
+ csThreadId :: Maybe Int
+ }
+ | OwnerOnly
+ { csOwner :: Text
+ }
+ deriving (Show, Eq)
+
+scopeChatId :: ConversationScope -> Maybe Int
+scopeChatId scope = case scope of
+ DirectConversation _ chatId -> Just chatId
+ GroupConversation chatId _ -> Just chatId
+ OwnerOnly _ -> Nothing
+
+scopeThreadId :: ConversationScope -> Maybe Int
+scopeThreadId scope = case scope of
+ GroupConversation _ threadId -> threadId
+ _ -> Nothing
+
+scopeMemoryOwner :: ConversationScope -> Text
+scopeMemoryOwner scope = case scope of
+ DirectConversation owner _ -> owner
+ GroupConversation chatId _ -> Memory.groupMemoryOwner chatId
+ OwnerOnly owner -> owner
+
+-- | Parse a chat ID from a context string.
+--
+-- Context formats:
+-- "telegram:12345" -> Just 12345
+-- "cli:session-abc" -> Nothing
+-- "12345" -> Just 12345 (legacy format)
+parseChatIdFromContext :: Text -> Maybe Int
+parseChatIdFromContext ctx
+ | "telegram:" `Text.isPrefixOf` ctx = readMaybe (Text.unpack (Text.drop 9 ctx))
+ | otherwise = readMaybe (Text.unpack ctx)
+
+-- | Create a passive-temporal context source from conversation history.
+passiveTemporalSource :: ConversationScope -> ContextSource TemporalResult
+passiveTemporalSource scope =
+ ContextSource <| \params -> do
+ messages <- case scope of
+ DirectConversation owner chatId ->
+ Memory.getRecentMessages owner chatId (spMaxItems params)
+ GroupConversation chatId threadId ->
+ Memory.getGroupRecentMessages chatId threadId (spMaxItems params)
+ OwnerOnly _ -> pure []
+
+ let formatted =
+ [ ( Memory.cmCreatedAt msg,
+ roleToText (Memory.cmRole msg),
+ fromMaybe "unknown" (Memory.cmSenderName msg),
+ Memory.cmContent msg
+ )
+ | msg <- reverse messages -- oldest first for context
+ ]
+ totalTokens = sum [Memory.cmTokensEstimate msg | msg <- messages]
+ recentContents = map Memory.cmContent messages
+
+ warning <- buildContextWarning scope (spObservation params) recentContents
+
+ pure
+ TemporalResult
+ { trMessages = formatted,
+ trTotalTokens = totalTokens,
+ trContextWarning = warning
+ }
+
+-- | Create a passive-semantic context source from chat history.
+passiveSemanticSource :: ConversationScope -> ContextSource SemanticResult
+passiveSemanticSource scope =
+ ContextSource <| \params -> do
+ let mChatId = scopeChatId scope
+ results <- case mChatId of
+ Nothing -> pure []
+ Just chatId -> Memory.searchChatHistoryInChat chatId (spObservation params) (spMaxItems params)
+
+ let filtered =
+ [ (Memory.cheCreatedAt entry, Memory.cheRole entry, Memory.cheContent entry, score)
+ | (entry, score) <- results,
+ score >= spThreshold params
+ ]
+ totalTokens = sum [estimateTokensSimple (Memory.cheContent entry) | (entry, _) <- results]
+
+ pure
+ SemanticResult
+ { srMatches = filtered,
+ srTotalTokens = totalTokens
+ }
+
+-- | Passive semantic context from long-term memories.
+passiveSemanticKnowledgeSource :: ConversationScope -> ContextSource KnowledgeResult
+passiveSemanticKnowledgeSource scope =
+ ContextSource <| \params -> do
+ let owner = scopeMemoryOwner scope
+ result <- Memory.recallMemoriesEnhanced owner (spObservation params) (spMaxItems params)
+
+ let memories = map Memory.smMemory (Memory.errMemories result)
+ formatted =
+ [ ( Memory.memoryContent mem,
+ sourceToContext (Memory.memorySource mem),
+ Memory.memoryTags mem
+ )
+ | mem <- memories
+ ]
+ contradictions =
+ [ (Memory.memoryContent m1, Memory.memoryContent m2)
+ | (m1, m2) <- Memory.errContradictions result
+ ]
+ totalTokens = sum [estimateTokensSimple (Memory.memoryContent mem) | mem <- memories]
+
+ pure
+ KnowledgeResult
+ { krMemories = formatted,
+ krContradictions = contradictions,
+ krSupersededCount = length (Memory.errSuperseded result),
+ krTotalTokens = totalTokens
+ }
+
+-- | Build a HydrationConfig with passive context sources.
+buildContextHydrationConfig :: Text -> [Engine.Tool] -> [Section] -> ConversationScope -> HydrationConfig
+buildContextHydrationConfig systemPrompt tools staticSections scope =
+ defaultHydrationConfig
+ { hcSystemPrompt = systemPrompt,
+ hcTools = map engineToolToToolDef tools,
+ hcStaticSections = staticSections,
+ hcTemporalSource = Just (passiveTemporalSource scope),
+ hcSemanticSource = Just (passiveSemanticSource scope),
+ hcKnowledgeSource = Just (passiveSemanticKnowledgeSource scope),
+ hcEstimateTokens = estimateTokensSimple
+ }
+
+-- | Active temporal tools (time-range lookups).
+activeTemporalTools :: ConversationScope -> [Engine.Tool]
+activeTemporalTools scope = case scope of
+ DirectConversation owner chatId -> [Memory.getMessagesByTimeTool owner chatId]
+ GroupConversation chatId threadId -> [Memory.getGroupMessagesByTimeTool chatId threadId]
+ OwnerOnly _ -> []
+
+-- | Active semantic tools (search and memory graph).
+activeSemanticTools :: ConversationScope -> [Engine.Tool]
+activeSemanticTools scope =
+ let owner = scopeMemoryOwner scope
+ chatTools = case scopeChatId scope of
+ Nothing -> []
+ Just chatId -> [Memory.searchChatHistoryTool chatId]
+ in chatTools
+ <> [ Memory.recallTool owner,
+ Memory.queryGraphTool owner,
+ Memory.lookupPastDiscussionsTool owner
+ ]
+
+-- | Memory write tools (remember/link).
+memoryWriteTools :: ConversationScope -> [Engine.Tool]
+memoryWriteTools scope =
+ let owner = scopeMemoryOwner scope
+ in [Memory.rememberTool owner, Memory.linkMemoriesTool owner]
+
+-- | Convert an Engine.Tool to IR.ToolDef.
+engineToolToToolDef :: Engine.Tool -> ToolDef
+engineToolToToolDef tool =
+ ToolDef
+ { tdName = Engine.toolName tool,
+ tdDescription = Engine.toolDescription tool,
+ tdSchema = Engine.toolJsonSchema tool,
+ tdPriority = High,
+ tdEmbedding = Nothing
+ }
+
+-- | Simple token estimation (roughly 4 characters per token).
+estimateTokensSimple :: Text -> Int
+estimateTokensSimple t = max 1 (Text.length t `div` 4)
+
+-- | Convert MessageRole to Text.
+roleToText :: Memory.MessageRole -> Text
+roleToText Memory.UserRole = "user"
+roleToText Memory.AssistantRole = "assistant"
+
+-- | Convert MemorySource to a context description.
+sourceToContext :: Memory.MemorySource -> Text
+sourceToContext = Memory.sourceContext
+
+-- | Create a static section for project context.
+mkProjectSection :: Text -> Text -> Section
+mkProjectSection projectName workingDir =
+ Section
+ { secId = "project",
+ secLabel = "Current Project",
+ secSource = SourceState "project",
+ secContent =
+ Text.unlines
+ [ "Project: " <> projectName,
+ "Working directory: " <> workingDir
+ ],
+ secTokens = 20,
+ secMinTokens = Nothing,
+ secPriority = High,
+ secRelevance = Nothing,
+ secRecency = Nothing,
+ secCompositionMode = Constraint,
+ secEmbedding = Nothing,
+ secHash = Nothing
+ }
+
+-- | Create a static section for current time.
+mkTimeSection :: Time.UTCTime -> Time.TimeZone -> Section
+mkTimeSection utcTime tz =
+ let localTime = Time.utcToLocalTime tz utcTime
+ timeStr = Text.pack (Time.formatTime Time.defaultTimeLocale "%A, %B %d, %Y at %H:%M" localTime)
+ in Section
+ { secId = "time",
+ secLabel = "Current Date and Time",
+ secSource = SourceState "clock",
+ secContent = timeStr,
+ secTokens = 15,
+ secMinTokens = Nothing,
+ secPriority = High,
+ secRelevance = Nothing,
+ secRecency = Just utcTime,
+ secCompositionMode = Contextual,
+ secEmbedding = Nothing,
+ secHash = Nothing
+ }
+
+-- | Build a context warning if recent context is likely insufficient.
+buildContextWarning :: ConversationScope -> Text -> [Text] -> IO (Maybe Text)
+buildContextWarning scope observation recentContents = do
+ let mChatId = scopeChatId scope
+ case mChatId of
+ Nothing -> pure Nothing
+ Just chatId -> do
+ now <- Time.getCurrentTime
+ queryEmbedding <- Memory.embedText observation
+ case queryEmbedding of
+ Left _ -> pure Nothing
+ Right qEmb -> do
+ (recentSimilarity, recentGapHours) <-
+ Memory.getRecentChatHistorySimilarity chatId qEmb (max 10 (length recentContents)) now
+ semanticHits <- Memory.searchChatHistoryInChatWithEmbedding chatId qEmb (contextWarningSemanticLimit * 2) now
+
+ let bestSemantic = listToMaybe semanticHits
+ bestSemanticScore = snd </ bestSemantic
+ bestSemanticIsRecent = maybe False (\(e, _) -> Memory.cheContent e `elem` recentContents) bestSemantic
+ recentScore = fromMaybe 0 recentSimilarity
+ topicDrift = case bestSemanticScore of
+ Just score ->
+ not bestSemanticIsRecent
+ && score
+ >= contextWarningSimilarityFloor
+ && score
+ - recentScore
+ >= contextWarningTopicDriftDelta
+ Nothing -> False
+ gapHours = fromMaybe 0 recentGapHours
+ baseConfidence
+ | gapHours >= contextWarningGapHours = ContextLow
+ | recentScore >= 0.75 = ContextHigh
+ | recentScore >= contextWarningSimilarityFloor = ContextMedium
+ | otherwise = ContextLow
+ confidence = if topicDrift then ContextLow else baseConfidence
+
+ pure <| case confidence of
+ ContextLow ->
+ Just
+ <| "## Context Confidence\n"
+ <> "LOW: recent context may be insufficient. Use search_chat_history or get_messages_by_time to backfill before responding.\n"
+ _ -> Nothing
+
+contextWarningSemanticLimit :: Int
+contextWarningSemanticLimit = 10
+
+contextWarningSimilarityFloor :: Float
+contextWarningSimilarityFloor = 0.6
+
+contextWarningTopicDriftDelta :: Float
+contextWarningTopicDriftDelta = 0.1
+
+contextWarningGapHours :: Float
+contextWarningGapHours = 8
+
+data ContextConfidence
+ = ContextHigh
+ | ContextMedium
+ | ContextLow
+ deriving (Show, Eq)
+
+--------------------------------------------------------------------------------
+-- Tests
+--------------------------------------------------------------------------------
+
+main :: IO ()
+main = Test.run test
+
+test :: Test.Tree
+test =
+ Test.group
+ "Omni.Agent.Context"
+ [ Test.unit "estimateTokensSimple basic" <| do
+ estimateTokensSimple "" Test.@=? 1
+ estimateTokensSimple "hello" Test.@=? 1
+ estimateTokensSimple "hello world test" Test.@=? 4
+ estimateTokensSimple (Text.replicate 100 "a") Test.@=? 25,
+ Test.unit "passiveTemporalSource returns DM messages" <| do
+ withTestDb <| \userId chatId -> do
+ let scope = DirectConversation userId chatId
+ source = passiveTemporalSource scope
+ params =
+ SourceParams
+ { spObservation = "test query",
+ spMaxTokens = 1000,
+ spMaxItems = 10,
+ spThreshold = 0.5
+ }
+ result <- runSource source params
+ length (trMessages result) Test.@=? 3
+ let contents = [content | (_, _, _, content) <- trMessages result]
+ case contents of
+ (msg1 : _ : msg3 : _) -> do
+ msg1 Test.@=? "First message"
+ msg3 Test.@=? "Third message"
+ _ -> Test.assertFailure "Expected 3 messages",
+ Test.unit "passiveTemporalSource returns group messages" <| do
+ withTestDb <| \_userId chatId -> do
+ let scope = GroupConversation chatId (Just testThreadId)
+ source = passiveTemporalSource scope
+ params =
+ SourceParams
+ { spObservation = "test query",
+ spMaxTokens = 1000,
+ spMaxItems = 10,
+ spThreshold = 0.5
+ }
+ result <- runSource source params
+ let contents = [content | (_, _, _, content) <- trMessages result]
+ contents Test.@=? ["Group first", "Group second"],
+ Test.unit "passiveSemanticKnowledgeSource runs without error" <| do
+ withTestDb <| \userId _chatId -> do
+ let scope = OwnerOnly userId
+ source = passiveSemanticKnowledgeSource scope
+ params =
+ SourceParams
+ { spObservation = "programming",
+ spMaxTokens = 1000,
+ spMaxItems = 10,
+ spThreshold = 0.5
+ }
+ result <- runSource source params
+ krTotalTokens result Test.@=? krTotalTokens result
+ ]
+
+-- | Run tests with a temporary database populated with test data.
+withTestDb :: (Text -> Int -> IO a) -> IO a
+withTestDb action = do
+ tmpDir <- Dir.getTemporaryDirectory
+ let testDbPath = tmpDir </> "context-test.db"
+
+ dbExists <- Dir.doesFileExist testDbPath
+ when dbExists <| Dir.removeFile testDbPath
+
+ Env.setEnv "MEMORY_DB_PATH" testDbPath
+
+ SQL.withConnection testDbPath <| \conn -> do
+ Memory.initMemoryDb conn
+ populateTestData conn
+
+ result <- action testUserId testChatId
+
+ Env.unsetEnv "MEMORY_DB_PATH"
+ Dir.removeFile testDbPath
+
+ pure result
+
+testUserId :: Text
+testUserId = "test-user-001"
+
+testChatId :: Int
+testChatId = 12345
+
+testThreadId :: Int
+testThreadId = 777
+
+populateTestData :: SQL.Connection -> IO ()
+populateTestData conn = do
+ now <- Time.getCurrentTime
+ let oneHourAgo = Time.addUTCTime (-3600) now
+ twoHoursAgo = Time.addUTCTime (-7200) now
+
+ SQL.execute
+ conn
+ "INSERT INTO users (id, name, created_at) VALUES (?, ?, ?)"
+ (testUserId, "Test User" :: Text, now)
+
+ SQL.execute
+ conn
+ "INSERT INTO conversation_messages (user_id, chat_id, role, sender_name, content, tokens_estimate, created_at) VALUES (?, ?, ?, ?, ?, ?, ?)"
+ (testUserId, testChatId, "user" :: Text, Just ("Test User" :: Text), "First message" :: Text, 10 :: Int, twoHoursAgo)
+
+ SQL.execute
+ conn
+ "INSERT INTO conversation_messages (user_id, chat_id, role, sender_name, content, tokens_estimate, created_at) VALUES (?, ?, ?, ?, ?, ?, ?)"
+ (testUserId, testChatId, "assistant" :: Text, Just ("Ava" :: Text), "Second message" :: Text, 10 :: Int, oneHourAgo)
+
+ SQL.execute
+ conn
+ "INSERT INTO conversation_messages (user_id, chat_id, role, sender_name, content, tokens_estimate, created_at) VALUES (?, ?, ?, ?, ?, ?, ?)"
+ (testUserId, testChatId, "user" :: Text, Just ("Test User" :: Text), "Third message" :: Text, 10 :: Int, now)
+
+ SQL.execute
+ conn
+ "INSERT INTO conversation_messages (user_id, chat_id, thread_id, role, sender_name, content, tokens_estimate, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?)"
+ ("group" :: Text, testChatId, testThreadId, "user" :: Text, Just ("Group User" :: Text), "Group first" :: Text, 10 :: Int, twoHoursAgo)
+
+ SQL.execute
+ conn
+ "INSERT INTO conversation_messages (user_id, chat_id, thread_id, role, sender_name, content, tokens_estimate, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?)"
+ ("group" :: Text, testChatId, testThreadId, "assistant" :: Text, Just ("Ava" :: Text), "Group second" :: Text, 10 :: Int, oneHourAgo)
+
+ SQL.execute
+ conn
+ "INSERT INTO memories (id, user_id, content, source_agent, source_session, source_context, confidence, created_at, last_accessed_at, tags) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"
+ ( "mem-001" :: Text,
+ testUserId,
+ "User likes programming in Haskell" :: Text,
+ "test" :: Text,
+ Nothing :: Maybe Text,
+ "test context" :: Text,
+ 0.9 :: Double,
+ now,
+ now,
+ "[\"programming\",\"haskell\"]" :: Text
+ )
+
+ SQL.execute
+ conn
+ "INSERT INTO memories (id, user_id, content, source_agent, source_session, source_context, confidence, created_at, last_accessed_at, tags) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"
+ ( "mem-002" :: Text,
+ testUserId,
+ "User works on the Omni project" :: Text,
+ "test" :: Text,
+ Nothing :: Maybe Text,
+ "test context" :: Text,
+ 0.8 :: Double,
+ now,
+ now,
+ "[\"work\",\"omni\"]" :: Text
+ )
diff --git a/Omni/Agent/Memory.hs b/Omni/Agent/Memory.hs
index c4db70c4..be0e40aa 100644
--- a/Omni/Agent/Memory.hs
+++ b/Omni/Agent/Memory.hs
@@ -78,11 +78,12 @@ module Omni.Agent.Memory
-- * Group Conversation History
saveGroupMessage,
getGroupRecentMessages,
- getGroupConversationContext,
-- * Group Memories
+ groupMemoryOwner,
storeGroupMemory,
recallGroupMemories,
+ recallGroupMemoriesEnhanced,
-- * Embeddings
embedText,
@@ -91,18 +92,16 @@ module Omni.Agent.Memory
saveChatHistoryEntry,
searchChatHistorySemantic,
searchChatHistoryInChat,
+ searchChatHistoryInChatWithEmbedding,
+ getRecentChatHistorySimilarity,
backfillChatHistory,
getChatHistoryStats,
getRecentChatHistory,
- -- * Adaptive Context (Hybrid Temporal + Semantic)
- getAdaptiveContext,
- AdaptiveContextConfig (..),
- defaultAdaptiveContextConfig,
-
-- * Chat History Tools
searchChatHistoryTool,
getMessagesByTimeTool,
+ getGroupMessagesByTimeTool,
-- * Agent Interaction Tracking
AgentInteraction (..),
@@ -2255,42 +2254,6 @@ getGroupRecentMessages chatId mThreadId limit =
\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 = formatEasternTime "%Y-%m-%d %H:%M" (cmCreatedAt m)
- prefix = case cmRole m of
- UserRole -> "[" <> timestamp <> "] " <> fromMaybe "User" (cmSenderName m) <> ": "
- AssistantRole -> "[" <> timestamp <> "] Assistant: "
- in prefix <> cmContent m
-
-- -----------------------------------------------------------------------------
-- Group Memories
-- -----------------------------------------------------------------------------
@@ -2299,6 +2262,9 @@ getGroupConversationContext chatId mThreadId maxTokens = do
groupUserId :: Int -> Text
groupUserId chatId = "group:" <> tshow chatId
+groupMemoryOwner :: Int -> Text
+groupMemoryOwner = groupUserId
+
-- | 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
@@ -2308,6 +2274,9 @@ storeGroupMemory chatId = storeMemory (groupUserId chatId)
recallGroupMemories :: Int -> Text -> Int -> IO [Memory]
recallGroupMemories chatId = recallMemories (groupUserId chatId)
+recallGroupMemoriesEnhanced :: Int -> Text -> Int -> IO EnhancedRecallResult
+recallGroupMemoriesEnhanced chatId = recallMemoriesEnhanced (groupUserId chatId)
+
-- | Save a chat history entry with embedding for semantic search.
saveChatHistoryEntry :: Int -> Maybe Text -> Text -> Maybe Text -> Text -> UTCTime -> IO ()
saveChatHistoryEntry chatId usrId role senderName content timestamp = do
@@ -2386,176 +2355,25 @@ searchChatHistoryInChatWithEmbedding chatId qEmb limit now = do
sorted = List.sortBy (\(_, s1) (_, s2) -> compare s2 s1) scored
pure (take limit sorted)
--- -----------------------------------------------------------------------------
--- Adaptive Context (Hybrid Temporal + Semantic)
--- -----------------------------------------------------------------------------
-
--- | Configuration for adaptive context retrieval.
-data AdaptiveContextConfig = AdaptiveContextConfig
- { -- | Maximum recent messages to include (temporal window)
- accTemporalLimit :: Int,
- -- | Maximum semantic matches to include
- accSemanticLimit :: Int,
- -- | Ratio of token budget for temporal context (e.g., 0.7 = 70%)
- accTemporalBudgetRatio :: Float,
- -- | Minimum similarity score for semantic matches
- accSimilarityThreshold :: Float,
- -- | Daily decay factor for semantic scores (e.g., 0.995)
- accRecencyDecay :: Float
- }
- deriving (Show, Eq, Generic)
-
-data ContextConfidence
- = ContextHigh
- | ContextMedium
- | ContextLow
- deriving (Show, Eq)
-
--- | Sensible defaults for adaptive context.
-defaultAdaptiveContextConfig :: AdaptiveContextConfig
-defaultAdaptiveContextConfig =
- AdaptiveContextConfig
- { accTemporalLimit = 24,
- accSemanticLimit = 5,
- accTemporalBudgetRatio = 0.85,
- accSimilarityThreshold = 0.6,
- accRecencyDecay = 0.995
- }
-
--- | Get adaptive context combining temporal and semantic retrieval.
--- Returns formatted context with explicit window labels and token count.
-getAdaptiveContext :: Text -> Int -> Int -> Text -> IO (Text, Int)
-getAdaptiveContext uid chatId maxTokens currentMessage = do
- let cfg = defaultAdaptiveContextConfig
- temporalBudget = floor (fromIntegral maxTokens * accTemporalBudgetRatio cfg)
- semanticBudget = maxTokens - temporalBudget
- adjustedTemporalBudget = max 0 (temporalBudget - 50)
-
- now <- getCurrentTime
- queryEmbedding <- embedText currentMessage
- (semanticHits, recentSimilarity, recentGapHours) <- case queryEmbedding of
- Left _ -> do
- semanticHits <- searchChatHistoryInChat chatId currentMessage (accSemanticLimit cfg * 2)
- pure (semanticHits, Nothing, Nothing)
- Right qEmb -> do
- (recentSim, recentGap) <- getRecentChatHistorySimilarity chatId qEmb (accTemporalLimit cfg) now
- semanticHits <- searchChatHistoryInChatWithEmbedding chatId qEmb (accSemanticLimit cfg * 2) now
- pure (semanticHits, recentSim, recentGap)
-
- -- 1. Temporal window: recent messages
- recentMsgs <- getRecentMessages uid chatId (accTemporalLimit cfg)
- let recentOldestFirst = reverse recentMsgs
- (selectedRecent, temporalTokens) = selectMessagesByBudget recentOldestFirst adjustedTemporalBudget
-
- -- 2. Semantic window: older relevant messages (not already in temporal)
- let recentContents = map cmContent selectedRecent
- oldButRelevant =
- [ (e, score)
- | (e, score) <- semanticHits,
- cheContent e `notElem` recentContents
+getRecentChatHistorySimilarity :: Int -> VS.Vector Float -> Int -> UTCTime -> IO (Maybe Float, Maybe Float)
+getRecentChatHistorySimilarity chatId qEmb limit now = do
+ entries <-
+ withMemoryDb <| \conn ->
+ SQL.query
+ conn
+ "SELECT id, chat_id, user_id, role, sender_name, content, embedding, created_at \
+ \FROM chat_history WHERE chat_id = ? AND embedding IS NOT NULL ORDER BY created_at DESC LIMIT ?"
+ (chatId, limit)
+ let scores =
+ [ cosineSimilarity qEmb emb
+ | e <- entries,
+ Just emb <- [cheEmbedding e]
]
- (selectedSemantic, semanticTokens) = selectSemanticByBudget oldButRelevant semanticBudget
-
- bestSemantic = listToMaybe semanticHits
- bestSemanticScore = snd </ bestSemantic
- bestSemanticIsRecent = maybe False (\(e, _) -> cheContent e `elem` recentContents) bestSemantic
- recentScore = fromMaybe 0 recentSimilarity
- topicDrift = case bestSemanticScore of
- Just score ->
- not bestSemanticIsRecent
- && score
- >= accSimilarityThreshold cfg
- && score
- - recentScore
- >= 0.1
- Nothing -> False
-
- gapHours = fromMaybe 0 recentGapHours
- baseConfidence
- | gapHours >= 8 = ContextLow
- | recentScore >= 0.75 = ContextHigh
- | recentScore >= 0.6 = ContextMedium
- | otherwise = ContextLow
- confidence = if topicDrift then ContextLow else baseConfidence
-
- contextWarning = case confidence of
- ContextLow ->
- "## Context Confidence\nLOW: recent context may be insufficient. Use search_chat_history or get_messages_by_time to backfill before responding.\n\n"
- _ -> ""
-
- -- Format with explicit window labels
- let temporalSection =
- if null selectedRecent
- then ""
- else
- "## Recent conversation (temporal context)\n"
- <> Text.unlines (map formatConversationMsg selectedRecent)
- <> "\n"
-
- semanticSection =
- if null selectedSemantic
- then ""
- else
- "## Related past messages (semantic context)\n"
- <> Text.unlines (map formatSemanticHit selectedSemantic)
- <> "\n"
-
- totalTokens = temporalTokens + semanticTokens
-
- pure (contextWarning <> temporalSection <> semanticSection, totalTokens)
- where
- selectMessagesByBudget :: [ConversationMessage] -> Int -> ([ConversationMessage], Int)
- selectMessagesByBudget 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))
-
- selectSemanticByBudget :: [(ChatHistoryEntry, Float)] -> Int -> ([(ChatHistoryEntry, Float)], Int)
- selectSemanticByBudget hits budget = go hits budget []
- where
- tokenCount (e, _) = estimateTokens (cheContent e)
- go [] _ acc = (reverse acc, sum (map tokenCount acc))
- go ((e, s) : rest) remaining acc
- | tokens <= remaining = go rest (remaining - tokens) ((e, s) : acc)
- | otherwise = (reverse acc, sum (map tokenCount acc))
- where
- tokens = estimateTokens (cheContent e)
-
- getRecentChatHistorySimilarity :: Int -> VS.Vector Float -> Int -> UTCTime -> IO (Maybe Float, Maybe Float)
- getRecentChatHistorySimilarity targetChatId qEmb limit now = do
- entries <-
- withMemoryDb <| \conn ->
- SQL.query
- conn
- "SELECT id, chat_id, user_id, role, sender_name, content, embedding, created_at \
- \FROM chat_history WHERE chat_id = ? AND embedding IS NOT NULL ORDER BY created_at DESC LIMIT ?"
- (targetChatId, limit)
- let scores =
- [ cosineSimilarity qEmb emb
- | e <- entries,
- Just emb <- [cheEmbedding e]
- ]
- bestScore = if null scores then Nothing else Just (List.maximum scores)
- gapHours = case entries of
- (e : _) -> Just (realToFrac (diffUTCTime now (cheCreatedAt e)) / 3600 :: Float)
- _ -> Nothing
- pure (bestScore, gapHours)
-
- formatConversationMsg m =
- let timestamp = formatEasternTime "%Y-%m-%d %H:%M" (cmCreatedAt m)
- prefix = case cmRole m of
- UserRole -> "[" <> timestamp <> "] " <> fromMaybe "User" (cmSenderName m) <> ": "
- AssistantRole -> "[" <> timestamp <> "] Assistant: "
- in prefix <> cmContent m
-
- formatSemanticHit (e, score) =
- let timestamp = formatEasternTime "%Y-%m-%d %H:%M" (cheCreatedAt e)
- scoreText = Text.pack (printf "%.0f%%" (score * 100 :: Float))
- rolePrefix = if cheRole e == "user" then fromMaybe "User" (cheSenderName e) else "Assistant"
- in "[" <> timestamp <> " | relevance: " <> scoreText <> "] " <> rolePrefix <> ": " <> cheContent e
+ bestScore = if null scores then Nothing else Just (List.maximum scores)
+ gapHours = case entries of
+ (e : _) -> Just (realToFrac (diffUTCTime now (cheCreatedAt e)) / 3600 :: Float)
+ _ -> Nothing
+ pure (bestScore, gapHours)
-- -----------------------------------------------------------------------------
-- Chat History Tools (Active Retrieval)
@@ -2656,6 +2474,14 @@ instance Aeson.FromJSON GetMessagesByTimeArgs where
-- | Tool for time-based message retrieval (active-temporal quadrant).
getMessagesByTimeTool :: Text -> Int -> Engine.Tool
getMessagesByTimeTool uid chatId =
+ messagesByTimeTool (executeMessagesByTime (getMessagesByTimeRange uid chatId))
+
+getGroupMessagesByTimeTool :: Int -> Maybe Int -> Engine.Tool
+getGroupMessagesByTimeTool chatId threadId =
+ messagesByTimeTool (executeMessagesByTime (getGroupMessagesByTimeRange chatId threadId))
+
+messagesByTimeTool :: (Aeson.Value -> IO Aeson.Value) -> Engine.Tool
+messagesByTimeTool execute =
Engine.Tool
{ Engine.toolName = "get_messages_by_time",
Engine.toolDescription =
@@ -2684,18 +2510,17 @@ getMessagesByTimeTool uid chatId =
],
"required" .= (["start_time"] :: [Text])
],
- Engine.toolExecute = executeGetMessagesByTime uid chatId
+ Engine.toolExecute = execute
}
-executeGetMessagesByTime :: Text -> Int -> Aeson.Value -> IO Aeson.Value
-executeGetMessagesByTime uid chatId v =
+executeMessagesByTime :: (UTCTime -> UTCTime -> Int -> IO [ConversationMessage]) -> Aeson.Value -> IO Aeson.Value
+executeMessagesByTime fetch v =
case Aeson.fromJSON v of
Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e])
Aeson.Success (args :: GetMessagesByTimeArgs) -> do
now <- getCurrentTime
let limit' = min 50 (fromMaybe 20 (gmtLimit args))
- -- Parse start time (support relative times)
let startTime = parseRelativeOrAbsoluteTime now (gmtStartTime args)
endTime = case gmtEndTime args of
Nothing -> Just now
@@ -2705,7 +2530,7 @@ executeGetMessagesByTime uid chatId v =
Nothing -> pure (Aeson.object ["error" .= ("Could not parse start_time" :: Text)])
Just start -> do
let end = fromMaybe now endTime
- msgs <- getMessagesByTimeRange uid chatId start end limit'
+ msgs <- fetch start end limit'
pure
( Aeson.object
[ "success" .= True,
@@ -2776,6 +2601,27 @@ getMessagesByTimeRange uid chatId startTime endTime limit' =
\ORDER BY created_at ASC LIMIT ?"
(uid, chatId, startTime, endTime, limit')
+getGroupMessagesByTimeRange :: Int -> Maybe Int -> UTCTime -> UTCTime -> Int -> IO [ConversationMessage]
+getGroupMessagesByTimeRange chatId mThreadId startTime endTime 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 = ? AND created_at >= ? AND created_at <= ? \
+ \ORDER BY created_at ASC LIMIT ?"
+ (chatId, threadId, startTime, endTime, 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 AND created_at >= ? AND created_at <= ? \
+ \ORDER BY created_at ASC LIMIT ?"
+ (chatId, startTime, endTime, limit')
+
-- | Save an agent interaction with embedding for semantic search.
saveAgentInteraction ::
Maybe Text -> -- User ID
diff --git a/Omni/Agent/Prompt/Hydrate.hs b/Omni/Agent/Prompt/Hydrate.hs
index 982c6836..22c99f00 100644
--- a/Omni/Agent/Prompt/Hydrate.hs
+++ b/Omni/Agent/Prompt/Hydrate.hs
@@ -104,7 +104,8 @@ data SourceParams = SourceParams
-- | Result from temporal source (recent messages).
data TemporalResult = TemporalResult
{ trMessages :: [(Time.UTCTime, Text, Text, Text)], -- (time, role, sender, content)
- trTotalTokens :: Int
+ trTotalTokens :: Int,
+ trContextWarning :: Maybe Text
}
deriving (Show, Eq, Generic)
@@ -355,10 +356,11 @@ buildStateSection secId' label content now =
-- | Build the temporal context section from recent messages.
buildTemporalSection :: TemporalResult -> Time.UTCTime -> Section
buildTemporalSection result now =
- let formatted = formatTemporalMessages (trMessages result)
+ let warningText = maybe "" (<> "\n") (trContextWarning result)
+ formatted = warningText <> formatTemporalMessages (trMessages result)
in Section
- { secId = "temporal",
- secLabel = "## Recent conversation (temporal context)",
+ { secId = "passive_temporal",
+ secLabel = "## Passive temporal context (recent messages)",
secSource = SourceTemporal,
secContent = formatted,
secTokens = trTotalTokens result,
@@ -384,8 +386,8 @@ buildSemanticSection result _now =
[] -> Nothing
matches -> Just (minimum (map (\(t, _, _, _) -> t) matches))
in Section
- { secId = "semantic",
- secLabel = "## Related past messages (semantic context)",
+ { secId = "passive_semantic_messages",
+ secLabel = "## Passive semantic context (related messages)",
secSource = SourceSemantic avgScore,
secContent = formatted,
secTokens = srTotalTokens result,
@@ -417,8 +419,8 @@ buildKnowledgeSection result _now =
else "\n(Note: " <> tshow (krSupersededCount result) <> " outdated memories were filtered out)"
fullContent = formatted <> contradictionWarning <> supersededNote
in Section
- { secId = "knowledge",
- secLabel = "## Long-term knowledge",
+ { secId = "passive_semantic_knowledge",
+ secLabel = "## Passive semantic context (long-term memories)",
secSource = SourceKnowledge,
secContent = fullContent,
secTokens = krTotalTokens result,
@@ -528,7 +530,8 @@ test =
let result =
TemporalResult
{ trMessages = [(now, "user", "Ben", "Hello")],
- trTotalTokens = 10
+ trTotalTokens = 10,
+ trContextWarning = Nothing
}
sec = buildTemporalSection result now
secCompositionMode sec Test.@=? Contextual
diff --git a/Omni/Agent/Prompt/MemorySources.hs b/Omni/Agent/Prompt/MemorySources.hs
index 61378516..23737617 100644
--- a/Omni/Agent/Prompt/MemorySources.hs
+++ b/Omni/Agent/Prompt/MemorySources.hs
@@ -1,474 +1,6 @@
-{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE NoImplicitPrelude #-}
--- | Context sources backed by the Memory system.
---
--- This module bridges the Memory.hs functions with the Hydrate.hs
--- ContextSource abstraction, enabling the Prompt IR pipeline to
--- pull context from the memory database.
---
--- : out memory-sources-test
--- : dep aeson
--- : dep directory
--- : dep filepath
--- : dep sqlite-simple
--- : dep time
--- : dep vector
-module Omni.Agent.Prompt.MemorySources
- ( -- * Context Sources
- memoryTemporalSource,
- memorySemanticSource,
- memoryKnowledgeSource,
+-- | Backwards-compatible re-exports for context hydration.
+module Omni.Agent.Prompt.MemorySources (module Omni.Agent.Context) where
- -- * HydrationConfig Builder
- buildHydrationConfig,
-
- -- * Section Builders
- mkProjectSection,
- mkTimeSection,
-
- -- * Tool Conversion
- engineToolToToolDef,
-
- -- * Helpers
- estimateTokensSimple,
-
- -- * Testing
- main,
- test,
- )
-where
-
-import Alpha
-import qualified Data.Aeson as Aeson
-import qualified Data.Text as Text
-import qualified Data.Time as Time
-import qualified Database.SQLite.Simple as SQL
-import qualified Omni.Agent.Engine as Engine
-import qualified Omni.Agent.Memory as Memory
-import Omni.Agent.Prompt.Hydrate
- ( ContextSource (..),
- HydrationConfig (..),
- KnowledgeResult (..),
- SemanticResult (..),
- SourceParams (..),
- TemporalResult (..),
- defaultHydrationConfig,
- )
-import Omni.Agent.Prompt.IR
- ( CompositionMode (..),
- Priority (..),
- Section (..),
- SectionSource (..),
- ToolDef (..),
- )
-import qualified Omni.Test as Test
-import qualified System.Directory as Dir
-import qualified System.Environment as Env
-import System.FilePath ((</>))
-
--- | Create a temporal context source from Memory.
---
--- Uses getRecentMessages to fetch recent conversation history.
-memoryTemporalSource ::
- -- | Memory identity
- Memory.MemoryIdentity ->
- ContextSource TemporalResult
-memoryTemporalSource memId =
- ContextSource <| \params -> do
- -- Parse chatId from context (format: "telegram:12345" -> 12345)
- let chatId = parseChatIdFromContext (Memory.miContext memId)
- -- Fetch recent messages up to maxItems
- messages <- Memory.getRecentMessages (Memory.miOwner memId) chatId (spMaxItems params)
-
- -- Convert to the expected format: (time, role, sender, content)
- let formatted =
- [ ( Memory.cmCreatedAt msg,
- roleToText (Memory.cmRole msg),
- fromMaybe "unknown" (Memory.cmSenderName msg),
- Memory.cmContent msg
- )
- | msg <- reverse messages -- oldest first for context
- ]
-
- -- Estimate total tokens
- totalTokens = sum [Memory.cmTokensEstimate msg | msg <- messages]
-
- pure
- TemporalResult
- { trMessages = formatted,
- trTotalTokens = totalTokens
- }
-
--- | Create a semantic context source from Memory.
---
--- Uses searchChatHistoryInChat to find similar past messages within the same chat.
--- This prevents DM conversations from leaking into group chat responses.
-memorySemanticSource ::
- -- | Memory identity (used to scope search to current chat)
- Memory.MemoryIdentity ->
- ContextSource SemanticResult
-memorySemanticSource memId =
- ContextSource <| \params -> do
- -- Parse chatId from context (format: "telegram:12345" -> 12345)
- let chatId = parseChatIdFromContext (Memory.miContext memId)
- -- Search for semantically similar messages within this chat only
- results <- Memory.searchChatHistoryInChat chatId (spObservation params) (spMaxItems params)
-
- -- Filter by similarity threshold and convert format
- let filtered =
- [ (Memory.cheCreatedAt entry, Memory.cheRole entry, Memory.cheContent entry, score)
- | (entry, score) <- results,
- score >= spThreshold params
- ]
-
- -- Estimate tokens (rough: 4 chars per token)
- totalTokens = sum [estimateTokensSimple (Memory.cheContent entry) | (entry, _) <- results]
-
- pure
- SemanticResult
- { srMatches = filtered,
- srTotalTokens = totalTokens
- }
-
--- | Create a knowledge context source from Memory.
---
--- Uses recallMemoriesEnhanced to fetch relevant long-term memories with:
--- - Recency weighting (recent memories get a boost)
--- - Superseded memory filtering (old info that's been replaced is hidden)
--- - Contradiction detection (warns when memories conflict)
-memoryKnowledgeSource ::
- -- | Memory identity
- Memory.MemoryIdentity ->
- ContextSource KnowledgeResult
-memoryKnowledgeSource memId =
- ContextSource <| \params -> do
- -- Recall memories with enhanced graph-aware retrieval
- result <- Memory.recallMemoriesEnhanced (Memory.miOwner memId) (spObservation params) (spMaxItems params)
-
- let memories = map Memory.smMemory (Memory.errMemories result)
-
- -- Convert to expected format: (content, context, tags)
- let formatted =
- [ ( Memory.memoryContent mem,
- sourceToContext (Memory.memorySource mem),
- Memory.memoryTags mem
- )
- | mem <- memories
- ]
-
- -- Convert contradictions to content pairs
- contradictions =
- [ (Memory.memoryContent m1, Memory.memoryContent m2)
- | (m1, m2) <- Memory.errContradictions result
- ]
-
- -- Estimate tokens
- totalTokens = sum [estimateTokensSimple (Memory.memoryContent mem) | mem <- memories]
-
- pure
- KnowledgeResult
- { krMemories = formatted,
- krContradictions = contradictions,
- krSupersededCount = length (Memory.errSuperseded result),
- krTotalTokens = totalTokens
- }
-
--- | Simple token estimation (roughly 4 characters per token).
-estimateTokensSimple :: Text -> Int
-estimateTokensSimple t = max 1 (Text.length t `div` 4)
-
--- | Convert MessageRole to Text.
-roleToText :: Memory.MessageRole -> Text
-roleToText Memory.UserRole = "user"
-roleToText Memory.AssistantRole = "assistant"
-
--- | Convert MemorySource to a context description.
-sourceToContext :: Memory.MemorySource -> Text
-sourceToContext = Memory.sourceContext
-
--- | Parse a chat ID from a context string.
---
--- Context formats:
--- "telegram:12345" -> 12345
--- "cli:session-abc" -> 0 (no chat ID)
--- "12345" -> 12345 (legacy format)
-parseChatIdFromContext :: Text -> Int
-parseChatIdFromContext ctx
- | "telegram:" `Text.isPrefixOf` ctx =
- fromMaybe 0 (readMaybe (Text.unpack (Text.drop 9 ctx)))
- | otherwise =
- fromMaybe 0 (readMaybe (Text.unpack ctx))
-
--- | Build a HydrationConfig for Telegram context.
---
--- This creates a fully configured HydrationConfig with Memory-backed
--- context sources. Use this to enable the IR-based context pipeline.
---
--- Example:
---
--- > cfg <- buildHydrationConfig
--- > systemPrompt
--- > tools
--- > [projectSection, timeSection]
--- > (MemoryIdentity userId ("telegram:" <> show chatId))
--- > let seqConfig = defaultSeqConfig provider seqTools
--- > { seqHydrationConfig = Just cfg }
-buildHydrationConfig ::
- -- | Base system prompt
- Text ->
- -- | Available tools (Engine.Tool format)
- [Engine.Tool] ->
- -- | Static sections (project context, time, etc.)
- [Section] ->
- -- | Identity for memory queries
- Memory.MemoryIdentity ->
- HydrationConfig
-buildHydrationConfig systemPrompt tools staticSections memIdentity =
- defaultHydrationConfig
- { hcSystemPrompt = systemPrompt,
- hcTools = map engineToolToToolDef tools,
- hcStaticSections = staticSections,
- hcTemporalSource = Just (memoryTemporalSource memIdentity),
- hcSemanticSource = Just (memorySemanticSource memIdentity),
- hcKnowledgeSource = Just (memoryKnowledgeSource memIdentity),
- hcEstimateTokens = estimateTokensSimple
- }
-
--- | Convert an Engine.Tool to IR.ToolDef.
-engineToolToToolDef :: Engine.Tool -> ToolDef
-engineToolToToolDef tool =
- ToolDef
- { tdName = Engine.toolName tool,
- tdDescription = Engine.toolDescription tool,
- tdSchema = Engine.toolJsonSchema tool,
- tdPriority = High, -- Default to High priority
- tdEmbedding = Nothing -- No embedding by default
- }
-
--- | Create a static section for project context.
-mkProjectSection :: Text -> Text -> Section
-mkProjectSection projectName workingDir =
- Section
- { secId = "project",
- secLabel = "Current Project",
- secSource = SourceState "project",
- secContent =
- Text.unlines
- [ "Project: " <> projectName,
- "Working directory: " <> workingDir
- ],
- secTokens = 20,
- secMinTokens = Nothing,
- secPriority = High,
- secRelevance = Nothing,
- secRecency = Nothing,
- secCompositionMode = Constraint,
- secEmbedding = Nothing,
- secHash = Nothing
- }
-
--- | Create a static section for current time.
-mkTimeSection :: Time.UTCTime -> Time.TimeZone -> Section
-mkTimeSection utcTime tz =
- let localTime = Time.utcToLocalTime tz utcTime
- timeStr = Text.pack (Time.formatTime Time.defaultTimeLocale "%A, %B %d, %Y at %H:%M" localTime)
- in Section
- { secId = "time",
- secLabel = "Current Date and Time",
- secSource = SourceState "clock",
- secContent = timeStr,
- secTokens = 15,
- secMinTokens = Nothing,
- secPriority = High,
- secRelevance = Nothing,
- secRecency = Just utcTime,
- secCompositionMode = Contextual,
- secEmbedding = Nothing,
- secHash = Nothing
- }
-
---------------------------------------------------------------------------------
--- Tests
---------------------------------------------------------------------------------
-
-main :: IO ()
-main = Test.run test
-
-test :: Test.Tree
-test =
- Test.group
- "Omni.Agent.Prompt.MemorySources"
- [ Test.unit "estimateTokensSimple basic" <| do
- estimateTokensSimple "" Test.@=? 1
- estimateTokensSimple "hello" Test.@=? 1
- estimateTokensSimple "hello world test" Test.@=? 4
- estimateTokensSimple (Text.replicate 100 "a") Test.@=? 25,
- Test.unit "roleToText converts correctly" <| do
- roleToText Memory.UserRole Test.@=? "user"
- roleToText Memory.AssistantRole Test.@=? "assistant",
- Test.unit "sourceToContext extracts context" <| do
- let src =
- Memory.MemorySource
- { Memory.sourceAgent = "test",
- Memory.sourceSession = Nothing,
- Memory.sourceContext = "from test"
- }
- sourceToContext src Test.@=? "from test",
- Test.unit "mkProjectSection creates valid section" <| do
- let section = mkProjectSection "omni" "/home/ben/omni"
- secId section Test.@=? "project"
- secPriority section Test.@=? High
- Text.isInfixOf "omni" (secContent section) Test.@=? True,
- Test.unit "mkTimeSection creates valid section" <| do
- now <- Time.getCurrentTime
- tz <- Time.getCurrentTimeZone
- let section = mkTimeSection now tz
- secId section Test.@=? "time"
- secPriority section Test.@=? High
- secRecency section Test.@=? Just now,
- Test.unit "engineToolToToolDef converts correctly" <| do
- let tool =
- Engine.Tool
- { Engine.toolName = "test_tool",
- Engine.toolDescription = "A test tool",
- Engine.toolJsonSchema = Aeson.object [],
- Engine.toolExecute = \_ -> pure (Aeson.String "ok")
- }
- toolDef = engineToolToToolDef tool
- tdName toolDef Test.@=? "test_tool"
- tdDescription toolDef Test.@=? "A test tool"
- tdPriority toolDef Test.@=? High,
- -- Database-backed tests
- Test.unit "memoryTemporalSource returns messages from db" <| do
- withTestDb <| \userId chatId -> do
- let memId = Memory.MemoryIdentity userId (tshow chatId)
- source = memoryTemporalSource memId
- params =
- SourceParams
- { spObservation = "test query",
- spMaxTokens = 1000,
- spMaxItems = 10,
- spThreshold = 0.5
- }
- result <- runSource source params
- -- Should have our test messages
- length (trMessages result) Test.@=? 3
- -- Messages should be in chronological order (oldest first)
- let contents = [content | (_, _, _, content) <- trMessages result]
- case contents of
- (msg1 : _ : msg3 : _) -> do
- msg1 Test.@=? "First message"
- msg3 Test.@=? "Third message"
- _ -> Test.assertFailure "Expected 3 messages",
- Test.unit "memoryKnowledgeSource runs without error" <| do
- withTestDb <| \userId chatId -> do
- let memId = Memory.MemoryIdentity userId (tshow chatId)
- source = memoryKnowledgeSource memId
- params =
- SourceParams
- { spObservation = "programming",
- spMaxTokens = 1000,
- spMaxItems = 10,
- spThreshold = 0.5
- }
- -- Just verify it runs without throwing
- -- (may return 0 memories since test data has no embeddings and
- -- embedding service may not be running)
- result <- runSource source params
- krTotalTokens result Test.@=? krTotalTokens result -- always true, just checking it runs
- ]
-
--- | Run tests with a temporary database populated with test data.
-withTestDb :: (Text -> Int -> IO a) -> IO a
-withTestDb action = do
- -- Create temp directory for test db
- tmpDir <- Dir.getTemporaryDirectory
- let testDbPath = tmpDir </> "memory-sources-test.db"
-
- -- Remove old test db if exists
- dbExists <- Dir.doesFileExist testDbPath
- when dbExists <| Dir.removeFile testDbPath
-
- -- Set environment variable to use test db
- Env.setEnv "MEMORY_DB_PATH" testDbPath
-
- -- Initialize and populate database
- SQL.withConnection testDbPath <| \conn -> do
- Memory.initMemoryDb conn
- populateTestData conn
-
- -- Run the test
- result <- action testUserId testChatId
-
- -- Cleanup
- Env.unsetEnv "MEMORY_DB_PATH"
- Dir.removeFile testDbPath
-
- pure result
-
-testUserId :: Text
-testUserId = "test-user-001"
-
-testChatId :: Int
-testChatId = 12345
-
--- | Populate test database with sample data.
-populateTestData :: SQL.Connection -> IO ()
-populateTestData conn = do
- now <- Time.getCurrentTime
- let oneHourAgo = Time.addUTCTime (-3600) now
- twoHoursAgo = Time.addUTCTime (-7200) now
-
- -- Create test user
- SQL.execute
- conn
- "INSERT INTO users (id, name, created_at) VALUES (?, ?, ?)"
- (testUserId, "Test User" :: Text, now)
-
- -- Create test conversation messages
- SQL.execute
- conn
- "INSERT INTO conversation_messages (user_id, chat_id, role, sender_name, content, tokens_estimate, created_at) VALUES (?, ?, ?, ?, ?, ?, ?)"
- (testUserId, testChatId, "user" :: Text, Just ("Test User" :: Text), "First message" :: Text, 10 :: Int, twoHoursAgo)
-
- SQL.execute
- conn
- "INSERT INTO conversation_messages (user_id, chat_id, role, sender_name, content, tokens_estimate, created_at) VALUES (?, ?, ?, ?, ?, ?, ?)"
- (testUserId, testChatId, "assistant" :: Text, Just ("Ava" :: Text), "Second message" :: Text, 10 :: Int, oneHourAgo)
-
- SQL.execute
- conn
- "INSERT INTO conversation_messages (user_id, chat_id, role, sender_name, content, tokens_estimate, created_at) VALUES (?, ?, ?, ?, ?, ?, ?)"
- (testUserId, testChatId, "user" :: Text, Just ("Test User" :: Text), "Third message" :: Text, 10 :: Int, now)
-
- -- Create test memories
- SQL.execute
- conn
- "INSERT INTO memories (id, user_id, content, source_agent, source_session, source_context, confidence, created_at, last_accessed_at, tags) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"
- ( "mem-001" :: Text,
- testUserId,
- "User likes programming in Haskell" :: Text,
- "test" :: Text,
- Nothing :: Maybe Text,
- "test context" :: Text,
- 0.9 :: Double,
- now,
- now,
- "[\"programming\",\"haskell\"]" :: Text
- )
-
- SQL.execute
- conn
- "INSERT INTO memories (id, user_id, content, source_agent, source_session, source_context, confidence, created_at, last_accessed_at, tags) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"
- ( "mem-002" :: Text,
- testUserId,
- "User works on the Omni project" :: Text,
- "test" :: Text,
- Nothing :: Maybe Text,
- "test context" :: Text,
- 0.8 :: Double,
- now,
- now,
- "[\"work\",\"omni\"]" :: Text
- )
+import Omni.Agent.Context
diff --git a/Omni/Ava/Telegram/Bot.hs b/Omni/Ava/Telegram/Bot.hs
index 5e35b7d4..32da99d6 100644
--- a/Omni/Ava/Telegram/Bot.hs
+++ b/Omni/Ava/Telegram/Bot.hs
@@ -97,11 +97,11 @@ import qualified Network.HTTP.Client as HTTPClient
import qualified Network.HTTP.Simple as HTTP
import qualified Omni.Agent.AuditLog as AuditLog
import qualified Omni.Agent.Auth as Auth
+import qualified Omni.Agent.Context as Context
import qualified Omni.Agent.Engine as Engine
import qualified Omni.Agent.Memory as Memory
import qualified Omni.Agent.Op.Bridge as OpBridge
import qualified Omni.Agent.Paths as Paths
-import qualified Omni.Agent.Prompt.MemorySources as MemorySources
import qualified Omni.Agent.Prompts.Core as Prompts
import qualified Omni.Agent.Provider as Provider
import qualified Omni.Agent.Tools as Tools
@@ -945,16 +945,6 @@ runHeartbeatTurn tgConfig modelVar projectVar providerVar engineCfg chatId = do
<> "\n\n## heartbeat.md\n"
<> heartbeatMd
- personalMemories <- Memory.recallMemories heartbeatUid heartbeatUserMessage 5
- let memoryContext =
- if null personalMemories
- then "No memories found."
- else Text.unlines <| ["[Personal] " <> Memory.memoryContent m | m <- personalMemories]
-
- (conversationContext, contextTokens) <-
- Memory.getAdaptiveContext heartbeatUid chatId maxConversationTokens heartbeatUserMessage
- putText <| "Conversation context: " <> tshow contextTokens <> " tokens (adaptive: temporal + semantic)"
-
-- Use DST-aware Eastern timezone
let localTime = Time.utcToEastern now
timeStr = Text.pack (formatTime defaultTimeLocale "%A, %B %d, %Y at %H:%M" localTime)
@@ -1061,10 +1051,6 @@ runHeartbeatTurn tgConfig modelVar projectVar providerVar engineCfg chatId = do
<> "\n\n## Current User\n"
<> "You are talking to: "
<> userName
- <> "\n\n## What you know about this user\n"
- <> memoryContext
- <> "\n\n"
- <> conversationContext
-- Get provider with failover chain
currentProviderType <- readTVarIO providerVar
@@ -1081,7 +1067,8 @@ runHeartbeatTurn tgConfig modelVar projectVar providerVar engineCfg chatId = do
pure <| Provider.ollamaProvider "qwen2.5:14b-instruct-q4_K_M"
-- Build tools (same as normal messages, minus image tools which need photo context)
- let coreTools =
+ let conversationScope = Context.DirectConversation heartbeatUid chatId
+ coreTools =
if isBenAuthorized userName
then
[ Tools.skillTool,
@@ -1090,11 +1077,9 @@ runHeartbeatTurn tgConfig modelVar projectVar providerVar engineCfg chatId = do
]
else []
memoryTools =
- [ Memory.rememberTool heartbeatUid,
- Memory.recallTool heartbeatUid,
- Memory.searchChatHistoryTool chatId,
- Memory.getMessagesByTimeTool heartbeatUid chatId
- ]
+ Context.memoryWriteTools conversationScope
+ <> Context.activeSemanticTools conversationScope
+ <> Context.activeTemporalTools conversationScope
searchTools = case Types.tgKagiApiKey tgConfig of
Just kagiKey -> [WebSearch.webSearchTool kagiKey]
Nothing -> []
@@ -1104,18 +1089,13 @@ runHeartbeatTurn tgConfig modelVar projectVar providerVar engineCfg chatId = do
-- Build hydration config for memory-augmented context
-- Use "ben" as owner so heartbeat can access the same memories as normal conversations
let easternTZ = Time.getEasternTimeZone now -- DST-aware Eastern timezone
- staticSections = [MemorySources.mkTimeSection now easternTZ]
- memIdentity =
- Memory.MemoryIdentity
- { Memory.miOwner = "ben", -- Same as normal messages so memories are shared
- Memory.miContext = "telegram:" <> tshow chatId
- }
+ staticSections = [Context.mkTimeSection now easternTZ]
hydrationCfg =
- MemorySources.buildHydrationConfig
+ Context.buildContextHydrationConfig
systemPrompt
tools
staticSections
- memIdentity
+ conversationScope
let agentCfg =
Engine.defaultAgentConfig
@@ -1458,16 +1438,9 @@ handleAuthorizedMessageBatchContinued tgConfig modelVar projectVar providerVar e
if not shouldEngage
then putText "Skipping group message (pre-filter said no)"
else do
- (conversationContext, contextTokens) <-
- if isGroup
- then do
- _ <- Memory.saveGroupMessage chatId threadId Memory.UserRole userName userMessage
- Memory.getGroupConversationContext chatId threadId maxConversationTokens
- else do
- _ <- Memory.saveMessage uid chatId Memory.UserRole (Just userName) userMessage
- -- Use adaptive context: combines temporal (recent) + semantic (relevant older) messages
- Memory.getAdaptiveContext (Memory.unUserId uid) chatId maxConversationTokens userMessage
- putText <| "Conversation context: " <> tshow contextTokens <> " tokens (adaptive: temporal + semantic)"
+ if isGroup
+ then void <| Memory.saveGroupMessage chatId threadId Memory.UserRole userName userMessage
+ else void <| Memory.saveMessage uid chatId Memory.UserRole (Just userName) userMessage
now <- getCurrentTime
_ <- forkIO <| Memory.saveChatHistoryEntry chatId (Just (Memory.unUserId uid)) "user" (Just userName) userMessage now
@@ -1484,7 +1457,7 @@ handleAuthorizedMessageBatchContinued tgConfig modelVar projectVar providerVar e
AuditLog.emptyMetadata
AuditLog.writeAvaLog entry
- processEngagedMessage tgConfig modelVar projectVar providerVar engineCfg msg uid userName chatId userMessage conversationContext photoBytes
+ processEngagedMessage tgConfig modelVar projectVar providerVar engineCfg msg uid userName chatId userMessage photoBytes
processEngagedMessage ::
Types.TelegramConfig ->
@@ -1497,10 +1470,9 @@ processEngagedMessage ::
Text ->
Int ->
Text ->
- Text ->
Maybe BL.ByteString -> -- photo bytes for image editing tool
IO ()
-processEngagedMessage tgConfig modelVar projectVar providerVar engineCfg msg uid userName chatId userMessage conversationContext photoBytes = do
+processEngagedMessage tgConfig modelVar projectVar providerVar engineCfg msg uid userName chatId userMessage photoBytes = do
let isGroup = Types.isGroupChat msg
lastToolArgsRef <- newIORef (Map.empty :: Map.Map Text Text)
@@ -1552,21 +1524,6 @@ processEngagedMessage tgConfig modelVar projectVar providerVar engineCfg msg uid
{ Engine.engineOnEvent = handleEvent
}
- personalMemories <- Memory.recallMemories (Memory.unUserId uid) userMessage 5
- groupMemories <-
- if isGroup
- then Memory.recallGroupMemories chatId userMessage 3
- else pure []
-
- let allMemories = personalMemories <> groupMemories
- memoryContext =
- if null allMemories
- then "No memories found."
- else
- Text.unlines
- <| ["[Personal] " <> Memory.memoryContent m | m <- personalMemories]
- <> ["[Group] " <> Memory.memoryContent m | m <- groupMemories]
-
now <- getCurrentTime
-- Use DST-aware Eastern timezone
let localTime = Time.utcToEastern now
@@ -1682,10 +1639,6 @@ processEngagedMessage tgConfig modelVar projectVar providerVar engineCfg msg uid
<> "\n\n## Current User\n"
<> "You are talking to: "
<> userName
- <> "\n\n## What you know about this user\n"
- <> memoryContext
- <> "\n\n"
- <> conversationContext
-- MINIMAL TOOLS - everything else via skills + run_bash
let -- Core tools (can't be replaced by skills)
@@ -1697,14 +1650,15 @@ processEngagedMessage tgConfig modelVar projectVar providerVar engineCfg msg uid
Tools.runBashTool -- execute commands (enables everything)
]
else []
- -- Memory tools (persistent context across conversations)
+ conversationScope =
+ if isGroup
+ then Context.GroupConversation chatId (Types.tmThreadId msg)
+ else Context.DirectConversation (Memory.unUserId uid) chatId
+ -- Memory tools (active semantic/temporal + memory write)
memoryTools =
- [ Memory.rememberTool (Memory.unUserId uid),
- Memory.recallTool (Memory.unUserId uid),
- -- Chat history tools (active retrieval for when passive context isn't enough)
- Memory.searchChatHistoryTool chatId,
- Memory.getMessagesByTimeTool (Memory.unUserId uid) chatId
- ]
+ Context.memoryWriteTools conversationScope
+ <> Context.activeSemanticTools conversationScope
+ <> Context.activeTemporalTools conversationScope
-- Web tools (need API keys)
searchTools = case Types.tgKagiApiKey tgConfig of
Just kagiKey -> [WebSearch.webSearchTool kagiKey]
@@ -1756,20 +1710,15 @@ processEngagedMessage tgConfig modelVar projectVar providerVar engineCfg msg uid
-- This adds dynamic context from memory sources alongside the static system prompt
let easternTZ = Time.getEasternTimeZone now
staticSections =
- [ MemorySources.mkTimeSection now easternTZ,
- MemorySources.mkProjectSection (projectName currentProject) (Text.pack (projectPath currentProject))
+ [ Context.mkTimeSection now easternTZ,
+ Context.mkProjectSection (projectName currentProject) (Text.pack (projectPath currentProject))
]
- memIdentity =
- Memory.MemoryIdentity
- { Memory.miOwner = Memory.unUserId uid,
- Memory.miContext = "telegram:" <> tshow chatId
- }
hydrationCfg =
- MemorySources.buildHydrationConfig
+ Context.buildContextHydrationConfig
systemPrompt
tools
staticSections
- memIdentity
+ conversationScope
result <-
withTypingIndicator tgConfig chatId