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.
- cortexdb_ag2-0.1.0/.gitignore +79 -0
- cortexdb_ag2-0.1.0/PKG-INFO +31 -0
- cortexdb_ag2-0.1.0/README.md +19 -0
- cortexdb_ag2-0.1.0/cortexdb_ag2/__init__.py +21 -0
- cortexdb_ag2-0.1.0/cortexdb_ag2/agent.py +183 -0
- cortexdb_ag2-0.1.0/cortexdb_ag2/tools.py +165 -0
- cortexdb_ag2-0.1.0/pyproject.toml +20 -0
- cortexdb_ag2-0.1.0/tests/__init__.py +0 -0
- cortexdb_ag2-0.1.0/tests/test_agent.py +182 -0
- cortexdb_ag2-0.1.0/tests/test_tools.py +155 -0
|
@@ -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()
|