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.
@@ -0,0 +1,14 @@
1
+ # Python-generated files
2
+ __pycache__/
3
+ *.py[oc]
4
+ build/
5
+ dist/
6
+ wheels/
7
+ *.egg-info
8
+ uv.lock
9
+
10
+ # Virtual environments
11
+ .venv
12
+
13
+ tmp/
14
+ .DS_Store
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
@@ -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,8 @@
1
+ from dataclasses import dataclass
2
+
3
+
4
+ @dataclass
5
+ class SQLEngineConfig:
6
+ table: str
7
+ add_query: int
8
+ get_query: float
@@ -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