cortexdb-ag2 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,79 @@
1
+ # Rust
2
+ /target
3
+ **/*.rs.bk
4
+
5
+ # Environment / secrets
6
+ .env
7
+ .env.local
8
+ .env*.local
9
+ *.pem
10
+ *.key
11
+ .npmrc
12
+
13
+ # SQLite database
14
+ *.sqlite
15
+ *.sqlite-wal
16
+ *.sqlite-shm
17
+
18
+ # OS
19
+ .DS_Store
20
+ Thumbs.db
21
+ desktop.ini
22
+
23
+ # IDE
24
+ .idea/
25
+ .vscode/
26
+ *.swp
27
+ *.swo
28
+
29
+ # Data directories
30
+ cortexdb_data*/
31
+ /data/
32
+ # Per-bench tenant stores (RocksDB + Tantivy + HNSW state; regeneratable per run)
33
+ /data_*/
34
+ # Experimental per-branch stores (not tracked on this branch but left gitignored
35
+ # so checkout from other branches doesn't surface them in git status)
36
+ /event_memory_store/
37
+ /llm_cache/
38
+
39
+ # Benchmark inputs and per-run outputs (kept local, regenerated each run)
40
+ benchmarks/longmemeval/data/
41
+ benchmarks/longmemeval/server_results/
42
+ benchmarks/longmemeval/fast_results/
43
+ benchmarks/longmemeval/micro_results/
44
+ benchmarks/longmemeval/server_logs/
45
+ benchmarks/longmemeval/*.log
46
+ benchmarks/locomo/locomo_results*.json
47
+ benchmarks/locomo/server_results/
48
+ benchmarks/locomo/*.log
49
+ /answer_out.json
50
+
51
+ # Local Claude Code state
52
+ .claude/
53
+ .tmp/
54
+
55
+ # Python
56
+ __pycache__/
57
+ *.pyc
58
+ .venv/
59
+ venv/
60
+
61
+ # Node
62
+ node_modules/
63
+ dist/
64
+ .next/
65
+
66
+ # Egg info
67
+ *.egg-info/
68
+
69
+ # Scratch/debug text files at root
70
+ /*.txt
71
+ /*.log
72
+
73
+ # Local debug / marketing / private content (not for repo)
74
+ harness/.reports/
75
+ harness_data_*/
76
+ blog/
77
+ sales/
78
+ videos/
79
+ local-instance/
@@ -0,0 +1,31 @@
1
+ Metadata-Version: 2.4
2
+ Name: cortexdb-ag2
3
+ Version: 0.1.0
4
+ Summary: AG2 integration for CortexDB — long-term memory for AI agent conversations
5
+ License-Expression: Apache-2.0
6
+ Requires-Python: >=3.10
7
+ Requires-Dist: ag2>=0.4
8
+ Requires-Dist: cortexdbai>=0.1.0
9
+ Provides-Extra: dev
10
+ Requires-Dist: pytest>=7.0; extra == 'dev'
11
+ Description-Content-Type: text/markdown
12
+
13
+ # cortexdb-ag2
14
+
15
+ AG2 (formerly AutoGen) integration for CortexDB long-term memory.
16
+
17
+ > **LLM provider note (audit BLK-2):** the canonical example shows GPT-4o
18
+ > in the AG2 LLM config. CortexDB itself is LLM-agnostic. Swap providers
19
+ > through AG2's standard `llm_config` — Anthropic, Gemini, Mistral, and
20
+ > Together AI all work. Nothing in `cortexdb-ag2` reads `OPENAI_API_KEY`.
21
+
22
+ ## Install
23
+
24
+ ```bash
25
+ pip install cortexdb-ag2
26
+ ```
27
+
28
+ ## API base URL
29
+
30
+ Defaults to `https://api-v1.cortexdb.ai` (audit FRI-8). Override with
31
+ `CORTEXDB_API_URL`.
@@ -0,0 +1,19 @@
1
+ # cortexdb-ag2
2
+
3
+ AG2 (formerly AutoGen) integration for CortexDB long-term memory.
4
+
5
+ > **LLM provider note (audit BLK-2):** the canonical example shows GPT-4o
6
+ > in the AG2 LLM config. CortexDB itself is LLM-agnostic. Swap providers
7
+ > through AG2's standard `llm_config` — Anthropic, Gemini, Mistral, and
8
+ > Together AI all work. Nothing in `cortexdb-ag2` reads `OPENAI_API_KEY`.
9
+
10
+ ## Install
11
+
12
+ ```bash
13
+ pip install cortexdb-ag2
14
+ ```
15
+
16
+ ## API base URL
17
+
18
+ Defaults to `https://api-v1.cortexdb.ai` (audit FRI-8). Override with
19
+ `CORTEXDB_API_URL`.
@@ -0,0 +1,21 @@
1
+ """CortexDB integration for AG2 (successor to AutoGen).
2
+
3
+ Provides a memory-augmented agent mixin and callable tools that connect
4
+ AG2's multi-agent framework to CortexDB's long-term memory system.
5
+ """
6
+
7
+ from cortexdb_ag2.agent import CortexDBAgent
8
+ from cortexdb_ag2.tools import (
9
+ cortexdb_forget_fn,
10
+ cortexdb_search_fn,
11
+ cortexdb_store_fn,
12
+ register_cortexdb_tools,
13
+ )
14
+
15
+ __all__ = [
16
+ "CortexDBAgent",
17
+ "cortexdb_search_fn",
18
+ "cortexdb_store_fn",
19
+ "cortexdb_forget_fn",
20
+ "register_cortexdb_tools",
21
+ ]
@@ -0,0 +1,183 @@
1
+ """AG2 agent with built-in CortexDB memory.
2
+
3
+ Provides a ConversableAgent subclass that automatically stores
4
+ conversation turns and retrieves relevant context from CortexDB
5
+ before generating responses. Compatible with AG2 (the successor
6
+ to Microsoft AutoGen).
7
+ """
8
+
9
+ from __future__ import annotations
10
+
11
+ from typing import Any, Optional, Union
12
+
13
+ from autogen import ConversableAgent
14
+
15
+ from cortexdb import Cortex
16
+
17
+
18
+ class CortexDBAgent(ConversableAgent):
19
+ """AG2 ConversableAgent with automatic CortexDB memory integration.
20
+
21
+ Extends AG2's ConversableAgent to automatically store conversation
22
+ turns in CortexDB and retrieve relevant past context before generating
23
+ responses. This enables the agent to maintain long-term memory across
24
+ conversations and sessions.
25
+
26
+ The agent injects retrieved context as a system-level prefix to the
27
+ conversation, giving the LLM access to relevant memories without
28
+ modifying the visible chat history.
29
+
30
+ Args:
31
+ name: The agent's name.
32
+ cortex_client: An initialized CortexDB client instance.
33
+ scope: The hierarchical scope path for memory isolation.
34
+ auto_store: Whether to automatically store conversation turns.
35
+ Defaults to ``True``.
36
+ auto_recall: Whether to automatically retrieve context before
37
+ responding. Defaults to ``True``.
38
+ **kwargs: Additional keyword arguments passed to ConversableAgent.
39
+
40
+ Example::
41
+
42
+ from cortexdb import Cortex
43
+ from cortexdb_ag2 import CortexDBAgent
44
+
45
+ client = Cortex("http://localhost:3141")
46
+ agent = CortexDBAgent(
47
+ name="memory_assistant",
48
+ cortex_client=client,
49
+ scope="user:default",
50
+ llm_config=llm_config,
51
+ system_message="You are a helpful assistant with long-term memory.",
52
+ )
53
+ """
54
+
55
+ def __init__(
56
+ self,
57
+ name: str,
58
+ cortex_client: Cortex,
59
+ scope: str = "user:default",
60
+ auto_store: bool = True,
61
+ auto_recall: bool = True,
62
+ **kwargs: Any,
63
+ ) -> None:
64
+ super().__init__(name=name, **kwargs)
65
+ self._cortex_client = cortex_client
66
+ self._scope = scope
67
+ self._auto_store = auto_store
68
+ self._auto_recall = auto_recall
69
+
70
+ def generate_reply(
71
+ self,
72
+ messages: Optional[list[dict[str, Any]]] = None,
73
+ sender: Optional[Any] = None,
74
+ **kwargs: Any,
75
+ ) -> Union[str, dict[str, Any], None]:
76
+ """Generate a reply with CortexDB memory augmentation.
77
+
78
+ Before generating a reply, retrieves relevant context from CortexDB
79
+ based on the latest message. After generating, stores the conversation
80
+ turn for future recall.
81
+
82
+ Args:
83
+ messages: The conversation messages to respond to.
84
+ sender: The agent that sent the message.
85
+ **kwargs: Additional keyword arguments.
86
+
87
+ Returns:
88
+ The generated reply string, a structured reply dict, or None.
89
+ """
90
+ if messages and self._auto_recall:
91
+ latest_message = messages[-1]
92
+ content = latest_message.get("content", "")
93
+ if content:
94
+ context = self._recall_context(str(content))
95
+ if context:
96
+ messages = self._inject_memory_context(messages, context)
97
+
98
+ reply = super().generate_reply(messages=messages, sender=sender, **kwargs)
99
+
100
+ if reply and self._auto_store and messages:
101
+ latest_content = messages[-1].get("content", "") if messages else ""
102
+ reply_text = reply if isinstance(reply, str) else reply.get("content", "")
103
+ if latest_content and reply_text:
104
+ self._store_turn(str(latest_content), str(reply_text))
105
+
106
+ return reply
107
+
108
+ def _recall_context(self, query: str) -> str:
109
+ """Retrieve relevant context from CortexDB.
110
+
111
+ Args:
112
+ query: The query to search for relevant memories.
113
+
114
+ Returns:
115
+ The context string from CortexDB, or empty string on failure.
116
+ """
117
+ try:
118
+ result = self._cortex_client.recall(
119
+ self._scope,
120
+ query=query,
121
+ )
122
+ except Exception:
123
+ return ""
124
+
125
+ return result.get("context_block", "") if result else ""
126
+
127
+ def _store_turn(self, user_message: str, assistant_reply: str) -> None:
128
+ """Store a conversation turn in CortexDB.
129
+
130
+ Args:
131
+ user_message: The user's message.
132
+ assistant_reply: The agent's reply.
133
+ """
134
+ content = f"User: {user_message}\nAssistant: {assistant_reply}"
135
+ try:
136
+ self._cortex_client.experience(self._scope, text=content)
137
+ except Exception:
138
+ pass
139
+
140
+ def _inject_memory_context(
141
+ self,
142
+ messages: list[dict[str, Any]],
143
+ context: str,
144
+ ) -> list[dict[str, Any]]:
145
+ """Inject retrieved memory context into the message list.
146
+
147
+ Prepends a system message containing the retrieved context
148
+ so the LLM can reference it when generating a response.
149
+
150
+ Args:
151
+ messages: The original conversation messages.
152
+ context: The retrieved context string from CortexDB.
153
+
154
+ Returns:
155
+ A new message list with the memory context injected.
156
+ """
157
+ context_message = {
158
+ "role": "system",
159
+ "content": (
160
+ f"Relevant context from long-term memory:\n\n{context}\n\n"
161
+ "Use this context to inform your response if relevant."
162
+ ),
163
+ }
164
+
165
+ augmented = list(messages)
166
+ augmented.insert(0, context_message)
167
+ return augmented
168
+
169
+ def clear_memory(self) -> dict[str, Any]:
170
+ """Clear all memories for this agent's scope.
171
+
172
+ Removes all stored memories from CortexDB for the configured
173
+ scope path.
174
+
175
+ Returns:
176
+ The forget response dict from CortexDB.
177
+ """
178
+ return self._cortex_client.forget(
179
+ self._scope,
180
+ confirm_all=True,
181
+ cascade="redact_events",
182
+ reason="AG2 agent memory clear requested",
183
+ )
@@ -0,0 +1,165 @@
1
+ """AG2 callable tools for interacting with CortexDB.
2
+
3
+ Provides factory functions that create CortexDB-backed callables
4
+ suitable for registration with AG2's ConversableAgent tool system.
5
+ """
6
+
7
+ from __future__ import annotations
8
+
9
+ from typing import Annotated, Any, Callable, Optional
10
+
11
+ from cortexdb import Cortex
12
+
13
+
14
+ def cortexdb_search_fn(
15
+ client: Cortex,
16
+ scope: str = "user:default",
17
+ ) -> Callable[..., str]:
18
+ """Create a search function for use as an AG2 tool.
19
+
20
+ Returns a callable that performs semantic recall over CortexDB's
21
+ memory store, suitable for registration with an AG2 agent.
22
+
23
+ Args:
24
+ client: An initialized CortexDB client instance.
25
+ scope: The hierarchical scope path for memory isolation.
26
+
27
+ Returns:
28
+ A callable function that accepts a ``query`` parameter and
29
+ returns the recalled context block.
30
+
31
+ Example::
32
+
33
+ from cortexdb import Cortex
34
+ from cortexdb_ag2 import cortexdb_search_fn
35
+
36
+ client = Cortex("http://localhost:3141")
37
+ search = cortexdb_search_fn(client, scope="user:default")
38
+
39
+ agent.register_for_llm(description="Search memories")(search)
40
+ user_proxy.register_for_execution()(search)
41
+ """
42
+
43
+ def search(
44
+ query: Annotated[str, "The search query to find relevant memories."],
45
+ ) -> str:
46
+ """Search CortexDB for relevant memories and past context."""
47
+ result = client.recall(scope, query=query)
48
+
49
+ context = result.get("context_block", "") if result else ""
50
+ if not context:
51
+ return "No relevant memories found."
52
+
53
+ return context
54
+
55
+ search.__name__ = "cortexdb_search"
56
+ search.__doc__ = "Search CortexDB for relevant memories and past context."
57
+ return search
58
+
59
+
60
+ def cortexdb_store_fn(
61
+ client: Cortex,
62
+ scope: str = "user:default",
63
+ ) -> Callable[..., str]:
64
+ """Create a store function for use as an AG2 tool.
65
+
66
+ Returns a callable that persists information into CortexDB's
67
+ long-term memory, suitable for registration with an AG2 agent.
68
+
69
+ Args:
70
+ client: An initialized CortexDB client instance.
71
+ scope: The hierarchical scope path for memory isolation.
72
+
73
+ Returns:
74
+ A callable function that accepts a ``content`` parameter and
75
+ returns a confirmation message.
76
+ """
77
+
78
+ def store(
79
+ content: Annotated[str, "The content to store as a memory."],
80
+ ) -> str:
81
+ """Store information in CortexDB's long-term memory."""
82
+ client.experience(scope, text=content)
83
+ return "Memory stored successfully in CortexDB."
84
+
85
+ store.__name__ = "cortexdb_store"
86
+ store.__doc__ = "Store information in CortexDB's long-term memory."
87
+ return store
88
+
89
+
90
+ def cortexdb_forget_fn(
91
+ client: Cortex,
92
+ scope: str = "user:default",
93
+ ) -> Callable[..., str]:
94
+ """Create a forget function for use as an AG2 tool.
95
+
96
+ Returns a callable that removes memories from CortexDB, suitable
97
+ for registration with an AG2 agent.
98
+
99
+ Args:
100
+ client: An initialized CortexDB client instance.
101
+ scope: The hierarchical scope path for memory isolation.
102
+
103
+ Returns:
104
+ A callable function that accepts a ``reason`` parameter
105
+ and returns a confirmation message.
106
+ """
107
+
108
+ def forget(
109
+ reason: Annotated[str, "The reason for forgetting these memories."],
110
+ ) -> str:
111
+ """Forget or remove memories from CortexDB."""
112
+ client.forget(
113
+ scope,
114
+ confirm_all=True,
115
+ cascade="redact_events",
116
+ reason=reason,
117
+ )
118
+ return f"Memories in scope '{scope}' have been forgotten."
119
+
120
+ forget.__name__ = "cortexdb_forget"
121
+ forget.__doc__ = "Forget or remove memories from CortexDB."
122
+ return forget
123
+
124
+
125
+ def register_cortexdb_tools(
126
+ agent: Any,
127
+ executor: Any,
128
+ client: Cortex,
129
+ scope: str = "user:default",
130
+ ) -> None:
131
+ """Register all CortexDB tools with an AG2 agent and executor.
132
+
133
+ Convenience function that creates and registers search, store, and
134
+ forget tools on both the LLM-facing agent and the execution proxy.
135
+
136
+ Args:
137
+ agent: The AG2 ConversableAgent that will call the tools.
138
+ executor: The AG2 agent (typically UserProxyAgent) that
139
+ executes tool calls.
140
+ client: An initialized CortexDB client instance.
141
+ scope: The hierarchical scope path for memory isolation.
142
+
143
+ Example::
144
+
145
+ from ag2 import ConversableAgent, UserProxyAgent
146
+ from cortexdb import Cortex
147
+ from cortexdb_ag2 import register_cortexdb_tools
148
+
149
+ client = Cortex("http://localhost:3141")
150
+ assistant = ConversableAgent("assistant", llm_config=llm_config)
151
+ user_proxy = UserProxyAgent("user_proxy")
152
+
153
+ register_cortexdb_tools(assistant, user_proxy, client, scope="user:default")
154
+ """
155
+ search = cortexdb_search_fn(client, scope)
156
+ store = cortexdb_store_fn(client, scope)
157
+ forget = cortexdb_forget_fn(client, scope)
158
+
159
+ for fn, desc in [
160
+ (search, "Search CortexDB for relevant memories and past context."),
161
+ (store, "Store information in CortexDB's long-term memory."),
162
+ (forget, "Forget or remove memories from CortexDB."),
163
+ ]:
164
+ agent.register_for_llm(description=desc)(fn)
165
+ executor.register_for_execution()(fn)
@@ -0,0 +1,20 @@
1
+ [build-system]
2
+ requires = ["hatchling"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "cortexdb-ag2"
7
+ version = "0.1.0"
8
+ description = "AG2 integration for CortexDB — long-term memory for AI agent conversations"
9
+ readme = "README.md"
10
+ requires-python = ">=3.10"
11
+ license = "Apache-2.0"
12
+ dependencies = [
13
+ "cortexdbai>=0.1.0",
14
+ "ag2>=0.4",
15
+ ]
16
+
17
+ [project.optional-dependencies]
18
+ dev = [
19
+ "pytest>=7.0",
20
+ ]
File without changes
@@ -0,0 +1,182 @@
1
+ """Tests for the CortexDBAgent AG2 integration.
2
+
3
+ Validates that the memory-augmented ConversableAgent subclass correctly
4
+ stores turns, recalls context, injects memories, and respects its
5
+ configuration flags against the v1 SDK surface.
6
+ """
7
+
8
+ from __future__ import annotations
9
+
10
+ import sys
11
+ import unittest
12
+ from unittest.mock import MagicMock, patch
13
+
14
+
15
+ class _MockConversableAgent:
16
+ """Minimal stand-in for ag2.ConversableAgent."""
17
+
18
+ def __init__(self, name: str = "", **kwargs):
19
+ self.name = name
20
+ self._kwargs = kwargs
21
+
22
+ def generate_reply(self, messages=None, sender=None, **kwargs):
23
+ return "mock reply"
24
+
25
+
26
+ mock_ag2 = MagicMock()
27
+ mock_ag2.ConversableAgent = _MockConversableAgent
28
+
29
+ mock_cortexdb = MagicMock()
30
+ mock_cortexdb.Cortex = MagicMock()
31
+
32
+ with patch.dict(sys.modules, {
33
+ "ag2": mock_ag2,
34
+ "cortexdb": mock_cortexdb,
35
+ }):
36
+ from cortexdb_ag2.agent import CortexDBAgent
37
+
38
+
39
+ class TestCortexDBAgentInit(unittest.TestCase):
40
+ """Tests for CortexDBAgent construction."""
41
+
42
+ def test_sets_scope(self) -> None:
43
+ client = MagicMock()
44
+ agent = CortexDBAgent(
45
+ name="test",
46
+ cortex_client=client,
47
+ scope="org:acme/user:bob",
48
+ )
49
+ self.assertEqual(agent._scope, "org:acme/user:bob")
50
+ self.assertIs(agent._cortex_client, client)
51
+
52
+ def test_defaults(self) -> None:
53
+ client = MagicMock()
54
+ agent = CortexDBAgent(name="test", cortex_client=client)
55
+ self.assertEqual(agent._scope, "user:default")
56
+ self.assertTrue(agent._auto_store)
57
+ self.assertTrue(agent._auto_recall)
58
+
59
+ def test_auto_flags(self) -> None:
60
+ client = MagicMock()
61
+ agent = CortexDBAgent(
62
+ name="test", cortex_client=client,
63
+ auto_store=False, auto_recall=False,
64
+ )
65
+ self.assertFalse(agent._auto_store)
66
+ self.assertFalse(agent._auto_recall)
67
+
68
+
69
+ class TestRecallContext(unittest.TestCase):
70
+ """Tests for _recall_context."""
71
+
72
+ def setUp(self) -> None:
73
+ self.client = MagicMock()
74
+ self.agent = CortexDBAgent(
75
+ name="test", cortex_client=self.client, scope="user:t1",
76
+ )
77
+
78
+ def test_calls_recall(self) -> None:
79
+ self.client.recall.return_value = {"context_block": "mem1"}
80
+ result = self.agent._recall_context("query")
81
+ self.client.recall.assert_called_once_with("user:t1", query="query")
82
+ self.assertEqual(result, "mem1")
83
+
84
+ def test_handles_empty_context(self) -> None:
85
+ self.client.recall.return_value = {"context_block": ""}
86
+ result = self.agent._recall_context("q")
87
+ self.assertEqual(result, "")
88
+
89
+ def test_handles_exception_gracefully(self) -> None:
90
+ self.client.recall.side_effect = ConnectionError("offline")
91
+ result = self.agent._recall_context("q")
92
+ self.assertEqual(result, "")
93
+
94
+
95
+ class TestStoreTurn(unittest.TestCase):
96
+ """Tests for _store_turn."""
97
+
98
+ def setUp(self) -> None:
99
+ self.client = MagicMock()
100
+ self.agent = CortexDBAgent(
101
+ name="test", cortex_client=self.client, scope="user:t1",
102
+ )
103
+
104
+ def test_calls_experience_with_formatted_content(self) -> None:
105
+ self.agent._store_turn("user msg", "assistant reply")
106
+ self.client.experience.assert_called_once_with(
107
+ "user:t1",
108
+ text="User: user msg\nAssistant: assistant reply",
109
+ )
110
+
111
+ def test_handles_exception_gracefully(self) -> None:
112
+ self.client.experience.side_effect = RuntimeError("fail")
113
+ self.agent._store_turn("u", "a")
114
+
115
+
116
+ class TestInjectMemoryContext(unittest.TestCase):
117
+ """Tests for _inject_memory_context."""
118
+
119
+ def setUp(self) -> None:
120
+ self.client = MagicMock()
121
+ self.agent = CortexDBAgent(name="test", cortex_client=self.client)
122
+
123
+ def test_prepends_system_message(self) -> None:
124
+ messages = [{"role": "user", "content": "hello"}]
125
+ result = self.agent._inject_memory_context(messages, "memory one")
126
+ self.assertEqual(len(result), 2)
127
+ self.assertEqual(result[0]["role"], "system")
128
+ self.assertIn("memory one", result[0]["content"])
129
+
130
+ def test_does_not_mutate_original(self) -> None:
131
+ messages = [{"role": "user", "content": "hi"}]
132
+ original_len = len(messages)
133
+ self.agent._inject_memory_context(messages, "m")
134
+ self.assertEqual(len(messages), original_len)
135
+
136
+
137
+ class TestClearMemory(unittest.TestCase):
138
+ """Tests for clear_memory."""
139
+
140
+ def test_calls_forget(self) -> None:
141
+ client = MagicMock()
142
+ agent = CortexDBAgent(
143
+ name="test", cortex_client=client, scope="user:t1",
144
+ )
145
+ agent.clear_memory()
146
+ client.forget.assert_called_once_with(
147
+ "user:t1",
148
+ confirm_all=True,
149
+ cascade="redact_events",
150
+ reason="AG2 agent memory clear requested",
151
+ )
152
+
153
+
154
+ class TestAutoStoreDisabled(unittest.TestCase):
155
+ """Tests that auto_store=False skips storing."""
156
+
157
+ def test_no_store_when_disabled(self) -> None:
158
+ client = MagicMock()
159
+ client.recall.return_value = {"context_block": ""}
160
+ agent = CortexDBAgent(
161
+ name="test", cortex_client=client, auto_store=False,
162
+ )
163
+ messages = [{"role": "user", "content": "hello"}]
164
+ agent.generate_reply(messages=messages)
165
+ client.experience.assert_not_called()
166
+
167
+
168
+ class TestAutoRecallDisabled(unittest.TestCase):
169
+ """Tests that auto_recall=False skips recall."""
170
+
171
+ def test_no_recall_when_disabled(self) -> None:
172
+ client = MagicMock()
173
+ agent = CortexDBAgent(
174
+ name="test", cortex_client=client, auto_recall=False,
175
+ )
176
+ messages = [{"role": "user", "content": "hello"}]
177
+ agent.generate_reply(messages=messages)
178
+ client.recall.assert_not_called()
179
+
180
+
181
+ if __name__ == "__main__":
182
+ unittest.main()
@@ -0,0 +1,155 @@
1
+ """Tests for AG2 CortexDB tools.
2
+
3
+ Validates that factory functions produce correctly configured callables
4
+ that delegate to the CortexDB client for search, store, and forget operations.
5
+ """
6
+
7
+ from __future__ import annotations
8
+
9
+ import sys
10
+ import unittest
11
+ from unittest.mock import MagicMock, patch
12
+
13
+
14
+ mock_cortexdb = MagicMock()
15
+ mock_cortex_class = MagicMock()
16
+ mock_cortexdb.Cortex = mock_cortex_class
17
+
18
+ mock_ag2 = MagicMock()
19
+ mock_ag2.ConversableAgent = type(
20
+ "ConversableAgent", (), {"__init__": lambda self, **kw: None},
21
+ )
22
+
23
+ with patch.dict(sys.modules, {"cortexdb": mock_cortexdb, "ag2": mock_ag2}):
24
+ from cortexdb_ag2.tools import (
25
+ cortexdb_forget_fn,
26
+ cortexdb_search_fn,
27
+ cortexdb_store_fn,
28
+ register_cortexdb_tools,
29
+ )
30
+
31
+
32
+ class TestCortexDBSearchFn(unittest.TestCase):
33
+ """Tests for the cortexdb_search_fn factory."""
34
+
35
+ def setUp(self) -> None:
36
+ self.client = MagicMock()
37
+
38
+ def test_returns_callable(self) -> None:
39
+ fn = cortexdb_search_fn(self.client)
40
+ self.assertTrue(callable(fn))
41
+
42
+ def test_function_name(self) -> None:
43
+ fn = cortexdb_search_fn(self.client)
44
+ self.assertEqual(fn.__name__, "cortexdb_search")
45
+
46
+ def test_calls_recall(self) -> None:
47
+ self.client.recall.return_value = {"context_block": ""}
48
+ fn = cortexdb_search_fn(self.client, scope="user:t1")
49
+ fn(query="hello")
50
+ self.client.recall.assert_called_once_with("user:t1", query="hello")
51
+
52
+ def test_with_results_returns_context(self) -> None:
53
+ self.client.recall.return_value = {"context_block": "memory one"}
54
+ fn = cortexdb_search_fn(self.client)
55
+ result = fn(query="test")
56
+ self.assertEqual(result, "memory one")
57
+
58
+ def test_with_no_results(self) -> None:
59
+ self.client.recall.return_value = {"context_block": ""}
60
+ fn = cortexdb_search_fn(self.client)
61
+ result = fn(query="nothing")
62
+ self.assertEqual(result, "No relevant memories found.")
63
+
64
+ def test_default_scope(self) -> None:
65
+ self.client.recall.return_value = {"context_block": ""}
66
+ fn = cortexdb_search_fn(self.client)
67
+ fn(query="q")
68
+ self.client.recall.assert_called_once_with("user:default", query="q")
69
+
70
+
71
+ class TestCortexDBStoreFn(unittest.TestCase):
72
+ """Tests for the cortexdb_store_fn factory."""
73
+
74
+ def setUp(self) -> None:
75
+ self.client = MagicMock()
76
+
77
+ def test_returns_callable(self) -> None:
78
+ fn = cortexdb_store_fn(self.client)
79
+ self.assertTrue(callable(fn))
80
+
81
+ def test_function_name(self) -> None:
82
+ fn = cortexdb_store_fn(self.client)
83
+ self.assertEqual(fn.__name__, "cortexdb_store")
84
+
85
+ def test_calls_experience(self) -> None:
86
+ fn = cortexdb_store_fn(self.client, scope="user:t1")
87
+ result = fn(content="important fact")
88
+ self.client.experience.assert_called_once_with(
89
+ "user:t1", text="important fact",
90
+ )
91
+ self.assertIn("stored successfully", result)
92
+
93
+ def test_default_scope(self) -> None:
94
+ fn = cortexdb_store_fn(self.client)
95
+ fn(content="data")
96
+ self.client.experience.assert_called_once_with(
97
+ "user:default", text="data",
98
+ )
99
+
100
+
101
+ class TestCortexDBForgetFn(unittest.TestCase):
102
+ """Tests for the cortexdb_forget_fn factory."""
103
+
104
+ def setUp(self) -> None:
105
+ self.client = MagicMock()
106
+
107
+ def test_returns_callable(self) -> None:
108
+ fn = cortexdb_forget_fn(self.client)
109
+ self.assertTrue(callable(fn))
110
+
111
+ def test_function_name(self) -> None:
112
+ fn = cortexdb_forget_fn(self.client)
113
+ self.assertEqual(fn.__name__, "cortexdb_forget")
114
+
115
+ def test_calls_forget_with_reason(self) -> None:
116
+ fn = cortexdb_forget_fn(self.client, scope="user:t1")
117
+ result = fn(reason="outdated")
118
+ self.client.forget.assert_called_once_with(
119
+ "user:t1",
120
+ confirm_all=True,
121
+ cascade="redact_events",
122
+ reason="outdated",
123
+ )
124
+ self.assertIn("user:t1", result)
125
+ self.assertIn("forgotten", result)
126
+
127
+
128
+ class TestRegisterCortexDBTools(unittest.TestCase):
129
+ """Tests for register_cortexdb_tools."""
130
+
131
+ def setUp(self) -> None:
132
+ self.client = MagicMock()
133
+ self.agent = MagicMock()
134
+ self.executor = MagicMock()
135
+ self.agent.register_for_llm.return_value = lambda fn: fn
136
+ self.executor.register_for_execution.return_value = lambda fn: fn
137
+
138
+ def test_registers_all_three_tools(self) -> None:
139
+ register_cortexdb_tools(self.agent, self.executor, self.client)
140
+ self.assertEqual(self.agent.register_for_llm.call_count, 3)
141
+ self.assertEqual(self.executor.register_for_execution.call_count, 3)
142
+
143
+ def test_descriptions_passed_to_agent(self) -> None:
144
+ register_cortexdb_tools(self.agent, self.executor, self.client)
145
+ descriptions = [
146
+ call.kwargs["description"]
147
+ for call in self.agent.register_for_llm.call_args_list
148
+ ]
149
+ self.assertTrue(any("Search" in d for d in descriptions))
150
+ self.assertTrue(any("Store" in d for d in descriptions))
151
+ self.assertTrue(any("Forget" in d for d in descriptions))
152
+
153
+
154
+ if __name__ == "__main__":
155
+ unittest.main()