axio-context-sqlite 0.5.1__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.
- axio_context_sqlite-0.5.1/.gitignore +12 -0
- axio_context_sqlite-0.5.1/PKG-INFO +145 -0
- axio_context_sqlite-0.5.1/README.md +132 -0
- axio_context_sqlite-0.5.1/pyproject.toml +49 -0
- axio_context_sqlite-0.5.1/src/axio_context_sqlite/__init__.py +5 -0
- axio_context_sqlite-0.5.1/src/axio_context_sqlite/py.typed +0 -0
- axio_context_sqlite-0.5.1/src/axio_context_sqlite/store.py +216 -0
- axio_context_sqlite-0.5.1/tests/test_store.py +164 -0
|
@@ -0,0 +1,145 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: axio-context-sqlite
|
|
3
|
+
Version: 0.5.1
|
|
4
|
+
Summary: SQLite-backed context store for Axio
|
|
5
|
+
Project-URL: Homepage, https://github.com/axio-agent/monorepo
|
|
6
|
+
Project-URL: Repository, https://github.com/axio-agent/monorepo
|
|
7
|
+
License: MIT
|
|
8
|
+
Keywords: agent,ai,context,llm,sqlite,storage
|
|
9
|
+
Requires-Python: >=3.12
|
|
10
|
+
Requires-Dist: aiosqlite>=0.20
|
|
11
|
+
Requires-Dist: axio
|
|
12
|
+
Description-Content-Type: text/markdown
|
|
13
|
+
|
|
14
|
+
# axio-context-sqlite
|
|
15
|
+
|
|
16
|
+
[](https://pypi.org/project/axio-context-sqlite/)
|
|
17
|
+
[](https://pypi.org/project/axio-context-sqlite/)
|
|
18
|
+
[](LICENSE)
|
|
19
|
+
|
|
20
|
+
SQLite-backed persistent context store for [axio](https://github.com/axio-agent/monorepo).
|
|
21
|
+
|
|
22
|
+
## Installation
|
|
23
|
+
|
|
24
|
+
```bash
|
|
25
|
+
pip install axio-context-sqlite
|
|
26
|
+
```
|
|
27
|
+
|
|
28
|
+
## Usage
|
|
29
|
+
|
|
30
|
+
Open a connection with `connect()`, then create a `SQLiteContextStore` bound to a
|
|
31
|
+
session. The caller owns the connection and is responsible for closing it.
|
|
32
|
+
|
|
33
|
+
<!-- name: test_readme_usage -->
|
|
34
|
+
```python
|
|
35
|
+
import asyncio
|
|
36
|
+
import tempfile
|
|
37
|
+
import pathlib
|
|
38
|
+
from axio_context_sqlite import connect, SQLiteContextStore
|
|
39
|
+
from axio.messages import Message
|
|
40
|
+
from axio.blocks import TextBlock
|
|
41
|
+
|
|
42
|
+
async def main() -> None:
|
|
43
|
+
conn = await connect(pathlib.Path(tempfile.mkdtemp()) / "chat.db")
|
|
44
|
+
try:
|
|
45
|
+
store = SQLiteContextStore(conn, session_id="my-session")
|
|
46
|
+
await store.append(Message(role="user", content=[TextBlock(text="Hello")]))
|
|
47
|
+
history = await store.get_history()
|
|
48
|
+
assert len(history) == 1
|
|
49
|
+
finally:
|
|
50
|
+
await conn.close()
|
|
51
|
+
|
|
52
|
+
asyncio.run(main())
|
|
53
|
+
```
|
|
54
|
+
|
|
55
|
+
`SQLiteContextStore` implements the `axio.context.ContextStore` ABC and persists
|
|
56
|
+
conversation history across process restarts. Multiple sessions can coexist in
|
|
57
|
+
the same database file, isolated by `session_id` and `project`.
|
|
58
|
+
|
|
59
|
+
### Agent integration
|
|
60
|
+
|
|
61
|
+
<!--
|
|
62
|
+
name: test_readme_agent
|
|
63
|
+
```python
|
|
64
|
+
from axio.testing import StubTransport, make_text_response
|
|
65
|
+
transport = StubTransport([make_text_response("Hi!")])
|
|
66
|
+
```
|
|
67
|
+
-->
|
|
68
|
+
<!-- name: test_readme_agent -->
|
|
69
|
+
```python
|
|
70
|
+
import asyncio
|
|
71
|
+
import tempfile
|
|
72
|
+
import pathlib
|
|
73
|
+
from axio.agent import Agent
|
|
74
|
+
from axio_context_sqlite import connect, SQLiteContextStore
|
|
75
|
+
|
|
76
|
+
async def main() -> None:
|
|
77
|
+
conn = await connect(pathlib.Path(tempfile.mkdtemp()) / "chat.db")
|
|
78
|
+
try:
|
|
79
|
+
ctx = SQLiteContextStore(conn, session_id="main")
|
|
80
|
+
agent = Agent(system="You are helpful.", tools=[], transport=transport)
|
|
81
|
+
result = await agent.run("Hello!", ctx)
|
|
82
|
+
assert result == "Hi!"
|
|
83
|
+
finally:
|
|
84
|
+
await conn.close()
|
|
85
|
+
|
|
86
|
+
asyncio.run(main())
|
|
87
|
+
```
|
|
88
|
+
|
|
89
|
+
### Listing sessions
|
|
90
|
+
|
|
91
|
+
<!-- name: test_readme_list_sessions -->
|
|
92
|
+
```python
|
|
93
|
+
import asyncio
|
|
94
|
+
import tempfile
|
|
95
|
+
import pathlib
|
|
96
|
+
from axio_context_sqlite import connect, SQLiteContextStore
|
|
97
|
+
from axio.messages import Message
|
|
98
|
+
from axio.blocks import TextBlock
|
|
99
|
+
|
|
100
|
+
async def main() -> None:
|
|
101
|
+
conn = await connect(pathlib.Path(tempfile.mkdtemp()) / "chat.db")
|
|
102
|
+
try:
|
|
103
|
+
store = SQLiteContextStore(conn, session_id="main", project="/myproject")
|
|
104
|
+
await store.append(Message(role="user", content=[TextBlock(text="hi")]))
|
|
105
|
+
sessions = await store.list_sessions()
|
|
106
|
+
for s in sessions:
|
|
107
|
+
print(s.session_id, s.preview, s.message_count)
|
|
108
|
+
assert len(sessions) == 1
|
|
109
|
+
finally:
|
|
110
|
+
await conn.close()
|
|
111
|
+
|
|
112
|
+
asyncio.run(main())
|
|
113
|
+
```
|
|
114
|
+
|
|
115
|
+
### Forking
|
|
116
|
+
|
|
117
|
+
`fork()` copies the current session's messages into a new session — useful for
|
|
118
|
+
branching conversations without affecting the original:
|
|
119
|
+
|
|
120
|
+
<!-- name: test_readme_fork -->
|
|
121
|
+
```python
|
|
122
|
+
import asyncio
|
|
123
|
+
import tempfile
|
|
124
|
+
import pathlib
|
|
125
|
+
from axio_context_sqlite import connect, SQLiteContextStore
|
|
126
|
+
from axio.messages import Message
|
|
127
|
+
from axio.blocks import TextBlock
|
|
128
|
+
|
|
129
|
+
async def main() -> None:
|
|
130
|
+
conn = await connect(pathlib.Path(tempfile.mkdtemp()) / "chat.db")
|
|
131
|
+
try:
|
|
132
|
+
store = SQLiteContextStore(conn, session_id="main")
|
|
133
|
+
await store.append(Message(role="user", content=[TextBlock(text="original")]))
|
|
134
|
+
branch = await store.fork()
|
|
135
|
+
assert branch.session_id != store.session_id
|
|
136
|
+
assert len(await branch.get_history()) == 1
|
|
137
|
+
finally:
|
|
138
|
+
await conn.close()
|
|
139
|
+
|
|
140
|
+
asyncio.run(main())
|
|
141
|
+
```
|
|
142
|
+
|
|
143
|
+
## License
|
|
144
|
+
|
|
145
|
+
MIT
|
|
@@ -0,0 +1,132 @@
|
|
|
1
|
+
# axio-context-sqlite
|
|
2
|
+
|
|
3
|
+
[](https://pypi.org/project/axio-context-sqlite/)
|
|
4
|
+
[](https://pypi.org/project/axio-context-sqlite/)
|
|
5
|
+
[](LICENSE)
|
|
6
|
+
|
|
7
|
+
SQLite-backed persistent context store for [axio](https://github.com/axio-agent/monorepo).
|
|
8
|
+
|
|
9
|
+
## Installation
|
|
10
|
+
|
|
11
|
+
```bash
|
|
12
|
+
pip install axio-context-sqlite
|
|
13
|
+
```
|
|
14
|
+
|
|
15
|
+
## Usage
|
|
16
|
+
|
|
17
|
+
Open a connection with `connect()`, then create a `SQLiteContextStore` bound to a
|
|
18
|
+
session. The caller owns the connection and is responsible for closing it.
|
|
19
|
+
|
|
20
|
+
<!-- name: test_readme_usage -->
|
|
21
|
+
```python
|
|
22
|
+
import asyncio
|
|
23
|
+
import tempfile
|
|
24
|
+
import pathlib
|
|
25
|
+
from axio_context_sqlite import connect, SQLiteContextStore
|
|
26
|
+
from axio.messages import Message
|
|
27
|
+
from axio.blocks import TextBlock
|
|
28
|
+
|
|
29
|
+
async def main() -> None:
|
|
30
|
+
conn = await connect(pathlib.Path(tempfile.mkdtemp()) / "chat.db")
|
|
31
|
+
try:
|
|
32
|
+
store = SQLiteContextStore(conn, session_id="my-session")
|
|
33
|
+
await store.append(Message(role="user", content=[TextBlock(text="Hello")]))
|
|
34
|
+
history = await store.get_history()
|
|
35
|
+
assert len(history) == 1
|
|
36
|
+
finally:
|
|
37
|
+
await conn.close()
|
|
38
|
+
|
|
39
|
+
asyncio.run(main())
|
|
40
|
+
```
|
|
41
|
+
|
|
42
|
+
`SQLiteContextStore` implements the `axio.context.ContextStore` ABC and persists
|
|
43
|
+
conversation history across process restarts. Multiple sessions can coexist in
|
|
44
|
+
the same database file, isolated by `session_id` and `project`.
|
|
45
|
+
|
|
46
|
+
### Agent integration
|
|
47
|
+
|
|
48
|
+
<!--
|
|
49
|
+
name: test_readme_agent
|
|
50
|
+
```python
|
|
51
|
+
from axio.testing import StubTransport, make_text_response
|
|
52
|
+
transport = StubTransport([make_text_response("Hi!")])
|
|
53
|
+
```
|
|
54
|
+
-->
|
|
55
|
+
<!-- name: test_readme_agent -->
|
|
56
|
+
```python
|
|
57
|
+
import asyncio
|
|
58
|
+
import tempfile
|
|
59
|
+
import pathlib
|
|
60
|
+
from axio.agent import Agent
|
|
61
|
+
from axio_context_sqlite import connect, SQLiteContextStore
|
|
62
|
+
|
|
63
|
+
async def main() -> None:
|
|
64
|
+
conn = await connect(pathlib.Path(tempfile.mkdtemp()) / "chat.db")
|
|
65
|
+
try:
|
|
66
|
+
ctx = SQLiteContextStore(conn, session_id="main")
|
|
67
|
+
agent = Agent(system="You are helpful.", tools=[], transport=transport)
|
|
68
|
+
result = await agent.run("Hello!", ctx)
|
|
69
|
+
assert result == "Hi!"
|
|
70
|
+
finally:
|
|
71
|
+
await conn.close()
|
|
72
|
+
|
|
73
|
+
asyncio.run(main())
|
|
74
|
+
```
|
|
75
|
+
|
|
76
|
+
### Listing sessions
|
|
77
|
+
|
|
78
|
+
<!-- name: test_readme_list_sessions -->
|
|
79
|
+
```python
|
|
80
|
+
import asyncio
|
|
81
|
+
import tempfile
|
|
82
|
+
import pathlib
|
|
83
|
+
from axio_context_sqlite import connect, SQLiteContextStore
|
|
84
|
+
from axio.messages import Message
|
|
85
|
+
from axio.blocks import TextBlock
|
|
86
|
+
|
|
87
|
+
async def main() -> None:
|
|
88
|
+
conn = await connect(pathlib.Path(tempfile.mkdtemp()) / "chat.db")
|
|
89
|
+
try:
|
|
90
|
+
store = SQLiteContextStore(conn, session_id="main", project="/myproject")
|
|
91
|
+
await store.append(Message(role="user", content=[TextBlock(text="hi")]))
|
|
92
|
+
sessions = await store.list_sessions()
|
|
93
|
+
for s in sessions:
|
|
94
|
+
print(s.session_id, s.preview, s.message_count)
|
|
95
|
+
assert len(sessions) == 1
|
|
96
|
+
finally:
|
|
97
|
+
await conn.close()
|
|
98
|
+
|
|
99
|
+
asyncio.run(main())
|
|
100
|
+
```
|
|
101
|
+
|
|
102
|
+
### Forking
|
|
103
|
+
|
|
104
|
+
`fork()` copies the current session's messages into a new session — useful for
|
|
105
|
+
branching conversations without affecting the original:
|
|
106
|
+
|
|
107
|
+
<!-- name: test_readme_fork -->
|
|
108
|
+
```python
|
|
109
|
+
import asyncio
|
|
110
|
+
import tempfile
|
|
111
|
+
import pathlib
|
|
112
|
+
from axio_context_sqlite import connect, SQLiteContextStore
|
|
113
|
+
from axio.messages import Message
|
|
114
|
+
from axio.blocks import TextBlock
|
|
115
|
+
|
|
116
|
+
async def main() -> None:
|
|
117
|
+
conn = await connect(pathlib.Path(tempfile.mkdtemp()) / "chat.db")
|
|
118
|
+
try:
|
|
119
|
+
store = SQLiteContextStore(conn, session_id="main")
|
|
120
|
+
await store.append(Message(role="user", content=[TextBlock(text="original")]))
|
|
121
|
+
branch = await store.fork()
|
|
122
|
+
assert branch.session_id != store.session_id
|
|
123
|
+
assert len(await branch.get_history()) == 1
|
|
124
|
+
finally:
|
|
125
|
+
await conn.close()
|
|
126
|
+
|
|
127
|
+
asyncio.run(main())
|
|
128
|
+
```
|
|
129
|
+
|
|
130
|
+
## License
|
|
131
|
+
|
|
132
|
+
MIT
|
|
@@ -0,0 +1,49 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "axio-context-sqlite"
|
|
3
|
+
version = "0.5.1"
|
|
4
|
+
description = "SQLite-backed context store for Axio"
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
requires-python = ">=3.12"
|
|
7
|
+
license = {text = "MIT"}
|
|
8
|
+
keywords = ["llm", "agent", "ai", "sqlite", "context", "storage"]
|
|
9
|
+
dependencies = ["axio", "aiosqlite>=0.20"]
|
|
10
|
+
|
|
11
|
+
[project.urls]
|
|
12
|
+
Homepage = "https://github.com/axio-agent/monorepo"
|
|
13
|
+
Repository = "https://github.com/axio-agent/monorepo"
|
|
14
|
+
|
|
15
|
+
[build-system]
|
|
16
|
+
requires = ["hatchling"]
|
|
17
|
+
build-backend = "hatchling.build"
|
|
18
|
+
|
|
19
|
+
[tool.hatch.build.targets.wheel]
|
|
20
|
+
packages = ["src/axio_context_sqlite"]
|
|
21
|
+
sources = ["src"]
|
|
22
|
+
|
|
23
|
+
[tool.pytest.ini_options]
|
|
24
|
+
asyncio_mode = "auto"
|
|
25
|
+
testpaths = ["tests", "README.md"]
|
|
26
|
+
|
|
27
|
+
[tool.ruff]
|
|
28
|
+
line-length = 119
|
|
29
|
+
target-version = "py312"
|
|
30
|
+
|
|
31
|
+
[tool.ruff.lint]
|
|
32
|
+
select = ["E", "F", "I", "UP"]
|
|
33
|
+
|
|
34
|
+
[tool.mypy]
|
|
35
|
+
strict = true
|
|
36
|
+
python_version = "3.12"
|
|
37
|
+
|
|
38
|
+
[dependency-groups]
|
|
39
|
+
dev = [
|
|
40
|
+
"pytest>=8",
|
|
41
|
+
"pytest-asyncio>=0.24",
|
|
42
|
+
"mypy>=1.14",
|
|
43
|
+
"ruff>=0.9",
|
|
44
|
+
"pytest-cov>=7.1.0",
|
|
45
|
+
"markdown-pytest>=0.6.0",
|
|
46
|
+
]
|
|
47
|
+
|
|
48
|
+
[tool.uv.sources]
|
|
49
|
+
axio = { workspace = true }
|
|
File without changes
|
|
@@ -0,0 +1,216 @@
|
|
|
1
|
+
"""SQLiteContextStore: persistent conversation storage backed by SQLite."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import asyncio
|
|
6
|
+
import base64
|
|
7
|
+
import gzip
|
|
8
|
+
import json
|
|
9
|
+
from pathlib import Path
|
|
10
|
+
from uuid import uuid4
|
|
11
|
+
|
|
12
|
+
import aiosqlite
|
|
13
|
+
from axio.context import ContextStore, SessionInfo
|
|
14
|
+
from axio.messages import Message
|
|
15
|
+
|
|
16
|
+
# Compress content payloads above this size (bytes of UTF-8 JSON).
|
|
17
|
+
COMPRESS_THRESHOLD = 512
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
def compress_payload(data: str) -> str:
|
|
21
|
+
raw = data.encode()
|
|
22
|
+
if len(raw) < COMPRESS_THRESHOLD:
|
|
23
|
+
return "plain:" + data
|
|
24
|
+
return "gzip:" + base64.b64encode(gzip.compress(raw, compresslevel=6)).decode()
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
def decompress_payload(data: str) -> str:
|
|
28
|
+
if data.startswith("gzip:"):
|
|
29
|
+
return gzip.decompress(base64.b64decode(data[5:])).decode()
|
|
30
|
+
if data.startswith("plain:"):
|
|
31
|
+
return data[6:]
|
|
32
|
+
# raw JSON
|
|
33
|
+
return data
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
async def connect(db_path: str | Path) -> aiosqlite.Connection:
|
|
37
|
+
"""Open (or create) a SQLite database and initialise the schema.
|
|
38
|
+
|
|
39
|
+
The caller is responsible for closing the returned connection.
|
|
40
|
+
"""
|
|
41
|
+
path = Path(db_path)
|
|
42
|
+
await asyncio.to_thread(path.parent.mkdir, parents=True, exist_ok=True)
|
|
43
|
+
conn = await aiosqlite.connect(str(path))
|
|
44
|
+
await conn.create_function("compress_payload", 1, compress_payload, deterministic=True)
|
|
45
|
+
await conn.create_function("decompress_payload", 1, decompress_payload, deterministic=True)
|
|
46
|
+
await conn.execute("PRAGMA journal_mode=WAL")
|
|
47
|
+
await conn.execute("PRAGMA busy_timeout=5000")
|
|
48
|
+
await conn.execute("PRAGMA synchronous=NORMAL")
|
|
49
|
+
await conn.execute(
|
|
50
|
+
"CREATE TABLE IF NOT EXISTS axio_context_messages ("
|
|
51
|
+
" id INTEGER PRIMARY KEY AUTOINCREMENT,"
|
|
52
|
+
" session_id TEXT NOT NULL,"
|
|
53
|
+
" project TEXT NOT NULL,"
|
|
54
|
+
" position INTEGER NOT NULL,"
|
|
55
|
+
" role TEXT NOT NULL,"
|
|
56
|
+
" content TEXT NOT NULL,"
|
|
57
|
+
" created_at TEXT NOT NULL DEFAULT (datetime('now')),"
|
|
58
|
+
" UNIQUE(session_id, position)"
|
|
59
|
+
")"
|
|
60
|
+
)
|
|
61
|
+
await conn.execute(
|
|
62
|
+
"CREATE INDEX IF NOT EXISTS idx_axio_context_messages_session ON axio_context_messages(session_id)"
|
|
63
|
+
)
|
|
64
|
+
await conn.execute(
|
|
65
|
+
"CREATE INDEX IF NOT EXISTS idx_axio_context_messages_project ON axio_context_messages(project)"
|
|
66
|
+
)
|
|
67
|
+
await conn.execute(
|
|
68
|
+
"CREATE TABLE IF NOT EXISTS axio_context_tokens ("
|
|
69
|
+
" session_id TEXT NOT NULL,"
|
|
70
|
+
" project TEXT NOT NULL,"
|
|
71
|
+
" input_tokens INTEGER NOT NULL DEFAULT 0,"
|
|
72
|
+
" output_tokens INTEGER NOT NULL DEFAULT 0,"
|
|
73
|
+
" PRIMARY KEY(session_id, project)"
|
|
74
|
+
")"
|
|
75
|
+
)
|
|
76
|
+
await conn.commit()
|
|
77
|
+
return conn
|
|
78
|
+
|
|
79
|
+
|
|
80
|
+
def _extract_preview(content_json: str, max_len: int = 80) -> str:
|
|
81
|
+
"""Extract text preview from serialized content JSON."""
|
|
82
|
+
try:
|
|
83
|
+
blocks = json.loads(content_json)
|
|
84
|
+
for b in blocks:
|
|
85
|
+
if b.get("type") == "text":
|
|
86
|
+
text: str = b["text"]
|
|
87
|
+
return text[:max_len] + ("..." if len(text) > max_len else "")
|
|
88
|
+
except (json.JSONDecodeError, KeyError):
|
|
89
|
+
pass
|
|
90
|
+
return "(no preview)"
|
|
91
|
+
|
|
92
|
+
|
|
93
|
+
class SQLiteContextStore(ContextStore):
|
|
94
|
+
"""Persistent conversation storage backed by SQLite.
|
|
95
|
+
|
|
96
|
+
The caller owns the connection and is responsible for closing it.
|
|
97
|
+
Use :func:`connect` to open a properly initialized connection.
|
|
98
|
+
"""
|
|
99
|
+
|
|
100
|
+
def __init__(
|
|
101
|
+
self,
|
|
102
|
+
conn: aiosqlite.Connection,
|
|
103
|
+
session_id: str,
|
|
104
|
+
project: str | None = None,
|
|
105
|
+
db_name: str = "axio_context",
|
|
106
|
+
) -> None:
|
|
107
|
+
self._conn = conn
|
|
108
|
+
self._db_name = db_name
|
|
109
|
+
self._session_id = session_id
|
|
110
|
+
self._project = project or str(Path.cwd().resolve())
|
|
111
|
+
|
|
112
|
+
@property
|
|
113
|
+
def session_id(self) -> str:
|
|
114
|
+
return self._session_id
|
|
115
|
+
|
|
116
|
+
async def append(self, message: Message) -> None:
|
|
117
|
+
content_json = json.dumps(message.to_dict()["content"])
|
|
118
|
+
await self._conn.execute(
|
|
119
|
+
"INSERT INTO axio_context_messages (session_id, project, position, role, content)"
|
|
120
|
+
"VALUES (?, ?, (SELECT COUNT(*) FROM axio_context_messages WHERE session_id = ?), ?, compress_payload(?))",
|
|
121
|
+
(self._session_id, self._project, self._session_id, message.role, content_json),
|
|
122
|
+
)
|
|
123
|
+
await self._conn.commit()
|
|
124
|
+
|
|
125
|
+
async def get_history(self) -> list[Message]:
|
|
126
|
+
async with self._conn.execute(
|
|
127
|
+
"SELECT role, decompress_payload(content) FROM axio_context_messages"
|
|
128
|
+
" WHERE session_id = ? ORDER BY position",
|
|
129
|
+
(self._session_id,),
|
|
130
|
+
) as cursor:
|
|
131
|
+
rows = await cursor.fetchall()
|
|
132
|
+
return [Message.from_dict({"role": role, "content": json.loads(content)}) for role, content in rows]
|
|
133
|
+
|
|
134
|
+
async def clear(self) -> None:
|
|
135
|
+
await self._conn.execute("DELETE FROM axio_context_messages WHERE session_id = ?", (self._session_id,))
|
|
136
|
+
await self._conn.execute(
|
|
137
|
+
"DELETE FROM axio_context_tokens WHERE session_id = ? AND project = ?",
|
|
138
|
+
(self._session_id, self._project),
|
|
139
|
+
)
|
|
140
|
+
await self._conn.commit()
|
|
141
|
+
|
|
142
|
+
async def fork(self) -> SQLiteContextStore:
|
|
143
|
+
new_id = uuid4().hex
|
|
144
|
+
await self._conn.execute(
|
|
145
|
+
"INSERT INTO axio_context_messages (session_id, project, position, role, content)"
|
|
146
|
+
"SELECT ?, project, position, role, content FROM axio_context_messages WHERE session_id = ?",
|
|
147
|
+
(new_id, self._session_id),
|
|
148
|
+
)
|
|
149
|
+
await self._conn.execute(
|
|
150
|
+
"INSERT OR IGNORE INTO axio_context_tokens (session_id, project, input_tokens, output_tokens) "
|
|
151
|
+
"SELECT ?, project, input_tokens, output_tokens FROM axio_context_tokens "
|
|
152
|
+
"WHERE session_id = ? AND project = ?",
|
|
153
|
+
(new_id, self._session_id, self._project),
|
|
154
|
+
)
|
|
155
|
+
await self._conn.commit()
|
|
156
|
+
return SQLiteContextStore(self._conn, new_id, self._project)
|
|
157
|
+
|
|
158
|
+
async def set_context_tokens(self, input_tokens: int, output_tokens: int) -> None:
|
|
159
|
+
await self._conn.execute(
|
|
160
|
+
"INSERT INTO axio_context_tokens (session_id, project, input_tokens, output_tokens)"
|
|
161
|
+
"VALUES (?, ?, ?, ?) "
|
|
162
|
+
"ON CONFLICT(session_id, project) DO UPDATE SET input_tokens=?, output_tokens=?",
|
|
163
|
+
(self._session_id, self._project, input_tokens, output_tokens, input_tokens, output_tokens),
|
|
164
|
+
)
|
|
165
|
+
await self._conn.commit()
|
|
166
|
+
|
|
167
|
+
async def add_context_tokens(self, input_tokens: int, output_tokens: int) -> None:
|
|
168
|
+
await self._conn.execute(
|
|
169
|
+
"INSERT INTO axio_context_tokens (session_id, project, input_tokens, output_tokens)"
|
|
170
|
+
"VALUES (?, ?, ?, ?) "
|
|
171
|
+
"ON CONFLICT(session_id, project) DO UPDATE "
|
|
172
|
+
"SET input_tokens = input_tokens + excluded.input_tokens, "
|
|
173
|
+
" output_tokens = output_tokens + excluded.output_tokens",
|
|
174
|
+
(self._session_id, self._project, input_tokens, output_tokens),
|
|
175
|
+
)
|
|
176
|
+
await self._conn.commit()
|
|
177
|
+
|
|
178
|
+
async def get_context_tokens(self) -> tuple[int, int]:
|
|
179
|
+
async with self._conn.execute(
|
|
180
|
+
"SELECT input_tokens, output_tokens FROM axio_context_tokens WHERE session_id = ? AND project = ?",
|
|
181
|
+
(self._session_id, self._project),
|
|
182
|
+
) as cursor:
|
|
183
|
+
row = await cursor.fetchone()
|
|
184
|
+
if row is None:
|
|
185
|
+
return 0, 0
|
|
186
|
+
return int(row[0]), int(row[1])
|
|
187
|
+
|
|
188
|
+
async def list_sessions(self) -> list[SessionInfo]:
|
|
189
|
+
"""List all sessions for a project, newest first."""
|
|
190
|
+
async with self._conn.execute(
|
|
191
|
+
"SELECT m.session_id, COUNT(*) as cnt, "
|
|
192
|
+
"(SELECT decompress_payload(content) FROM axio_context_messages WHERE session_id = m.session_id "
|
|
193
|
+
"AND role = 'user' ORDER BY position LIMIT 1) as first_content, "
|
|
194
|
+
"MIN(m.created_at) as created, "
|
|
195
|
+
"COALESCE(ct.input_tokens, 0), COALESCE(ct.output_tokens, 0) "
|
|
196
|
+
"FROM axio_context_messages m "
|
|
197
|
+
"LEFT JOIN axio_context_tokens ct ON ct.session_id = m.session_id AND ct.project = m.project "
|
|
198
|
+
"WHERE m.project = ? "
|
|
199
|
+
"GROUP BY m.session_id ORDER BY created DESC",
|
|
200
|
+
(self._project,),
|
|
201
|
+
) as cursor:
|
|
202
|
+
rows = await cursor.fetchall()
|
|
203
|
+
result: list[SessionInfo] = []
|
|
204
|
+
for session_id, count, first_content, created_at, in_tok, out_tok in rows:
|
|
205
|
+
preview = _extract_preview(first_content) if first_content else "(no preview)"
|
|
206
|
+
result.append(
|
|
207
|
+
SessionInfo(
|
|
208
|
+
session_id=session_id,
|
|
209
|
+
message_count=count,
|
|
210
|
+
preview=preview,
|
|
211
|
+
created_at=created_at,
|
|
212
|
+
input_tokens=int(in_tok),
|
|
213
|
+
output_tokens=int(out_tok),
|
|
214
|
+
)
|
|
215
|
+
)
|
|
216
|
+
return result
|
|
@@ -0,0 +1,164 @@
|
|
|
1
|
+
"""Tests for SQLiteContextStore."""
|
|
2
|
+
|
|
3
|
+
from collections.abc import AsyncGenerator
|
|
4
|
+
from pathlib import Path
|
|
5
|
+
|
|
6
|
+
import aiosqlite
|
|
7
|
+
import pytest
|
|
8
|
+
from axio.blocks import TextBlock
|
|
9
|
+
from axio.messages import Message
|
|
10
|
+
|
|
11
|
+
from axio_context_sqlite import SQLiteContextStore, connect
|
|
12
|
+
from axio_context_sqlite.store import COMPRESS_THRESHOLD, compress_payload, decompress_payload
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
def _msg(role: str, text: str) -> Message:
|
|
16
|
+
return Message(role=role, content=[TextBlock(text=text)]) # type: ignore[arg-type]
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
@pytest.fixture
|
|
20
|
+
def db_path(tmp_path: Path) -> Path:
|
|
21
|
+
return tmp_path / "test.db"
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
@pytest.fixture
|
|
25
|
+
async def conn(db_path: Path) -> AsyncGenerator[aiosqlite.Connection, None]:
|
|
26
|
+
c = await connect(db_path)
|
|
27
|
+
yield c
|
|
28
|
+
await c.close()
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
@pytest.fixture
|
|
32
|
+
async def store(conn: aiosqlite.Connection) -> SQLiteContextStore:
|
|
33
|
+
return SQLiteContextStore(conn, "session-1", "test-project")
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
async def test_append_and_get_history(store: SQLiteContextStore) -> None:
|
|
37
|
+
await store.append(_msg("user", "Hello"))
|
|
38
|
+
await store.append(_msg("assistant", "Hi!"))
|
|
39
|
+
history = await store.get_history()
|
|
40
|
+
assert len(history) == 2
|
|
41
|
+
assert history[0].role == "user"
|
|
42
|
+
assert history[1].role == "assistant"
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
async def test_clear(store: SQLiteContextStore) -> None:
|
|
46
|
+
await store.append(_msg("user", "Hello"))
|
|
47
|
+
await store.clear()
|
|
48
|
+
assert await store.get_history() == []
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
async def test_fork(store: SQLiteContextStore) -> None:
|
|
52
|
+
await store.append(_msg("user", "Hello"))
|
|
53
|
+
forked = await store.fork()
|
|
54
|
+
await forked.append(_msg("assistant", "Hi!"))
|
|
55
|
+
assert len(await store.get_history()) == 1
|
|
56
|
+
assert len(await forked.get_history()) == 2
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
async def test_set_get_context_tokens(store: SQLiteContextStore) -> None:
|
|
60
|
+
await store.set_context_tokens(100, 50)
|
|
61
|
+
inp, out = await store.get_context_tokens()
|
|
62
|
+
assert inp == 100
|
|
63
|
+
assert out == 50
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
async def test_add_context_tokens(store: SQLiteContextStore) -> None:
|
|
67
|
+
await store.set_context_tokens(100, 50)
|
|
68
|
+
await store.add_context_tokens(20, 10)
|
|
69
|
+
inp, out = await store.get_context_tokens()
|
|
70
|
+
assert inp == 120
|
|
71
|
+
assert out == 60
|
|
72
|
+
|
|
73
|
+
|
|
74
|
+
async def test_list_sessions(db_path: Path) -> None:
|
|
75
|
+
c = await connect(db_path)
|
|
76
|
+
try:
|
|
77
|
+
s1 = SQLiteContextStore(c, "sess-a", "proj")
|
|
78
|
+
s2 = SQLiteContextStore(c, "sess-b", "proj")
|
|
79
|
+
await s1.append(_msg("user", "First session"))
|
|
80
|
+
await s2.append(_msg("user", "Second session"))
|
|
81
|
+
sessions = await s1.list_sessions()
|
|
82
|
+
assert len(sessions) == 2
|
|
83
|
+
ids = {s.session_id for s in sessions}
|
|
84
|
+
assert ids == {"sess-a", "sess-b"}
|
|
85
|
+
previews = {s.preview for s in sessions}
|
|
86
|
+
assert "First session" in previews
|
|
87
|
+
assert "Second session" in previews
|
|
88
|
+
finally:
|
|
89
|
+
await c.close()
|
|
90
|
+
|
|
91
|
+
|
|
92
|
+
class TestCompressPayload:
|
|
93
|
+
def test_small_stays_plain(self) -> None:
|
|
94
|
+
data = "hello"
|
|
95
|
+
result = compress_payload(data)
|
|
96
|
+
assert result.startswith("plain:")
|
|
97
|
+
assert decompress_payload(result) == data
|
|
98
|
+
|
|
99
|
+
def test_large_gets_compressed(self) -> None:
|
|
100
|
+
data = "x" * COMPRESS_THRESHOLD
|
|
101
|
+
result = compress_payload(data)
|
|
102
|
+
assert result.startswith("gzip:")
|
|
103
|
+
assert decompress_payload(result) == data
|
|
104
|
+
|
|
105
|
+
def test_legacy_no_prefix(self) -> None:
|
|
106
|
+
assert decompress_payload('{"foo": 1}') == '{"foo": 1}'
|
|
107
|
+
|
|
108
|
+
def test_roundtrip(self) -> None:
|
|
109
|
+
data = "[" + '{"type":"text","text":"a"},' * 50 + "]"
|
|
110
|
+
assert decompress_payload(compress_payload(data)) == data
|
|
111
|
+
|
|
112
|
+
|
|
113
|
+
async def test_large_message_compressed_on_disk(db_path: Path) -> None:
|
|
114
|
+
"""Large content is stored compressed; get_history still returns original."""
|
|
115
|
+
big_text = "word " * 200 # well above threshold
|
|
116
|
+
|
|
117
|
+
c = await connect(db_path)
|
|
118
|
+
try:
|
|
119
|
+
s = SQLiteContextStore(c, "big-session", "proj")
|
|
120
|
+
await s.append(_msg("user", big_text))
|
|
121
|
+
finally:
|
|
122
|
+
await c.close()
|
|
123
|
+
|
|
124
|
+
# Inspect raw bytes on disk — should start with gzip:
|
|
125
|
+
async with aiosqlite.connect(str(db_path)) as raw:
|
|
126
|
+
async with raw.execute("SELECT content FROM axio_context_messages") as cur:
|
|
127
|
+
row = await cur.fetchone()
|
|
128
|
+
assert row is not None
|
|
129
|
+
assert row[0].startswith("gzip:")
|
|
130
|
+
|
|
131
|
+
# But get_history transparently decompresses
|
|
132
|
+
c2 = await connect(db_path)
|
|
133
|
+
try:
|
|
134
|
+
s2 = SQLiteContextStore(c2, "big-session", "proj")
|
|
135
|
+
history = await s2.get_history()
|
|
136
|
+
finally:
|
|
137
|
+
await c2.close()
|
|
138
|
+
assert len(history) == 1
|
|
139
|
+
block = history[0].content[0]
|
|
140
|
+
assert isinstance(block, TextBlock)
|
|
141
|
+
assert block.text == big_text
|
|
142
|
+
|
|
143
|
+
|
|
144
|
+
async def test_close_and_reopen(db_path: Path) -> None:
|
|
145
|
+
"""Data persists after close and reopen."""
|
|
146
|
+
c = await connect(db_path)
|
|
147
|
+
try:
|
|
148
|
+
s = SQLiteContextStore(c, "persist-session", "proj")
|
|
149
|
+
await s.append(_msg("user", "Persistent message"))
|
|
150
|
+
await s.set_context_tokens(42, 7)
|
|
151
|
+
finally:
|
|
152
|
+
await c.close()
|
|
153
|
+
|
|
154
|
+
c2 = await connect(db_path)
|
|
155
|
+
try:
|
|
156
|
+
s2 = SQLiteContextStore(c2, "persist-session", "proj")
|
|
157
|
+
history = await s2.get_history()
|
|
158
|
+
assert len(history) == 1
|
|
159
|
+
assert history[0].role == "user"
|
|
160
|
+
inp, out = await s2.get_context_tokens()
|
|
161
|
+
assert inp == 42
|
|
162
|
+
assert out == 7
|
|
163
|
+
finally:
|
|
164
|
+
await c2.close()
|