memx-ai 0.1.8__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.
- memx_ai-0.1.8/.gitignore +14 -0
- memx_ai-0.1.8/LICENSE +21 -0
- memx_ai-0.1.8/PKG-INFO +183 -0
- memx_ai-0.1.8/README.md +158 -0
- memx_ai-0.1.8/memx/__init__.py +0 -0
- memx_ai-0.1.8/memx/engine/__init__.py +15 -0
- memx_ai-0.1.8/memx/engine/config.py +8 -0
- memx_ai-0.1.8/memx/engine/mongodb.py +34 -0
- memx_ai-0.1.8/memx/engine/postgres.py +122 -0
- memx_ai-0.1.8/memx/engine/sqlite.py +106 -0
- memx_ai-0.1.8/memx/memory/__init__.py +14 -0
- memx_ai-0.1.8/memx/memory/mongodb.py +66 -0
- memx_ai-0.1.8/memx/memory/postgres.py +97 -0
- memx_ai-0.1.8/memx/memory/redis.py +0 -0
- memx_ai-0.1.8/memx/memory/simple.py +101 -0
- memx_ai-0.1.8/memx/memory/sqlite.py +105 -0
- memx_ai-0.1.8/pyproject.toml +81 -0
- memx_ai-0.1.8/tests/conftest.py +0 -0
- memx_ai-0.1.8/tests/test_postgres.py +0 -0
- memx_ai-0.1.8/tests/test_sqlite.py +0 -0
memx_ai-0.1.8/.gitignore
ADDED
memx_ai-0.1.8/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
The MIT License (MIT)
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2025 Pablo Galilea
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in
|
|
13
|
+
all copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
|
|
21
|
+
THE SOFTWARE.
|
memx_ai-0.1.8/PKG-INFO
ADDED
|
@@ -0,0 +1,183 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: memx-ai
|
|
3
|
+
Version: 0.1.8
|
|
4
|
+
Summary: Lightweight and extensible memory layer for LLMs
|
|
5
|
+
Author-email: Pablo Galilea <pgalilea@me.com>
|
|
6
|
+
License-Expression: MIT
|
|
7
|
+
License-File: LICENSE
|
|
8
|
+
Requires-Python: >=3.11
|
|
9
|
+
Requires-Dist: aiofiles>=25.1.0
|
|
10
|
+
Requires-Dist: orjson>=3.11.4
|
|
11
|
+
Provides-Extra: all
|
|
12
|
+
Requires-Dist: aiosqlite>=0.21.0; extra == 'all'
|
|
13
|
+
Requires-Dist: psycopg>=3.2.12; extra == 'all'
|
|
14
|
+
Requires-Dist: pymongo>=4.15.3; extra == 'all'
|
|
15
|
+
Requires-Dist: sqlalchemy[asyncio,postgresql-psycopg]>=2.0.44; extra == 'all'
|
|
16
|
+
Provides-Extra: mongodb
|
|
17
|
+
Requires-Dist: pymongo>=4.15.3; extra == 'mongodb'
|
|
18
|
+
Provides-Extra: postgres
|
|
19
|
+
Requires-Dist: psycopg>=3.2.12; extra == 'postgres'
|
|
20
|
+
Requires-Dist: sqlalchemy[asyncio,postgresql-psycopg]>=2.0.44; extra == 'postgres'
|
|
21
|
+
Provides-Extra: sqlite
|
|
22
|
+
Requires-Dist: aiosqlite>=0.21.0; extra == 'sqlite'
|
|
23
|
+
Requires-Dist: sqlalchemy[asyncio]>=2.0.44; extra == 'sqlite'
|
|
24
|
+
Description-Content-Type: text/markdown
|
|
25
|
+
|
|
26
|
+
<p align="center">
|
|
27
|
+
<a href="https://github.com/pgalilea/memx"><img src="https://i.ibb.co/JjYq8fzW/memx.png" alt="memx - memory layer"></a>
|
|
28
|
+
</p>
|
|
29
|
+
|
|
30
|
+
<br/>
|
|
31
|
+
Lightweight and extensible memory layer for LLMs.
|
|
32
|
+
<br/><br/>
|
|
33
|
+
|
|
34
|
+
**Important Disclaimer**: This library is intended to be production-ready, but currently is in active development. Fix the version and run your own tests :)
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
## 🔥 Key Features
|
|
38
|
+
- **Framework agnostic**: Use your preferred AI agent framework.
|
|
39
|
+
- **Own infrastructure**: Use your preferred cloud provider. No third-party api keys; your data, your rules.
|
|
40
|
+
- **Multiple backends**: Move from your local *POC* to production deployment, seamlessly (SQLite, MongoDB, PostgreSQL).
|
|
41
|
+
- **Sync and async api**: Highly compatible with modern and *legacy* frameworks.
|
|
42
|
+
- **No forced schema**: As long it is a list of json serializable objects.
|
|
43
|
+
- **Resumable memory**: Perfect for chat applications and REST APIs
|
|
44
|
+
- **Robust**: Get production-ready code with minimal effort.
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
## ⚙️ Installation
|
|
48
|
+
|
|
49
|
+
From pypi
|
|
50
|
+
```bash
|
|
51
|
+
pip install memx-ai
|
|
52
|
+
```
|
|
53
|
+
Or clone the repo and install it
|
|
54
|
+
```bash
|
|
55
|
+
pip install .
|
|
56
|
+
```
|
|
57
|
+
|
|
58
|
+
## 🚀 Quickstart
|
|
59
|
+
|
|
60
|
+
### OpenAI
|
|
61
|
+
Simple conversation with [OpenAI Python library](https://github.com/openai/openai-python)
|
|
62
|
+
```Python
|
|
63
|
+
# https://platform.openai.com/docs/guides/conversation-state?api-mode=responses
|
|
64
|
+
# tested on openai==2.6.1
|
|
65
|
+
|
|
66
|
+
from openai import OpenAI
|
|
67
|
+
from memx.memory.sqlite import SQLiteMemory
|
|
68
|
+
|
|
69
|
+
sqlite_uri = "sqlite+aiosqlite:///message-storage.db"
|
|
70
|
+
engine = SQLiteEngine(sqlite_uri, "memx-messages", start_up=True)
|
|
71
|
+
m1 = engine.create_session() # create a new session
|
|
72
|
+
|
|
73
|
+
client = OpenAI()
|
|
74
|
+
|
|
75
|
+
m1.sync.add([{"role": "user", "content": "tell me a good joke about programmers"}])
|
|
76
|
+
|
|
77
|
+
first_response = client.responses.create(
|
|
78
|
+
model="gpt-4o-mini", input=m1.sync.get(), store=False
|
|
79
|
+
)
|
|
80
|
+
|
|
81
|
+
print(first_response.output_text)
|
|
82
|
+
|
|
83
|
+
m1.sync.add(
|
|
84
|
+
[{"role": r.role, "content": r.content[0].text} for r in first_response.output]
|
|
85
|
+
)
|
|
86
|
+
|
|
87
|
+
m1.sync.add([{"role": "user", "content": "tell me another"}])
|
|
88
|
+
|
|
89
|
+
second_response = client.responses.create(
|
|
90
|
+
model="gpt-4o-mini", input=m1.sync.get(), store=False
|
|
91
|
+
)
|
|
92
|
+
|
|
93
|
+
m1.sync.add(
|
|
94
|
+
[{"role": r.role, "content": r.content[0].text} for r in second_response.output]
|
|
95
|
+
)
|
|
96
|
+
|
|
97
|
+
print(f"\n\n{second_response.output_text}")
|
|
98
|
+
|
|
99
|
+
print(m1.sync.get())
|
|
100
|
+
```
|
|
101
|
+
### Pydantic AI
|
|
102
|
+
Message history with async [Pydantic AI](https://ai.pydantic.dev/) + OpenAI
|
|
103
|
+
|
|
104
|
+
```Python
|
|
105
|
+
# Reference: https://ai.pydantic.dev/message-history/
|
|
106
|
+
|
|
107
|
+
import asyncio
|
|
108
|
+
|
|
109
|
+
import orjson
|
|
110
|
+
from pydantic_ai import Agent, ModelMessagesTypeAdapter
|
|
111
|
+
|
|
112
|
+
from memx.engine.sqlite import SQLiteEngine
|
|
113
|
+
|
|
114
|
+
agent = Agent("openai:gpt-4o-mini")
|
|
115
|
+
|
|
116
|
+
|
|
117
|
+
async def main():
|
|
118
|
+
sqlite_uri = "sqlite+aiosqlite:///message_store.db"
|
|
119
|
+
engine = SQLiteEngine(sqlite_uri, "memx-messages", start_up=True)
|
|
120
|
+
m1 = engine.create_session() # create a new session
|
|
121
|
+
|
|
122
|
+
result1 = await agent.run('Where does "hello world" come from?')
|
|
123
|
+
|
|
124
|
+
# it is your responsibility to add the messages as a list[dict]
|
|
125
|
+
messages = orjson.loads(result1.new_messages_json())
|
|
126
|
+
|
|
127
|
+
await m1.add(messages) # messages: list[dict] must be json serializable
|
|
128
|
+
|
|
129
|
+
session_id = m1.get_id()
|
|
130
|
+
print("Messages added with session_id: ", session_id)
|
|
131
|
+
|
|
132
|
+
# resume the conversation from 'another' memory
|
|
133
|
+
m2 = await engine.get_session(session_id)
|
|
134
|
+
old_messages = ModelMessagesTypeAdapter.validate_python(await m2.get())
|
|
135
|
+
|
|
136
|
+
print("Past messages:\n", old_messages)
|
|
137
|
+
|
|
138
|
+
result2 = await agent.run(
|
|
139
|
+
"Could you tell me more about the authors?", message_history=old_messages
|
|
140
|
+
)
|
|
141
|
+
print("\n\nContext aware result:\n", result2.output)
|
|
142
|
+
|
|
143
|
+
|
|
144
|
+
if __name__ == "__main__":
|
|
145
|
+
asyncio.run(main())
|
|
146
|
+
|
|
147
|
+
|
|
148
|
+
```
|
|
149
|
+
|
|
150
|
+
You can change the memory backend with minimal modifications. Same api to add and get messages.
|
|
151
|
+
```Python
|
|
152
|
+
from memx.memory.mongodb import MongoDBMemory
|
|
153
|
+
from memx.memory.postgres import PostgresMemory
|
|
154
|
+
from memx.memory.sqlite import SQLiteMemory
|
|
155
|
+
|
|
156
|
+
# SQLite backend
|
|
157
|
+
sqlite_uri = "sqlite+aiosqlite:///message_store.db"
|
|
158
|
+
e1 = SQLiteMemory(sqlite_uri, "memx-messages", start_up=True)
|
|
159
|
+
m1 = e1.create_session() # memory session ready to go
|
|
160
|
+
|
|
161
|
+
# PostgreSQL backend
|
|
162
|
+
pg_uri = "postgresql+psycopg://admin:1234@localhost:5433/test-database"
|
|
163
|
+
e2 = PostgresMemory(pg_uri, "memx-messages", start_up=True)
|
|
164
|
+
m2 = e2.create_session()
|
|
165
|
+
|
|
166
|
+
# MongoDB backend
|
|
167
|
+
mongodb_uri = "mongodb://admin:1234@localhost:27017"
|
|
168
|
+
e3 = MongoDBMemory(uri=mongodb_uri, database="memx-test", "memx-messages")
|
|
169
|
+
m3 = e3.create_session()
|
|
170
|
+
|
|
171
|
+
```
|
|
172
|
+
|
|
173
|
+
[More examples...](examples/)
|
|
174
|
+
|
|
175
|
+
## Tasks
|
|
176
|
+
- [x] Add mongodb backend
|
|
177
|
+
- [x] Add SQLite backend
|
|
178
|
+
- [x] Add Postgres backend
|
|
179
|
+
- [ ] Add redis backend
|
|
180
|
+
- [ ] Add tests
|
|
181
|
+
- [x] Publish on pypi
|
|
182
|
+
- [ ] Add full sync support
|
|
183
|
+
- [ ] Add docstrings
|
memx_ai-0.1.8/README.md
ADDED
|
@@ -0,0 +1,158 @@
|
|
|
1
|
+
<p align="center">
|
|
2
|
+
<a href="https://github.com/pgalilea/memx"><img src="https://i.ibb.co/JjYq8fzW/memx.png" alt="memx - memory layer"></a>
|
|
3
|
+
</p>
|
|
4
|
+
|
|
5
|
+
<br/>
|
|
6
|
+
Lightweight and extensible memory layer for LLMs.
|
|
7
|
+
<br/><br/>
|
|
8
|
+
|
|
9
|
+
**Important Disclaimer**: This library is intended to be production-ready, but currently is in active development. Fix the version and run your own tests :)
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
## 🔥 Key Features
|
|
13
|
+
- **Framework agnostic**: Use your preferred AI agent framework.
|
|
14
|
+
- **Own infrastructure**: Use your preferred cloud provider. No third-party api keys; your data, your rules.
|
|
15
|
+
- **Multiple backends**: Move from your local *POC* to production deployment, seamlessly (SQLite, MongoDB, PostgreSQL).
|
|
16
|
+
- **Sync and async api**: Highly compatible with modern and *legacy* frameworks.
|
|
17
|
+
- **No forced schema**: As long it is a list of json serializable objects.
|
|
18
|
+
- **Resumable memory**: Perfect for chat applications and REST APIs
|
|
19
|
+
- **Robust**: Get production-ready code with minimal effort.
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
## ⚙️ Installation
|
|
23
|
+
|
|
24
|
+
From pypi
|
|
25
|
+
```bash
|
|
26
|
+
pip install memx-ai
|
|
27
|
+
```
|
|
28
|
+
Or clone the repo and install it
|
|
29
|
+
```bash
|
|
30
|
+
pip install .
|
|
31
|
+
```
|
|
32
|
+
|
|
33
|
+
## 🚀 Quickstart
|
|
34
|
+
|
|
35
|
+
### OpenAI
|
|
36
|
+
Simple conversation with [OpenAI Python library](https://github.com/openai/openai-python)
|
|
37
|
+
```Python
|
|
38
|
+
# https://platform.openai.com/docs/guides/conversation-state?api-mode=responses
|
|
39
|
+
# tested on openai==2.6.1
|
|
40
|
+
|
|
41
|
+
from openai import OpenAI
|
|
42
|
+
from memx.memory.sqlite import SQLiteMemory
|
|
43
|
+
|
|
44
|
+
sqlite_uri = "sqlite+aiosqlite:///message-storage.db"
|
|
45
|
+
engine = SQLiteEngine(sqlite_uri, "memx-messages", start_up=True)
|
|
46
|
+
m1 = engine.create_session() # create a new session
|
|
47
|
+
|
|
48
|
+
client = OpenAI()
|
|
49
|
+
|
|
50
|
+
m1.sync.add([{"role": "user", "content": "tell me a good joke about programmers"}])
|
|
51
|
+
|
|
52
|
+
first_response = client.responses.create(
|
|
53
|
+
model="gpt-4o-mini", input=m1.sync.get(), store=False
|
|
54
|
+
)
|
|
55
|
+
|
|
56
|
+
print(first_response.output_text)
|
|
57
|
+
|
|
58
|
+
m1.sync.add(
|
|
59
|
+
[{"role": r.role, "content": r.content[0].text} for r in first_response.output]
|
|
60
|
+
)
|
|
61
|
+
|
|
62
|
+
m1.sync.add([{"role": "user", "content": "tell me another"}])
|
|
63
|
+
|
|
64
|
+
second_response = client.responses.create(
|
|
65
|
+
model="gpt-4o-mini", input=m1.sync.get(), store=False
|
|
66
|
+
)
|
|
67
|
+
|
|
68
|
+
m1.sync.add(
|
|
69
|
+
[{"role": r.role, "content": r.content[0].text} for r in second_response.output]
|
|
70
|
+
)
|
|
71
|
+
|
|
72
|
+
print(f"\n\n{second_response.output_text}")
|
|
73
|
+
|
|
74
|
+
print(m1.sync.get())
|
|
75
|
+
```
|
|
76
|
+
### Pydantic AI
|
|
77
|
+
Message history with async [Pydantic AI](https://ai.pydantic.dev/) + OpenAI
|
|
78
|
+
|
|
79
|
+
```Python
|
|
80
|
+
# Reference: https://ai.pydantic.dev/message-history/
|
|
81
|
+
|
|
82
|
+
import asyncio
|
|
83
|
+
|
|
84
|
+
import orjson
|
|
85
|
+
from pydantic_ai import Agent, ModelMessagesTypeAdapter
|
|
86
|
+
|
|
87
|
+
from memx.engine.sqlite import SQLiteEngine
|
|
88
|
+
|
|
89
|
+
agent = Agent("openai:gpt-4o-mini")
|
|
90
|
+
|
|
91
|
+
|
|
92
|
+
async def main():
|
|
93
|
+
sqlite_uri = "sqlite+aiosqlite:///message_store.db"
|
|
94
|
+
engine = SQLiteEngine(sqlite_uri, "memx-messages", start_up=True)
|
|
95
|
+
m1 = engine.create_session() # create a new session
|
|
96
|
+
|
|
97
|
+
result1 = await agent.run('Where does "hello world" come from?')
|
|
98
|
+
|
|
99
|
+
# it is your responsibility to add the messages as a list[dict]
|
|
100
|
+
messages = orjson.loads(result1.new_messages_json())
|
|
101
|
+
|
|
102
|
+
await m1.add(messages) # messages: list[dict] must be json serializable
|
|
103
|
+
|
|
104
|
+
session_id = m1.get_id()
|
|
105
|
+
print("Messages added with session_id: ", session_id)
|
|
106
|
+
|
|
107
|
+
# resume the conversation from 'another' memory
|
|
108
|
+
m2 = await engine.get_session(session_id)
|
|
109
|
+
old_messages = ModelMessagesTypeAdapter.validate_python(await m2.get())
|
|
110
|
+
|
|
111
|
+
print("Past messages:\n", old_messages)
|
|
112
|
+
|
|
113
|
+
result2 = await agent.run(
|
|
114
|
+
"Could you tell me more about the authors?", message_history=old_messages
|
|
115
|
+
)
|
|
116
|
+
print("\n\nContext aware result:\n", result2.output)
|
|
117
|
+
|
|
118
|
+
|
|
119
|
+
if __name__ == "__main__":
|
|
120
|
+
asyncio.run(main())
|
|
121
|
+
|
|
122
|
+
|
|
123
|
+
```
|
|
124
|
+
|
|
125
|
+
You can change the memory backend with minimal modifications. Same api to add and get messages.
|
|
126
|
+
```Python
|
|
127
|
+
from memx.memory.mongodb import MongoDBMemory
|
|
128
|
+
from memx.memory.postgres import PostgresMemory
|
|
129
|
+
from memx.memory.sqlite import SQLiteMemory
|
|
130
|
+
|
|
131
|
+
# SQLite backend
|
|
132
|
+
sqlite_uri = "sqlite+aiosqlite:///message_store.db"
|
|
133
|
+
e1 = SQLiteMemory(sqlite_uri, "memx-messages", start_up=True)
|
|
134
|
+
m1 = e1.create_session() # memory session ready to go
|
|
135
|
+
|
|
136
|
+
# PostgreSQL backend
|
|
137
|
+
pg_uri = "postgresql+psycopg://admin:1234@localhost:5433/test-database"
|
|
138
|
+
e2 = PostgresMemory(pg_uri, "memx-messages", start_up=True)
|
|
139
|
+
m2 = e2.create_session()
|
|
140
|
+
|
|
141
|
+
# MongoDB backend
|
|
142
|
+
mongodb_uri = "mongodb://admin:1234@localhost:27017"
|
|
143
|
+
e3 = MongoDBMemory(uri=mongodb_uri, database="memx-test", "memx-messages")
|
|
144
|
+
m3 = e3.create_session()
|
|
145
|
+
|
|
146
|
+
```
|
|
147
|
+
|
|
148
|
+
[More examples...](examples/)
|
|
149
|
+
|
|
150
|
+
## Tasks
|
|
151
|
+
- [x] Add mongodb backend
|
|
152
|
+
- [x] Add SQLite backend
|
|
153
|
+
- [x] Add Postgres backend
|
|
154
|
+
- [ ] Add redis backend
|
|
155
|
+
- [ ] Add tests
|
|
156
|
+
- [x] Publish on pypi
|
|
157
|
+
- [ ] Add full sync support
|
|
158
|
+
- [ ] Add docstrings
|
|
File without changes
|
|
@@ -0,0 +1,15 @@
|
|
|
1
|
+
from abc import ABC, abstractmethod
|
|
2
|
+
|
|
3
|
+
from memx.memory import BaseMemory
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
class BaseEngine(ABC):
|
|
7
|
+
@abstractmethod
|
|
8
|
+
def create_session(self) -> BaseMemory:
|
|
9
|
+
"""Create a memory session."""
|
|
10
|
+
pass
|
|
11
|
+
|
|
12
|
+
@abstractmethod
|
|
13
|
+
def get_session(self, session_id: str) -> BaseMemory | None:
|
|
14
|
+
"""Get a memory session from backend."""
|
|
15
|
+
pass
|
|
@@ -0,0 +1,34 @@
|
|
|
1
|
+
from datetime import UTC, datetime
|
|
2
|
+
from uuid import uuid4
|
|
3
|
+
|
|
4
|
+
from pymongo import AsyncMongoClient, MongoClient
|
|
5
|
+
from pymongo.server_api import ServerApi
|
|
6
|
+
|
|
7
|
+
from memx.engine import BaseEngine
|
|
8
|
+
from memx.memory.mongodb import MongoDBMemory
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
class MongoDBEngine(BaseEngine):
|
|
12
|
+
def __init__(self, uri: str, database: str, collection: str):
|
|
13
|
+
"""MongoDB memory engine."""
|
|
14
|
+
|
|
15
|
+
self.client = MongoClient(uri)
|
|
16
|
+
self.async_client = AsyncMongoClient(
|
|
17
|
+
uri,
|
|
18
|
+
server_api=ServerApi(version="1", strict=True, deprecation_errors=True),
|
|
19
|
+
)
|
|
20
|
+
|
|
21
|
+
self.db = self.client[database]
|
|
22
|
+
self.async_db = self.async_client.get_database(database)
|
|
23
|
+
|
|
24
|
+
self.sync_collection = self.db[collection]
|
|
25
|
+
self.async_collection = self.async_db[collection]
|
|
26
|
+
|
|
27
|
+
def create_session(self) -> MongoDBMemory:
|
|
28
|
+
return MongoDBMemory(self.async_collection, self.sync_collection)
|
|
29
|
+
|
|
30
|
+
async def get_session(self, id: str) -> MongoDBMemory | None:
|
|
31
|
+
result = await self.async_collection.find_one({"session_id": id})
|
|
32
|
+
|
|
33
|
+
if result:
|
|
34
|
+
return MongoDBMemory(self.async_collection, self.sync_collection, id)
|
|
@@ -0,0 +1,122 @@
|
|
|
1
|
+
from datetime import UTC, datetime
|
|
2
|
+
from textwrap import dedent
|
|
3
|
+
from uuid import uuid4
|
|
4
|
+
|
|
5
|
+
import orjson
|
|
6
|
+
from sqlalchemy import create_engine, text
|
|
7
|
+
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
|
8
|
+
from sqlalchemy.orm import Session, sessionmaker
|
|
9
|
+
|
|
10
|
+
from memx.engine import BaseEngine
|
|
11
|
+
from memx.engine.config import SQLEngineConfig
|
|
12
|
+
from memx.memory.postgres import PostgresMemory
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
class PostgresEngine(BaseEngine):
|
|
16
|
+
def __init__(self, uri: str, table: str, schema: str = "public", start_up: bool = None):
|
|
17
|
+
"""."""
|
|
18
|
+
|
|
19
|
+
self.table_name = f'"{table.strip()}"'
|
|
20
|
+
self.init_queries()
|
|
21
|
+
|
|
22
|
+
driver, _ = uri.split(":", 1)
|
|
23
|
+
if driver.strip() != "postgresql+psycopg":
|
|
24
|
+
raise ValueError("For the moment, only 'postgresql+psycopg' driver is supported")
|
|
25
|
+
|
|
26
|
+
common_args = {
|
|
27
|
+
"autocommit": False,
|
|
28
|
+
"autoflush": False,
|
|
29
|
+
"expire_on_commit": True,
|
|
30
|
+
}
|
|
31
|
+
|
|
32
|
+
self.async_engine = create_async_engine(
|
|
33
|
+
uri,
|
|
34
|
+
connect_args={"options": f"-csearch_path={schema}"},
|
|
35
|
+
)
|
|
36
|
+
self.AsyncSession = async_sessionmaker(
|
|
37
|
+
**common_args,
|
|
38
|
+
bind=self.async_engine,
|
|
39
|
+
class_=AsyncSession,
|
|
40
|
+
) # type: ignore
|
|
41
|
+
|
|
42
|
+
self.sync_engine = create_engine(
|
|
43
|
+
uri,
|
|
44
|
+
connect_args={"options": f"-csearch_path={schema}"},
|
|
45
|
+
)
|
|
46
|
+
self.SyncSession = sessionmaker(
|
|
47
|
+
**common_args,
|
|
48
|
+
bind=self.sync_engine,
|
|
49
|
+
class_=Session,
|
|
50
|
+
) # type: ignore
|
|
51
|
+
|
|
52
|
+
if start_up:
|
|
53
|
+
self.start_up() # blocking operation
|
|
54
|
+
|
|
55
|
+
def create_session(self) -> PostgresMemory:
|
|
56
|
+
"""Get or create a memory session."""
|
|
57
|
+
|
|
58
|
+
engine_config = SQLEngineConfig(
|
|
59
|
+
table=self.table_name,
|
|
60
|
+
add_query=self.add_sql,
|
|
61
|
+
get_query=self.get_sql,
|
|
62
|
+
)
|
|
63
|
+
return PostgresMemory(self.AsyncSession, self.SyncSession, engine_config)
|
|
64
|
+
|
|
65
|
+
async def get_session(self, id: str) -> PostgresMemory | None:
|
|
66
|
+
"""Get a memory session."""
|
|
67
|
+
|
|
68
|
+
async with self.AsyncSession() as session:
|
|
69
|
+
result = (
|
|
70
|
+
await session.execute(
|
|
71
|
+
text(self.get_session_sql),
|
|
72
|
+
{"session_id": id},
|
|
73
|
+
)
|
|
74
|
+
).first()
|
|
75
|
+
|
|
76
|
+
if result[0] == 1: # type: ignore
|
|
77
|
+
engine_config = SQLEngineConfig(
|
|
78
|
+
table=self.table_name,
|
|
79
|
+
add_query=self.add_sql,
|
|
80
|
+
get_query=self.get_sql,
|
|
81
|
+
)
|
|
82
|
+
return PostgresMemory(self.AsyncSession, self.SyncSession, engine_config, id)
|
|
83
|
+
|
|
84
|
+
def init_queries(self):
|
|
85
|
+
"""."""
|
|
86
|
+
|
|
87
|
+
self.table_sql = dedent(f"""
|
|
88
|
+
CREATE TABLE IF NOT EXISTS {self.table_name} (
|
|
89
|
+
session_id uuid PRIMARY KEY,
|
|
90
|
+
message JSONB,
|
|
91
|
+
created_at TIMESTAMP WITH TIME ZONE DEFAULT (NOW() AT TIME ZONE 'UTC'),
|
|
92
|
+
updated_at TIMESTAMP WITH TIME ZONE DEFAULT (NOW() AT TIME ZONE 'UTC')
|
|
93
|
+
);
|
|
94
|
+
""")
|
|
95
|
+
|
|
96
|
+
self.add_sql = dedent(f"""
|
|
97
|
+
INSERT INTO {self.table_name} (session_id, message, updated_at)
|
|
98
|
+
VALUES (:session_id, cast(:message as jsonb), :updated_at)
|
|
99
|
+
ON CONFLICT (session_id)
|
|
100
|
+
DO UPDATE SET
|
|
101
|
+
message = COALESCE({self.table_name}.message, '[]'::jsonb) || EXCLUDED.message,
|
|
102
|
+
updated_at = EXCLUDED.updated_at;
|
|
103
|
+
""")
|
|
104
|
+
|
|
105
|
+
self.get_sql = dedent(f"""
|
|
106
|
+
SELECT * FROM {self.table_name}
|
|
107
|
+
WHERE session_id = :session_id;
|
|
108
|
+
""")
|
|
109
|
+
|
|
110
|
+
self.get_session_sql = dedent(f"""
|
|
111
|
+
SELECT EXISTS(
|
|
112
|
+
SELECT 1 FROM {self.table_name}
|
|
113
|
+
WHERE session_id=:session_id
|
|
114
|
+
) as r;
|
|
115
|
+
""")
|
|
116
|
+
|
|
117
|
+
def start_up(self):
|
|
118
|
+
"""Create the table if it doesn't exist."""
|
|
119
|
+
|
|
120
|
+
with self.SyncSession() as session:
|
|
121
|
+
session.execute(text(self.table_sql))
|
|
122
|
+
session.commit()
|
|
@@ -0,0 +1,106 @@
|
|
|
1
|
+
from textwrap import dedent
|
|
2
|
+
|
|
3
|
+
from sqlalchemy import create_engine, text
|
|
4
|
+
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
|
5
|
+
from sqlalchemy.orm import Session, sessionmaker
|
|
6
|
+
|
|
7
|
+
from memx.engine import BaseEngine
|
|
8
|
+
from memx.engine.config import SQLEngineConfig
|
|
9
|
+
from memx.memory.sqlite import SQLiteMemory
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
class SQLiteEngine(BaseEngine):
|
|
13
|
+
def __init__(self, uri: str, table: str, start_up: bool = False):
|
|
14
|
+
"""SQLite memory engine."""
|
|
15
|
+
|
|
16
|
+
self.table_name = f"'{table.strip()}'"
|
|
17
|
+
self.init_queries()
|
|
18
|
+
|
|
19
|
+
self.async_engine = create_async_engine(uri, echo=False, future=True)
|
|
20
|
+
self.AsyncSession = async_sessionmaker(
|
|
21
|
+
bind=self.async_engine,
|
|
22
|
+
expire_on_commit=False,
|
|
23
|
+
class_=AsyncSession,
|
|
24
|
+
)
|
|
25
|
+
|
|
26
|
+
drivers, others = uri.split(":", 1) # type: ignore[reportUnusedVariable]
|
|
27
|
+
self.sync_engine = create_engine(
|
|
28
|
+
f"sqlite:{others}",
|
|
29
|
+
echo=False,
|
|
30
|
+
connect_args={"check_same_thread": True},
|
|
31
|
+
)
|
|
32
|
+
|
|
33
|
+
self.SyncSession = sessionmaker(
|
|
34
|
+
autocommit=False,
|
|
35
|
+
autoflush=False,
|
|
36
|
+
bind=self.sync_engine,
|
|
37
|
+
class_=Session,
|
|
38
|
+
)
|
|
39
|
+
|
|
40
|
+
if start_up:
|
|
41
|
+
self.start_up() # blocking operation
|
|
42
|
+
|
|
43
|
+
def create_session(self) -> SQLiteMemory:
|
|
44
|
+
"""Create a local memory session."""
|
|
45
|
+
|
|
46
|
+
engine_config = SQLEngineConfig(
|
|
47
|
+
table=self.table_name,
|
|
48
|
+
add_query=self.add_sql,
|
|
49
|
+
get_query=self.get_sql,
|
|
50
|
+
)
|
|
51
|
+
return SQLiteMemory(self.AsyncSession, self.SyncSession, engine_config)
|
|
52
|
+
|
|
53
|
+
async def get_session(self, id: str) -> SQLiteMemory | None:
|
|
54
|
+
"""Get a memory session."""
|
|
55
|
+
|
|
56
|
+
async with self.AsyncSession() as session:
|
|
57
|
+
result = (
|
|
58
|
+
await session.execute(
|
|
59
|
+
text(self.get_session_sql),
|
|
60
|
+
{"session_id": id},
|
|
61
|
+
)
|
|
62
|
+
).first()
|
|
63
|
+
|
|
64
|
+
if result[0] == 1: # type: ignore
|
|
65
|
+
engine_config = SQLEngineConfig(
|
|
66
|
+
table=self.table_name,
|
|
67
|
+
add_query=self.add_sql,
|
|
68
|
+
get_query=self.get_sql,
|
|
69
|
+
)
|
|
70
|
+
return SQLiteMemory(self.AsyncSession, self.SyncSession, engine_config, id)
|
|
71
|
+
|
|
72
|
+
def init_queries(self):
|
|
73
|
+
"""."""
|
|
74
|
+
|
|
75
|
+
self.table_sql = dedent(f"""
|
|
76
|
+
CREATE TABLE IF NOT EXISTS {self.table_name} (
|
|
77
|
+
session_id TEXT,
|
|
78
|
+
message JSONB,
|
|
79
|
+
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
|
80
|
+
);
|
|
81
|
+
CREATE INDEX IF NOT EXISTS session_index ON {self.table_name} (session_id);
|
|
82
|
+
""")
|
|
83
|
+
|
|
84
|
+
self.add_sql = dedent(f"""
|
|
85
|
+
INSERT INTO {self.table_name} (session_id, message, created_at)
|
|
86
|
+
VALUES (:session_id, :message, :created_at);
|
|
87
|
+
""")
|
|
88
|
+
|
|
89
|
+
self.get_sql = dedent(f"""
|
|
90
|
+
SELECT message FROM {self.table_name}
|
|
91
|
+
WHERE session_id = :session_id
|
|
92
|
+
ORDER BY created_at ASC;
|
|
93
|
+
""")
|
|
94
|
+
|
|
95
|
+
self.get_session_sql = dedent(f"""
|
|
96
|
+
SELECT EXISTS(
|
|
97
|
+
SELECT 1 FROM {self.table_name}
|
|
98
|
+
WHERE session_id=:session_id
|
|
99
|
+
) as r;
|
|
100
|
+
""")
|
|
101
|
+
|
|
102
|
+
def start_up(self):
|
|
103
|
+
"""Create the table if it doesn't exist."""
|
|
104
|
+
|
|
105
|
+
with self.sync_engine.begin() as conn:
|
|
106
|
+
conn.connection.executescript(self.table_sql)
|
|
@@ -0,0 +1,14 @@
|
|
|
1
|
+
from abc import ABC, abstractmethod
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
class BaseMemory(ABC):
|
|
5
|
+
@abstractmethod
|
|
6
|
+
def add(self, messages: list[dict]):
|
|
7
|
+
pass
|
|
8
|
+
|
|
9
|
+
@abstractmethod
|
|
10
|
+
def get(self) -> list[dict]:
|
|
11
|
+
pass
|
|
12
|
+
|
|
13
|
+
def get_id(self) -> str:
|
|
14
|
+
return self._session_id # type: ignore
|
|
@@ -0,0 +1,66 @@
|
|
|
1
|
+
from datetime import UTC, datetime
|
|
2
|
+
from uuid import uuid4
|
|
3
|
+
|
|
4
|
+
from pymongo.asynchronous.collection import AsyncCollection
|
|
5
|
+
from pymongo.collection import Collection
|
|
6
|
+
|
|
7
|
+
from memx.memory import BaseMemory
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
class MongoDBMemory(BaseMemory):
|
|
11
|
+
def __init__(
|
|
12
|
+
self,
|
|
13
|
+
async_collection: AsyncCollection,
|
|
14
|
+
sync_collection: Collection,
|
|
15
|
+
session_id: str = None,
|
|
16
|
+
):
|
|
17
|
+
self.async_collection = async_collection
|
|
18
|
+
self.sync_collection = sync_collection
|
|
19
|
+
|
|
20
|
+
self.sync = _sync(self) # to group sync methods
|
|
21
|
+
|
|
22
|
+
if session_id:
|
|
23
|
+
self._session_id = session_id
|
|
24
|
+
else:
|
|
25
|
+
self._session_id = str(uuid4())
|
|
26
|
+
|
|
27
|
+
async def add(self, messages: list[dict]):
|
|
28
|
+
ts_now = datetime.now(UTC)
|
|
29
|
+
|
|
30
|
+
await self.async_collection.find_one_and_update(
|
|
31
|
+
{"session_id": self._session_id},
|
|
32
|
+
{
|
|
33
|
+
"$push": {"messages": {"$each": messages}},
|
|
34
|
+
"$setOnInsert": {"created_at": ts_now},
|
|
35
|
+
"$set": {"updated_at": ts_now},
|
|
36
|
+
},
|
|
37
|
+
upsert=True,
|
|
38
|
+
)
|
|
39
|
+
|
|
40
|
+
async def get(self) -> list[dict]:
|
|
41
|
+
doc = await self.async_collection.find_one({"session_id": self._session_id})
|
|
42
|
+
|
|
43
|
+
return (doc or {}).get("messages", [])
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
class _sync(BaseMemory):
|
|
47
|
+
def __init__(self, parent: "MongoDBMemory"):
|
|
48
|
+
self.pm = parent # parent memory (?)
|
|
49
|
+
|
|
50
|
+
def add(self, messages: list[dict]):
|
|
51
|
+
ts_now = datetime.now(UTC)
|
|
52
|
+
|
|
53
|
+
self.pm.sync_collection.find_one_and_update(
|
|
54
|
+
{"session_id": self.pm._session_id},
|
|
55
|
+
{
|
|
56
|
+
"$push": {"messages": {"$each": messages}},
|
|
57
|
+
"$setOnInsert": {"created_at": ts_now},
|
|
58
|
+
"$set": {"updated_at": ts_now},
|
|
59
|
+
},
|
|
60
|
+
upsert=True,
|
|
61
|
+
)
|
|
62
|
+
|
|
63
|
+
def get(self) -> list[dict]:
|
|
64
|
+
doc = self.pm.sync_collection.find_one({"session_id": self.pm._session_id})
|
|
65
|
+
|
|
66
|
+
return (doc or {}).get("messages", [])
|
|
@@ -0,0 +1,97 @@
|
|
|
1
|
+
from datetime import UTC, datetime
|
|
2
|
+
from uuid import uuid4
|
|
3
|
+
|
|
4
|
+
import orjson
|
|
5
|
+
from sqlalchemy import text
|
|
6
|
+
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
|
7
|
+
from sqlalchemy.orm import Session, sessionmaker
|
|
8
|
+
|
|
9
|
+
from memx.engine.config import SQLEngineConfig
|
|
10
|
+
from memx.memory import BaseMemory
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
class PostgresMemory(BaseMemory):
|
|
14
|
+
def __init__(
|
|
15
|
+
self,
|
|
16
|
+
async_session_maker: async_sessionmaker[AsyncSession], # type: ignore
|
|
17
|
+
sync_session_maker: sessionmaker[Session],
|
|
18
|
+
engine_config: SQLEngineConfig,
|
|
19
|
+
session_id: str = None,
|
|
20
|
+
):
|
|
21
|
+
self.AsyncSession = async_session_maker
|
|
22
|
+
self.SyncSession = sync_session_maker
|
|
23
|
+
|
|
24
|
+
self.engine_config = engine_config
|
|
25
|
+
|
|
26
|
+
self.sync = _sync(self) # to group sync methods
|
|
27
|
+
|
|
28
|
+
if session_id:
|
|
29
|
+
self._session_id = session_id
|
|
30
|
+
else:
|
|
31
|
+
self._session_id = str(uuid4())
|
|
32
|
+
|
|
33
|
+
async def add(self, messages: list[dict]):
|
|
34
|
+
# TODO: refactor this with sqlite
|
|
35
|
+
await self._pre_add()
|
|
36
|
+
|
|
37
|
+
ts_now = datetime.now(UTC)
|
|
38
|
+
data = {
|
|
39
|
+
"session_id": self._session_id,
|
|
40
|
+
"message": orjson.dumps(messages).decode("utf-8"),
|
|
41
|
+
"updated_at": ts_now,
|
|
42
|
+
}
|
|
43
|
+
|
|
44
|
+
async with self.AsyncSession() as session:
|
|
45
|
+
await session.execute(text(self.engine_config.add_query), data)
|
|
46
|
+
await session.commit()
|
|
47
|
+
|
|
48
|
+
async def get(self) -> list[dict]:
|
|
49
|
+
async with self.AsyncSession() as session:
|
|
50
|
+
result = await session.execute(
|
|
51
|
+
text(self.engine_config.get_query),
|
|
52
|
+
{"session_id": self._session_id},
|
|
53
|
+
)
|
|
54
|
+
|
|
55
|
+
result = result.first()
|
|
56
|
+
result = getattr(result, "message", [])
|
|
57
|
+
|
|
58
|
+
return result
|
|
59
|
+
|
|
60
|
+
async def _pre_add(self):
|
|
61
|
+
pass
|
|
62
|
+
|
|
63
|
+
|
|
64
|
+
class _sync(BaseMemory):
|
|
65
|
+
def __init__(self, parent: "PostgresMemory"):
|
|
66
|
+
self.pm = parent # parent memory (?)
|
|
67
|
+
|
|
68
|
+
def add(self, messages: list[dict]):
|
|
69
|
+
# TODO: refactor this with sqlite
|
|
70
|
+
|
|
71
|
+
self._pre_add()
|
|
72
|
+
|
|
73
|
+
ts_now = datetime.now(UTC)
|
|
74
|
+
data = {
|
|
75
|
+
"session_id": self.pm._session_id,
|
|
76
|
+
"message": orjson.dumps(messages).decode("utf-8"),
|
|
77
|
+
"updated_at": ts_now,
|
|
78
|
+
}
|
|
79
|
+
|
|
80
|
+
with self.pm.SyncSession() as session:
|
|
81
|
+
session.execute(text(self.pm.engine_config.add_query), data)
|
|
82
|
+
session.commit()
|
|
83
|
+
|
|
84
|
+
def get(self) -> list[dict]:
|
|
85
|
+
with self.pm.SyncSession() as session:
|
|
86
|
+
result = session.execute(
|
|
87
|
+
text(self.pm.engine_config.get_query),
|
|
88
|
+
{"session_id": self.pm._session_id},
|
|
89
|
+
)
|
|
90
|
+
|
|
91
|
+
result = result.first()
|
|
92
|
+
result = getattr(result, "message", [])
|
|
93
|
+
|
|
94
|
+
return result
|
|
95
|
+
|
|
96
|
+
def _pre_add(self):
|
|
97
|
+
pass
|
|
File without changes
|
|
@@ -0,0 +1,101 @@
|
|
|
1
|
+
import asyncio
|
|
2
|
+
import pickle
|
|
3
|
+
from pathlib import Path
|
|
4
|
+
from uuid import uuid4
|
|
5
|
+
|
|
6
|
+
import aiofiles
|
|
7
|
+
|
|
8
|
+
from memx.memory import BaseMemory
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
class DiskMemory(BaseMemory):
|
|
12
|
+
def __init__(self, session_id: str = None, dir: str = None):
|
|
13
|
+
file_id = session_id if session_id else str(uuid4())
|
|
14
|
+
file_dir = Path(dir) if dir else Path.home() / ".memx"
|
|
15
|
+
|
|
16
|
+
self.file_path = Path(file_dir) / f"{file_id}.pkl"
|
|
17
|
+
file_dir.mkdir(parents=True, exist_ok=True)
|
|
18
|
+
|
|
19
|
+
if not self.file_path.is_file():
|
|
20
|
+
with open(self.file_path, "wb") as f:
|
|
21
|
+
pickle.dump([], f, protocol=pickle.HIGHEST_PROTOCOL)
|
|
22
|
+
|
|
23
|
+
self._session_id = file_id
|
|
24
|
+
|
|
25
|
+
self.sync = _sync(self) # to group sync methods
|
|
26
|
+
|
|
27
|
+
async def add(self, messages: list[dict]):
|
|
28
|
+
# read the file
|
|
29
|
+
async with aiofiles.open(self.file_path, "rb") as f:
|
|
30
|
+
pickled_data = await f.read()
|
|
31
|
+
|
|
32
|
+
stored_messages: list[dict] = await asyncio.to_thread(
|
|
33
|
+
pickle.loads, pickled_data
|
|
34
|
+
)
|
|
35
|
+
|
|
36
|
+
# extend the messages
|
|
37
|
+
stored_messages.extend(messages)
|
|
38
|
+
|
|
39
|
+
# write the file
|
|
40
|
+
pickled_data = await asyncio.to_thread(
|
|
41
|
+
pickle.dumps, stored_messages, protocol=pickle.HIGHEST_PROTOCOL
|
|
42
|
+
)
|
|
43
|
+
|
|
44
|
+
async with aiofiles.open(self.file_path, "wb") as f:
|
|
45
|
+
await f.write(pickled_data)
|
|
46
|
+
|
|
47
|
+
async def get(self) -> list[dict]:
|
|
48
|
+
async with aiofiles.open(self.file_path, "rb") as f:
|
|
49
|
+
pickled_data = await f.read()
|
|
50
|
+
|
|
51
|
+
stored_messages: list[dict] = await asyncio.to_thread(
|
|
52
|
+
pickle.loads, pickled_data
|
|
53
|
+
)
|
|
54
|
+
|
|
55
|
+
return stored_messages
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
class _sync(BaseMemory):
|
|
59
|
+
def __init__(self, parent: "DiskMemory"):
|
|
60
|
+
self.pm = parent # parent memory (?)
|
|
61
|
+
|
|
62
|
+
def add(self, messages: list[dict]):
|
|
63
|
+
with open(self.pm.file_path, "rb") as f:
|
|
64
|
+
stored_messages: list[dict] = pickle.load(f)
|
|
65
|
+
|
|
66
|
+
stored_messages.extend(messages)
|
|
67
|
+
|
|
68
|
+
with open(self.pm.file_path, "wb") as f:
|
|
69
|
+
pickle.dump(stored_messages, f, protocol=pickle.HIGHEST_PROTOCOL)
|
|
70
|
+
|
|
71
|
+
def get(self) -> list[dict]:
|
|
72
|
+
with open(self.pm.file_path, "rb") as f:
|
|
73
|
+
stored_messages: list[dict] = pickle.load(f)
|
|
74
|
+
|
|
75
|
+
return stored_messages
|
|
76
|
+
|
|
77
|
+
|
|
78
|
+
class InMemory(BaseMemory):
|
|
79
|
+
# TODO: add .sync just for consistency
|
|
80
|
+
def __init__(self, session_id: str = None):
|
|
81
|
+
global __memx_in_memory__
|
|
82
|
+
if "__memx_in_memory__" not in globals().keys():
|
|
83
|
+
__memx_in_memory__ = {}
|
|
84
|
+
|
|
85
|
+
if session_id:
|
|
86
|
+
self._messages = __memx_in_memory__[session_id]
|
|
87
|
+
_session_id = session_id
|
|
88
|
+
else:
|
|
89
|
+
_session_id = str(uuid4())
|
|
90
|
+
__memx_in_memory__[_session_id] = []
|
|
91
|
+
self._messages = __memx_in_memory__[_session_id]
|
|
92
|
+
|
|
93
|
+
self._session_id = _session_id
|
|
94
|
+
|
|
95
|
+
def add(self, messages: list[dict]):
|
|
96
|
+
self._messages.extend(messages)
|
|
97
|
+
|
|
98
|
+
def get(
|
|
99
|
+
self,
|
|
100
|
+
) -> list[str]:
|
|
101
|
+
return self._messages
|
|
@@ -0,0 +1,105 @@
|
|
|
1
|
+
from datetime import UTC, datetime
|
|
2
|
+
from uuid import uuid4
|
|
3
|
+
|
|
4
|
+
import orjson
|
|
5
|
+
from sqlalchemy import Result, text
|
|
6
|
+
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
|
7
|
+
from sqlalchemy.orm import Session, sessionmaker
|
|
8
|
+
|
|
9
|
+
from memx.engine.config import SQLEngineConfig
|
|
10
|
+
from memx.memory import BaseMemory
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
class SQLiteMemory(BaseMemory):
|
|
14
|
+
def __init__(
|
|
15
|
+
self,
|
|
16
|
+
async_session_maker: async_sessionmaker[AsyncSession], # type: ignore
|
|
17
|
+
sync_session_maker: sessionmaker[Session],
|
|
18
|
+
engine_config: SQLEngineConfig,
|
|
19
|
+
session_id: str = None,
|
|
20
|
+
):
|
|
21
|
+
self.AsyncSession = async_session_maker
|
|
22
|
+
self.SyncSession = sync_session_maker
|
|
23
|
+
|
|
24
|
+
self.engine_config = engine_config
|
|
25
|
+
|
|
26
|
+
self.sync = _sync(self) # to group sync methods
|
|
27
|
+
|
|
28
|
+
if session_id:
|
|
29
|
+
self._session_id = session_id
|
|
30
|
+
else:
|
|
31
|
+
self._session_id = str(uuid4())
|
|
32
|
+
|
|
33
|
+
async def add(self, messages: list[dict]):
|
|
34
|
+
await self._pre_add()
|
|
35
|
+
|
|
36
|
+
data = self._format_messages(messages)
|
|
37
|
+
|
|
38
|
+
async with self.AsyncSession() as session:
|
|
39
|
+
await session.execute(text(self.engine_config.add_query), data)
|
|
40
|
+
await session.commit()
|
|
41
|
+
|
|
42
|
+
async def get(self) -> list[dict]:
|
|
43
|
+
async with self.AsyncSession() as session:
|
|
44
|
+
result = await session.execute(
|
|
45
|
+
text(self.engine_config.get_query),
|
|
46
|
+
{"session_id": self._session_id},
|
|
47
|
+
)
|
|
48
|
+
|
|
49
|
+
messages = _merge_messages(result)
|
|
50
|
+
|
|
51
|
+
return messages
|
|
52
|
+
|
|
53
|
+
async def _pre_add(self):
|
|
54
|
+
pass
|
|
55
|
+
|
|
56
|
+
def _format_messages(self, messages: list[dict]) -> dict:
|
|
57
|
+
ts_now = datetime.now(UTC)
|
|
58
|
+
data = {
|
|
59
|
+
"session_id": self._session_id,
|
|
60
|
+
"message": orjson.dumps(messages).decode("utf-8"),
|
|
61
|
+
"created_at": ts_now,
|
|
62
|
+
}
|
|
63
|
+
|
|
64
|
+
return data
|
|
65
|
+
|
|
66
|
+
|
|
67
|
+
class _sync(BaseMemory):
|
|
68
|
+
def __init__(self, parent: "SQLiteMemory"):
|
|
69
|
+
self.pm = parent # parent memory (?)
|
|
70
|
+
|
|
71
|
+
def add(self, messages: list[dict]):
|
|
72
|
+
self._pre_add()
|
|
73
|
+
|
|
74
|
+
data = self.pm._format_messages(messages)
|
|
75
|
+
|
|
76
|
+
with self.pm.SyncSession() as session:
|
|
77
|
+
session.execute(text(self.pm.engine_config.add_query), data)
|
|
78
|
+
session.commit()
|
|
79
|
+
|
|
80
|
+
def get(self) -> list[dict]:
|
|
81
|
+
with self.pm.SyncSession() as session:
|
|
82
|
+
result = session.execute(
|
|
83
|
+
text(self.pm.engine_config.get_query),
|
|
84
|
+
{"session_id": self.pm._session_id},
|
|
85
|
+
)
|
|
86
|
+
|
|
87
|
+
messages = _merge_messages(result)
|
|
88
|
+
|
|
89
|
+
return messages
|
|
90
|
+
|
|
91
|
+
def _pre_add(self):
|
|
92
|
+
pass
|
|
93
|
+
|
|
94
|
+
|
|
95
|
+
def _merge_messages(msg_result: Result) -> list[dict]:
|
|
96
|
+
"""."""
|
|
97
|
+
|
|
98
|
+
# list.extend is the fastest approach
|
|
99
|
+
result = [dict(row._mapping) for row in msg_result.fetchall()]
|
|
100
|
+
messages = []
|
|
101
|
+
|
|
102
|
+
for r in result:
|
|
103
|
+
messages.extend(orjson.loads(r["message"]))
|
|
104
|
+
|
|
105
|
+
return messages
|
|
@@ -0,0 +1,81 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "memx-ai"
|
|
7
|
+
version = "0.1.8"
|
|
8
|
+
description = "Lightweight and extensible memory layer for LLMs"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
license = "MIT"
|
|
11
|
+
license-files = ["LICENSE"]
|
|
12
|
+
requires-python = ">=3.11"
|
|
13
|
+
authors = [
|
|
14
|
+
{ name = "Pablo Galilea", email = "pgalilea@me.com" },
|
|
15
|
+
]
|
|
16
|
+
dependencies = [
|
|
17
|
+
"aiofiles>=25.1.0",
|
|
18
|
+
"orjson>=3.11.4",
|
|
19
|
+
]
|
|
20
|
+
|
|
21
|
+
[dependency-groups]
|
|
22
|
+
dev = [
|
|
23
|
+
"pytest>=8.4.2",
|
|
24
|
+
"ruff>=0.14.2"
|
|
25
|
+
]
|
|
26
|
+
|
|
27
|
+
[project.optional-dependencies]
|
|
28
|
+
all = [
|
|
29
|
+
"aiosqlite>=0.21.0",
|
|
30
|
+
"psycopg>=3.2.12",
|
|
31
|
+
"pymongo>=4.15.3",
|
|
32
|
+
"sqlalchemy[asyncio, postgresql-psycopg]>=2.0.44"
|
|
33
|
+
]
|
|
34
|
+
|
|
35
|
+
sqlite = [
|
|
36
|
+
"sqlalchemy[asyncio]>=2.0.44",
|
|
37
|
+
"aiosqlite>=0.21.0"
|
|
38
|
+
]
|
|
39
|
+
|
|
40
|
+
postgres = [
|
|
41
|
+
"sqlalchemy[postgresql-psycopg, asyncio]>=2.0.44",
|
|
42
|
+
"psycopg>=3.2.12"
|
|
43
|
+
]
|
|
44
|
+
|
|
45
|
+
mongodb = [
|
|
46
|
+
"pymongo>=4.15.3"
|
|
47
|
+
]
|
|
48
|
+
|
|
49
|
+
[tool.hatch.build.targets.wheel]
|
|
50
|
+
packages = ["memx"]
|
|
51
|
+
only-include = ["memx"]
|
|
52
|
+
|
|
53
|
+
[tool.hatch.build]
|
|
54
|
+
include = [
|
|
55
|
+
"memx/**/*.py",
|
|
56
|
+
"tests"
|
|
57
|
+
]
|
|
58
|
+
|
|
59
|
+
[tool.hatch.build.targets.wheel.shared-data]
|
|
60
|
+
"README.md" = "README.md"
|
|
61
|
+
|
|
62
|
+
[tool.ruff]
|
|
63
|
+
select = [
|
|
64
|
+
"E", # pycodestyle errors
|
|
65
|
+
"W", # pycodestyle warnings
|
|
66
|
+
"F", # pyflakes
|
|
67
|
+
"I", # isort
|
|
68
|
+
"B", # flake8-bugbear
|
|
69
|
+
"C4", # flake8-comprehensions
|
|
70
|
+
"UP", # pyupgrade
|
|
71
|
+
]
|
|
72
|
+
line-length = 99
|
|
73
|
+
fixable = ["ALL"]
|
|
74
|
+
unfixable = ["F401", "F403"]
|
|
75
|
+
ignore = ["F401", "F403"]
|
|
76
|
+
|
|
77
|
+
[tool.pyright]
|
|
78
|
+
pythonVersion = "3.11"
|
|
79
|
+
typeCheckingMode = "basic"
|
|
80
|
+
reportInvalidTypeForm = "none"
|
|
81
|
+
reportArgumentType = false
|
|
File without changes
|
|
File without changes
|
|
File without changes
|