test_memory_settings.py
1 """ 2 Tests for memory enable/disable settings and memory scoping by user+agent. 3 4 These tests verify: 5 1. The behavior of ConversationAdapters.is_memory_enabled() for different combinations of: 6 - ServerChatSettings.memory_mode (DISABLED, ENABLED_DEFAULT_OFF, ENABLED_DEFAULT_ON) 7 - UserConversationConfig.enable_memory (True, False, or not set) 8 9 2. Memory scoping by user and agent: 10 - Memories are scoped to user + agent 11 - Default agent has access to ALL memories across all agents for a user 12 - Non-default agents only see their own memories 13 """ 14 15 import pytest 16 from unittest.mock import MagicMock 17 18 from khoj.database.adapters import ConversationAdapters, UserMemoryAdapters 19 from khoj.database.models import ServerChatSettings, UserConversationConfig 20 from khoj.routers.helpers import get_user_config 21 from tests.helpers import ( 22 acreate_user, 23 acreate_subscription, 24 acreate_chat_model, 25 acreate_default_agent, 26 acreate_agent, 27 acreate_test_memory, 28 ServerChatSettingsFactory, 29 SubscriptionFactory, 30 UserFactory, 31 ) 32 33 34 # ---------------------------------------------------------------------------------------------------- 35 # Test is_memory_enabled with no server config (default behavior) 36 # ---------------------------------------------------------------------------------------------------- 37 @pytest.mark.django_db 38 def test_memory_enabled_no_server_config_no_user_config(): 39 """When no server config and no user config exists, memory should be enabled (default on).""" 40 user = UserFactory() 41 SubscriptionFactory(user=user) 42 43 result = ConversationAdapters.is_memory_enabled(user) 44 45 assert result is True 46 47 48 @pytest.mark.django_db 49 def test_memory_enabled_no_server_config_user_enabled(): 50 """When no server config but user has explicitly enabled memory.""" 51 user = UserFactory() 52 SubscriptionFactory(user=user) 53 user_config = UserConversationConfig.objects.create(user=user, enable_memory=True) 54 55 result = ConversationAdapters.is_memory_enabled(user) 56 57 assert result is True 58 59 60 @pytest.mark.django_db 61 def test_memory_enabled_no_server_config_user_disabled(): 62 """When no server config but user has explicitly disabled memory.""" 63 user = UserFactory() 64 SubscriptionFactory(user=user) 65 user_config = UserConversationConfig.objects.create(user=user, enable_memory=False) 66 67 result = ConversationAdapters.is_memory_enabled(user) 68 69 assert result is False 70 71 72 # ---------------------------------------------------------------------------------------------------- 73 # Test is_memory_enabled with server mode DISABLED 74 # ---------------------------------------------------------------------------------------------------- 75 @pytest.mark.django_db 76 def test_memory_disabled_server_disabled_no_user_config(): 77 """When server disables memory, it should override everything - no user config.""" 78 user = UserFactory() 79 SubscriptionFactory(user=user) 80 ServerChatSettingsFactory(memory_mode=ServerChatSettings.MemoryMode.DISABLED) 81 82 result = ConversationAdapters.is_memory_enabled(user) 83 84 assert result is False 85 86 87 @pytest.mark.django_db 88 def test_memory_disabled_server_disabled_user_enabled(): 89 """When server disables memory, it should override user preference (enabled).""" 90 user = UserFactory() 91 SubscriptionFactory(user=user) 92 ServerChatSettingsFactory(memory_mode=ServerChatSettings.MemoryMode.DISABLED) 93 UserConversationConfig.objects.create(user=user, enable_memory=True) 94 95 result = ConversationAdapters.is_memory_enabled(user) 96 97 assert result is False 98 99 100 @pytest.mark.django_db 101 def test_memory_disabled_server_disabled_user_disabled(): 102 """When server disables memory, user disabled too - should be disabled.""" 103 user = UserFactory() 104 SubscriptionFactory(user=user) 105 ServerChatSettingsFactory(memory_mode=ServerChatSettings.MemoryMode.DISABLED) 106 UserConversationConfig.objects.create(user=user, enable_memory=False) 107 108 result = ConversationAdapters.is_memory_enabled(user) 109 110 assert result is False 111 112 113 # ---------------------------------------------------------------------------------------------------- 114 # Test is_memory_enabled with server mode ENABLED_DEFAULT_OFF 115 # ---------------------------------------------------------------------------------------------------- 116 @pytest.mark.django_db 117 def test_memory_enabled_default_off_no_user_config(): 118 """When server is enabled_default_off and no user config, memory should be off.""" 119 user = UserFactory() 120 SubscriptionFactory(user=user) 121 ServerChatSettingsFactory(memory_mode=ServerChatSettings.MemoryMode.ENABLED_DEFAULT_OFF) 122 123 result = ConversationAdapters.is_memory_enabled(user) 124 125 assert result is False 126 127 128 @pytest.mark.django_db 129 def test_memory_enabled_default_off_user_enabled(): 130 """When server is enabled_default_off and user opts in, memory should be on.""" 131 user = UserFactory() 132 SubscriptionFactory(user=user) 133 ServerChatSettingsFactory(memory_mode=ServerChatSettings.MemoryMode.ENABLED_DEFAULT_OFF) 134 UserConversationConfig.objects.create(user=user, enable_memory=True) 135 136 result = ConversationAdapters.is_memory_enabled(user) 137 138 assert result is True 139 140 141 @pytest.mark.django_db 142 def test_memory_enabled_default_off_user_disabled(): 143 """When server is enabled_default_off and user explicitly disabled, memory should be off.""" 144 user = UserFactory() 145 SubscriptionFactory(user=user) 146 ServerChatSettingsFactory(memory_mode=ServerChatSettings.MemoryMode.ENABLED_DEFAULT_OFF) 147 UserConversationConfig.objects.create(user=user, enable_memory=False) 148 149 result = ConversationAdapters.is_memory_enabled(user) 150 151 assert result is False 152 153 154 # ---------------------------------------------------------------------------------------------------- 155 # Test is_memory_enabled with server mode ENABLED_DEFAULT_ON 156 # ---------------------------------------------------------------------------------------------------- 157 @pytest.mark.django_db 158 def test_memory_enabled_default_on_no_user_config(): 159 """When server is enabled_default_on and no user config, memory should be on.""" 160 user = UserFactory() 161 SubscriptionFactory(user=user) 162 ServerChatSettingsFactory(memory_mode=ServerChatSettings.MemoryMode.ENABLED_DEFAULT_ON) 163 164 result = ConversationAdapters.is_memory_enabled(user) 165 166 assert result is True 167 168 169 @pytest.mark.django_db 170 def test_memory_enabled_default_on_user_enabled(): 171 """When server is enabled_default_on and user enabled, memory should be on.""" 172 user = UserFactory() 173 SubscriptionFactory(user=user) 174 ServerChatSettingsFactory(memory_mode=ServerChatSettings.MemoryMode.ENABLED_DEFAULT_ON) 175 UserConversationConfig.objects.create(user=user, enable_memory=True) 176 177 result = ConversationAdapters.is_memory_enabled(user) 178 179 assert result is True 180 181 182 @pytest.mark.django_db 183 def test_memory_enabled_default_on_user_disabled(): 184 """When server is enabled_default_on and user opts out, memory should be off.""" 185 user = UserFactory() 186 SubscriptionFactory(user=user) 187 ServerChatSettingsFactory(memory_mode=ServerChatSettings.MemoryMode.ENABLED_DEFAULT_ON) 188 UserConversationConfig.objects.create(user=user, enable_memory=False) 189 190 result = ConversationAdapters.is_memory_enabled(user) 191 192 assert result is False 193 194 195 # ---------------------------------------------------------------------------------------------------- 196 # Test get_user_config returns correct enable_memory and server_memory_mode 197 # ---------------------------------------------------------------------------------------------------- 198 @pytest.mark.django_db 199 def test_get_user_config_memory_no_server_config(): 200 """get_user_config should return default values when no server config.""" 201 user = UserFactory() 202 SubscriptionFactory(user=user) 203 request = MagicMock() 204 request.url = MagicMock() 205 request.url.path = "/api/config" 206 request.session = {} 207 208 config = get_user_config(user, request, is_detailed=True) 209 210 assert config["enable_memory"] is True 211 assert config["server_memory_mode"] == "enabled_default_on" 212 213 214 @pytest.mark.django_db 215 def test_get_user_config_memory_server_disabled(): 216 """get_user_config should reflect server disabled mode.""" 217 user = UserFactory() 218 SubscriptionFactory(user=user) 219 ServerChatSettingsFactory(memory_mode=ServerChatSettings.MemoryMode.DISABLED) 220 request = MagicMock() 221 request.url = MagicMock() 222 request.url.path = "/api/config" 223 request.session = {} 224 225 config = get_user_config(user, request, is_detailed=True) 226 227 assert config["enable_memory"] is False 228 assert config["server_memory_mode"] == "disabled" 229 230 231 @pytest.mark.django_db 232 def test_get_user_config_memory_server_enabled_default_off_user_opted_in(): 233 """get_user_config should show user opted in when server is default off.""" 234 user = UserFactory() 235 SubscriptionFactory(user=user) 236 ServerChatSettingsFactory(memory_mode=ServerChatSettings.MemoryMode.ENABLED_DEFAULT_OFF) 237 UserConversationConfig.objects.create(user=user, enable_memory=True) 238 request = MagicMock() 239 request.url = MagicMock() 240 request.url.path = "/api/config" 241 request.session = {} 242 243 config = get_user_config(user, request, is_detailed=True) 244 245 assert config["enable_memory"] is True 246 assert config["server_memory_mode"] == "enabled_default_off" 247 248 249 @pytest.mark.django_db 250 def test_get_user_config_memory_server_enabled_default_on_user_opted_out(): 251 """get_user_config should show user opted out when server is default on.""" 252 user = UserFactory() 253 SubscriptionFactory(user=user) 254 ServerChatSettingsFactory(memory_mode=ServerChatSettings.MemoryMode.ENABLED_DEFAULT_ON) 255 UserConversationConfig.objects.create(user=user, enable_memory=False) 256 request = MagicMock() 257 request.url = MagicMock() 258 request.url.path = "/api/config" 259 request.session = {} 260 261 config = get_user_config(user, request, is_detailed=True) 262 263 assert config["enable_memory"] is False 264 assert config["server_memory_mode"] == "enabled_default_on" 265 266 267 # ---------------------------------------------------------------------------------------------------- 268 # Test memory scoping by user and agent 269 # ---------------------------------------------------------------------------------------------------- 270 271 272 @pytest.mark.anyio 273 @pytest.mark.django_db(transaction=True) 274 async def test_pull_memories_default_agent_sees_all_memories(): 275 """Default agent should see ALL memories for the user, including those from other agents.""" 276 # Setup 277 user = await acreate_user() 278 await acreate_subscription(user) 279 chat_model = await acreate_chat_model() 280 281 # Create default agent 282 default_agent = await acreate_default_agent() 283 assert default_agent is not None 284 285 # Create a custom agent 286 custom_agent = await acreate_agent("Custom Agent", chat_model, "A custom agent") 287 288 # Create memories for different agents 289 await acreate_test_memory(user, agent=None, raw_text="memory without agent") 290 await acreate_test_memory(user, agent=default_agent, raw_text="memory for default agent") 291 await acreate_test_memory(user, agent=custom_agent, raw_text="memory for custom agent") 292 293 # Act: Pull memories with default agent 294 memories = await UserMemoryAdapters.pull_memories(user=user, agent=default_agent) 295 296 # Assert: Default agent sees ALL memories 297 memory_texts = [m.raw for m in memories] 298 assert "memory without agent" in memory_texts 299 assert "memory for default agent" in memory_texts 300 assert "memory for custom agent" in memory_texts 301 assert len(memories) == 3 302 303 304 @pytest.mark.anyio 305 @pytest.mark.django_db(transaction=True) 306 async def test_pull_memories_custom_agent_sees_only_own_memories(): 307 """Custom (non-default) agent should only see its own memories.""" 308 # Setup 309 user = await acreate_user() 310 await acreate_subscription(user) 311 chat_model = await acreate_chat_model() 312 313 # Create default agent 314 default_agent = await acreate_default_agent() 315 assert default_agent is not None 316 317 # Create custom agents 318 custom_agent_1 = await acreate_agent("Custom Agent 1", chat_model, "First custom agent") 319 custom_agent_2 = await acreate_agent("Custom Agent 2", chat_model, "Second custom agent") 320 321 # Create memories for different agents 322 await acreate_test_memory(user, agent=None, raw_text="memory without agent") 323 await acreate_test_memory(user, agent=default_agent, raw_text="memory for default agent") 324 await acreate_test_memory(user, agent=custom_agent_1, raw_text="memory for custom agent 1") 325 await acreate_test_memory(user, agent=custom_agent_2, raw_text="memory for custom agent 2") 326 327 # Act: Pull memories with custom_agent_1 328 memories = await UserMemoryAdapters.pull_memories(user=user, agent=custom_agent_1) 329 330 # Assert: Custom agent 1 only sees its own memories 331 memory_texts = [m.raw for m in memories] 332 assert "memory for custom agent 1" in memory_texts 333 assert "memory without agent" not in memory_texts 334 assert "memory for default agent" not in memory_texts 335 assert "memory for custom agent 2" not in memory_texts 336 assert len(memories) == 1 337 338 339 @pytest.mark.anyio 340 @pytest.mark.django_db(transaction=True) 341 async def test_pull_memories_no_agent_same_as_default_agent(): 342 """Pulling memories with agent=None should behave same as default agent (see all).""" 343 # Setup 344 user = await acreate_user() 345 await acreate_subscription(user) 346 chat_model = await acreate_chat_model() 347 348 # Create default agent 349 default_agent = await acreate_default_agent() 350 assert default_agent is not None 351 352 # Create a custom agent 353 custom_agent = await acreate_agent("Custom Agent", chat_model, "A custom agent") 354 355 # Create memories 356 await acreate_test_memory(user, agent=None, raw_text="memory without agent") 357 await acreate_test_memory(user, agent=default_agent, raw_text="memory for default agent") 358 await acreate_test_memory(user, agent=custom_agent, raw_text="memory for custom agent") 359 360 # Act: Pull memories with agent=None 361 memories = await UserMemoryAdapters.pull_memories(user=user, agent=None) 362 363 # Assert: Should see all memories (same as default agent behavior) 364 memory_texts = [m.raw for m in memories] 365 assert "memory without agent" in memory_texts 366 assert "memory for default agent" in memory_texts 367 assert "memory for custom agent" in memory_texts 368 assert len(memories) == 3 369 370 371 @pytest.mark.anyio 372 @pytest.mark.django_db(transaction=True) 373 async def test_save_memory_with_custom_agent_scopes_to_agent(): 374 """Memories saved with a custom agent should be scoped to that agent.""" 375 # Setup 376 user = await acreate_user() 377 await acreate_subscription(user) 378 chat_model = await acreate_chat_model() 379 380 # Create default agent 381 default_agent = await acreate_default_agent() 382 assert default_agent is not None 383 384 # Create custom agent 385 custom_agent = await acreate_agent("Custom Agent", chat_model, "A custom agent") 386 387 # Create memory with custom agent (directly in DB to avoid embeddings) 388 memory = await acreate_test_memory(user, agent=custom_agent, raw_text="custom agent memory") 389 390 # Assert: Memory is scoped to the custom agent 391 assert memory.agent == custom_agent 392 assert memory.user == user 393 394 # Verify custom agent can see it 395 custom_memories = await UserMemoryAdapters.pull_memories(user=user, agent=custom_agent) 396 assert len(custom_memories) == 1 397 assert custom_memories[0].raw == "custom agent memory" 398 399 400 @pytest.mark.anyio 401 @pytest.mark.django_db(transaction=True) 402 async def test_save_memory_with_default_agent_has_no_agent_scope(): 403 """Memories saved with default agent should have agent=None (global scope).""" 404 # Setup 405 user = await acreate_user() 406 await acreate_subscription(user) 407 await acreate_chat_model() # Required for default agent creation 408 409 # Create default agent 410 default_agent = await acreate_default_agent() 411 assert default_agent is not None 412 413 # Create memory with default agent (directly in DB) 414 # Based on save_memory logic: if agent == default_agent, agent is not set 415 memory = await acreate_test_memory(user, agent=None, raw_text="default agent memory") 416 417 # Assert: Memory has no agent (global scope) 418 assert memory.agent is None 419 assert memory.user == user 420 421 422 @pytest.mark.anyio 423 @pytest.mark.django_db(transaction=True) 424 async def test_memories_isolated_between_users(): 425 """Memories should be isolated between different users.""" 426 # Setup 427 user1 = await acreate_user() 428 user2 = await acreate_user() 429 await acreate_subscription(user1) 430 await acreate_subscription(user2) 431 432 # Create default agent 433 await acreate_default_agent() 434 435 # Create memories for each user 436 await acreate_test_memory(user1, agent=None, raw_text="user1 memory") 437 await acreate_test_memory(user2, agent=None, raw_text="user2 memory") 438 439 # Act: Pull memories for each user 440 user1_memories = await UserMemoryAdapters.pull_memories(user=user1) 441 user2_memories = await UserMemoryAdapters.pull_memories(user=user2) 442 443 # Assert: Each user only sees their own memories 444 assert len(user1_memories) == 1 445 assert user1_memories[0].raw == "user1 memory" 446 447 assert len(user2_memories) == 1 448 assert user2_memories[0].raw == "user2 memory" 449 450 451 @pytest.mark.anyio 452 @pytest.mark.django_db(transaction=True) 453 async def test_custom_agent_cannot_see_other_custom_agent_memories(): 454 """One custom agent should not see another custom agent's memories.""" 455 # Setup 456 user = await acreate_user() 457 await acreate_subscription(user) 458 chat_model = await acreate_chat_model() 459 460 # Create default agent 461 await acreate_default_agent() 462 463 # Create two custom agents 464 agent_accountant = await acreate_agent("Accountant", chat_model, "Financial advisor") 465 agent_chef = await acreate_agent("Chef", chat_model, "Cooking expert") 466 467 # Create memories for each agent 468 await acreate_test_memory(user, agent=agent_accountant, raw_text="user spent $500 on groceries") 469 await acreate_test_memory(user, agent=agent_chef, raw_text="user likes Italian food") 470 471 # Act & Assert: Accountant only sees financial memories 472 accountant_memories = await UserMemoryAdapters.pull_memories(user=user, agent=agent_accountant) 473 assert len(accountant_memories) == 1 474 assert accountant_memories[0].raw == "user spent $500 on groceries" 475 476 # Act & Assert: Chef only sees food memories 477 chef_memories = await UserMemoryAdapters.pull_memories(user=user, agent=agent_chef) 478 assert len(chef_memories) == 1 479 assert chef_memories[0].raw == "user likes Italian food"