xtr-ai-orm-conversation-store 6.6.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 xterr
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 all
13
+ 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 THE
21
+ SOFTWARE.
@@ -0,0 +1,95 @@
1
+ Metadata-Version: 2.4
2
+ Name: xtr-ai-orm-conversation-store
3
+ Version: 6.6.0
4
+ Summary: Agent conversations stored in a relational database through xtr-orm, for xtr-ai.
5
+ Keywords: ai,llm,agent,pydantic-ai,conversation,sqlalchemy,database
6
+ Author: Xterr
7
+ Author-email: Xterr <me@xterr.dev>
8
+ License-Expression: MIT
9
+ License-File: LICENSE
10
+ Classifier: Development Status :: 3 - Alpha
11
+ Classifier: Framework :: AsyncIO
12
+ Classifier: Intended Audience :: Developers
13
+ Classifier: Programming Language :: Python :: 3.11
14
+ Classifier: Programming Language :: Python :: 3.12
15
+ Classifier: Programming Language :: Python :: 3.13
16
+ Classifier: Programming Language :: Python :: 3.14
17
+ Classifier: Typing :: Typed
18
+ Requires-Dist: pydantic-ai-slim>=2.66,<3
19
+ Requires-Dist: sqlalchemy[asyncio]>=2.0.41
20
+ Requires-Dist: typing-extensions>=4.12
21
+ Requires-Dist: xtr-ai>=6.0,<7
22
+ Requires-Dist: xtr-orm>=6.0,<7
23
+ Requires-Python: >=3.11
24
+ Description-Content-Type: text/markdown
25
+
26
+ # xtr-ai-orm-conversation-store
27
+
28
+ Agent conversations kept in two tables on an [xtr-orm](../xtr-orm) connection, for
29
+ [xtr-ai](../xtr-ai).
30
+
31
+ ## Install
32
+
33
+ ```sh
34
+ uv add "xtr-ai[orm]"
35
+ ```
36
+
37
+ The `orm` extra installs this package.
38
+
39
+ ## Use
40
+
41
+ `OrmConversationStore` is a pydantic-ai `ConversationStore`. Give it the connection's primary
42
+ engine:
43
+
44
+ ```python
45
+ from pydantic_ai.conversation import Conversation
46
+
47
+ from xtr_ai_orm_conversation_store import OrmConversationStore
48
+
49
+ store = OrmConversationStore(engine)
50
+ conversation = await store.load(conversation_id) or Conversation(conversation_id=conversation_id)
51
+ result = await agent.run(prompt, conversation=conversation)
52
+ await store.save(result)
53
+ ```
54
+
55
+ In an application, configure it in xtr-ai instead, as
56
+ `conversation_stores={"default": OrmConversationStoreConfig(connection="default")}`, and inject
57
+ `ConversationStore`. See xtr-ai's [Use in an application](../xtr-ai#use-in-an-application).
58
+
59
+ ## Tables
60
+
61
+ | Table | Holds |
62
+ |---|---|
63
+ | `ai_conversation` | One row per conversation: where its history ends, usage, pending tool calls, scope, branches, fork origin, revision |
64
+ | `ai_conversation_entry` | One row per message, pointing at the one before it; shared by a conversation and its forks |
65
+
66
+ `table_prefix` changes `ai_`. `ConversationMappingDriver` adds the tables to a connection's
67
+ definitions, so `orm:migrations:diff` writes them. xtr-ai registers it when the store is
68
+ configured.
69
+
70
+ ## Behaviour
71
+
72
+ - **Own transactions.** Every call runs in a short transaction of its own and commits before it
73
+ returns. A store never joins the caller's unit of work.
74
+ - A long model call holds no database connection.
75
+ - Rolling back the caller's work never undoes a saved turn.
76
+ - A conversation cannot be saved atomically with the application's own rows.
77
+ - **Conflicts.** A record and the messages it adds are written in one transaction. A save over a
78
+ newer revision raises `ConversationConflict` and leaves neither behind.
79
+ - **History.** A conversation's history is read in one recursive query.
80
+ - **Deleting.** Deleting a conversation deletes the messages only it reaches. Messages a fork,
81
+ or its origin, still reaches stay until the last of them is deleted.
82
+ - A fork saved while its origin is being deleted can lose the messages it shares with it,
83
+ except on SQLite, where writes take turns. Don't fork a conversation you are deleting.
84
+ - **Databases.** It runs on PostgreSQL, MySQL 8 and SQLite, through their async drivers.
85
+
86
+ ## Development
87
+
88
+ ```sh
89
+ uv sync
90
+ uv run ruff check && uv run ruff format --check && uv run basedpyright && uv run ty check && uv run pytest
91
+ ```
92
+
93
+ ## License
94
+
95
+ MIT
@@ -0,0 +1,70 @@
1
+ # xtr-ai-orm-conversation-store
2
+
3
+ Agent conversations kept in two tables on an [xtr-orm](../xtr-orm) connection, for
4
+ [xtr-ai](../xtr-ai).
5
+
6
+ ## Install
7
+
8
+ ```sh
9
+ uv add "xtr-ai[orm]"
10
+ ```
11
+
12
+ The `orm` extra installs this package.
13
+
14
+ ## Use
15
+
16
+ `OrmConversationStore` is a pydantic-ai `ConversationStore`. Give it the connection's primary
17
+ engine:
18
+
19
+ ```python
20
+ from pydantic_ai.conversation import Conversation
21
+
22
+ from xtr_ai_orm_conversation_store import OrmConversationStore
23
+
24
+ store = OrmConversationStore(engine)
25
+ conversation = await store.load(conversation_id) or Conversation(conversation_id=conversation_id)
26
+ result = await agent.run(prompt, conversation=conversation)
27
+ await store.save(result)
28
+ ```
29
+
30
+ In an application, configure it in xtr-ai instead, as
31
+ `conversation_stores={"default": OrmConversationStoreConfig(connection="default")}`, and inject
32
+ `ConversationStore`. See xtr-ai's [Use in an application](../xtr-ai#use-in-an-application).
33
+
34
+ ## Tables
35
+
36
+ | Table | Holds |
37
+ |---|---|
38
+ | `ai_conversation` | One row per conversation: where its history ends, usage, pending tool calls, scope, branches, fork origin, revision |
39
+ | `ai_conversation_entry` | One row per message, pointing at the one before it; shared by a conversation and its forks |
40
+
41
+ `table_prefix` changes `ai_`. `ConversationMappingDriver` adds the tables to a connection's
42
+ definitions, so `orm:migrations:diff` writes them. xtr-ai registers it when the store is
43
+ configured.
44
+
45
+ ## Behaviour
46
+
47
+ - **Own transactions.** Every call runs in a short transaction of its own and commits before it
48
+ returns. A store never joins the caller's unit of work.
49
+ - A long model call holds no database connection.
50
+ - Rolling back the caller's work never undoes a saved turn.
51
+ - A conversation cannot be saved atomically with the application's own rows.
52
+ - **Conflicts.** A record and the messages it adds are written in one transaction. A save over a
53
+ newer revision raises `ConversationConflict` and leaves neither behind.
54
+ - **History.** A conversation's history is read in one recursive query.
55
+ - **Deleting.** Deleting a conversation deletes the messages only it reaches. Messages a fork,
56
+ or its origin, still reaches stay until the last of them is deleted.
57
+ - A fork saved while its origin is being deleted can lose the messages it shares with it,
58
+ except on SQLite, where writes take turns. Don't fork a conversation you are deleting.
59
+ - **Databases.** It runs on PostgreSQL, MySQL 8 and SQLite, through their async drivers.
60
+
61
+ ## Development
62
+
63
+ ```sh
64
+ uv sync
65
+ uv run ruff check && uv run ruff format --check && uv run basedpyright && uv run ty check && uv run pytest
66
+ ```
67
+
68
+ ## License
69
+
70
+ MIT
@@ -0,0 +1,160 @@
1
+ [project]
2
+ name = "xtr-ai-orm-conversation-store"
3
+ version = "6.6.0"
4
+ description = "Agent conversations stored in a relational database through xtr-orm, for xtr-ai."
5
+ readme = "README.md"
6
+ requires-python = ">=3.11"
7
+ license = "MIT"
8
+ license-files = ["LICENSE"]
9
+ keywords = [
10
+ "ai",
11
+ "llm",
12
+ "agent",
13
+ "pydantic-ai",
14
+ "conversation",
15
+ "sqlalchemy",
16
+ "database",
17
+ ]
18
+ classifiers = [
19
+ "Development Status :: 3 - Alpha",
20
+ "Framework :: AsyncIO",
21
+ "Intended Audience :: Developers",
22
+ "Programming Language :: Python :: 3.11",
23
+ "Programming Language :: Python :: 3.12",
24
+ "Programming Language :: Python :: 3.13",
25
+ "Programming Language :: Python :: 3.14",
26
+ "Typing :: Typed",
27
+ ]
28
+ dependencies = [
29
+ "pydantic-ai-slim>=2.66,<3",
30
+ "sqlalchemy[asyncio]>=2.0.41",
31
+ "typing-extensions>=4.12",
32
+ "xtr-ai>=6.0,<7",
33
+ "xtr-orm>=6.0,<7",
34
+ ]
35
+
36
+ [[project.authors]]
37
+ name = "Xterr"
38
+ email = "me@xterr.dev"
39
+
40
+ [dependency-groups]
41
+ dev = [
42
+ "basedpyright>=1.21",
43
+ "ruff>=0.8",
44
+ "pytest>=8",
45
+ "pytest-cov>=5",
46
+ "anyio>=4.0",
47
+ "ty>=0.0.83",
48
+ "aiosqlite>=0.20",
49
+ ]
50
+
51
+ [build-system]
52
+ requires = ["uv_build>=0.9.18,<0.10.0"]
53
+ build-backend = "uv_build"
54
+
55
+ [tool.basedpyright]
56
+ typeCheckingMode = "all"
57
+ pythonVersion = "3.11"
58
+ pythonPlatform = "All"
59
+ include = [
60
+ "src",
61
+ "tests",
62
+ ]
63
+ exclude = [
64
+ "**/__pycache__",
65
+ "**/.venv",
66
+ "**/build",
67
+ "**/dist",
68
+ ".tmp",
69
+ ]
70
+ reportUnusedCallResult = "warning"
71
+ reportUnnecessaryTypeIgnoreComment = "error"
72
+ reportUnusedVariable = "error"
73
+ reportMissingParameterType = "error"
74
+ reportPrivateUsage = "error"
75
+ reportImplicitStringConcatenation = "none"
76
+ reportExplicitAny = "none"
77
+ reportAny = "none"
78
+
79
+ [[tool.basedpyright.executionEnvironments]]
80
+ root = "tests"
81
+ extraPaths = ["."]
82
+ reportPrivateUsage = "none"
83
+
84
+ [tool.ruff]
85
+ target-version = "py311"
86
+ line-length = 100
87
+ src = [
88
+ "src",
89
+ "tests",
90
+ ]
91
+
92
+ [tool.ruff.lint]
93
+ select = ["ALL"]
94
+ ignore = [
95
+ "COM812",
96
+ "ISC001",
97
+ "D203",
98
+ "D213",
99
+ "CPY001",
100
+ "FBT001",
101
+ "FBT002",
102
+ "TD002",
103
+ "TD003",
104
+ "FIX002",
105
+ "TRY003",
106
+ "EM101",
107
+ "EM102",
108
+ "ANN401",
109
+ ]
110
+ fixable = ["ALL"]
111
+ unfixable = []
112
+
113
+ [tool.ruff.lint.per-file-ignores]
114
+ "tests/**/*.py" = [
115
+ "S101",
116
+ "ARG",
117
+ "PLR2004",
118
+ "SLF001",
119
+ "D",
120
+ ]
121
+
122
+ [tool.ruff.lint.pydocstyle]
123
+ convention = "google"
124
+
125
+ [tool.ruff.format]
126
+ quote-style = "double"
127
+ indent-style = "space"
128
+ docstring-code-format = true
129
+ docstring-code-line-length = "dynamic"
130
+
131
+ [tool.ty.src]
132
+ include = [
133
+ "src",
134
+ "tests",
135
+ ]
136
+
137
+ [tool.pytest.ini_options]
138
+ minversion = "8.0"
139
+ testpaths = ["tests"]
140
+ addopts = [
141
+ "-ra",
142
+ "--strict-config",
143
+ "--strict-markers",
144
+ ]
145
+ filterwarnings = [
146
+ "error",
147
+ "ignore:The declarative_mixin decorator:sqlalchemy.exc.SADeprecationWarning",
148
+ ]
149
+
150
+ [tool.coverage.run]
151
+ source = ["src"]
152
+ branch = true
153
+
154
+ [tool.coverage.report]
155
+ fail_under = 95
156
+ exclude_lines = [
157
+ "pragma: no cover",
158
+ "if TYPE_CHECKING:",
159
+ "raise NotImplementedError",
160
+ ]
@@ -0,0 +1,137 @@
1
+ [project]
2
+ name = "xtr-ai-orm-conversation-store"
3
+ version = "6.6.0"
4
+ description = "Agent conversations stored in a relational database through xtr-orm, for xtr-ai."
5
+ readme = "README.md"
6
+ authors = [
7
+ { name = "Xterr", email = "me@xterr.dev" }
8
+ ]
9
+ requires-python = ">=3.11"
10
+ license = "MIT"
11
+ license-files = ["LICENSE"]
12
+ keywords = ["ai", "llm", "agent", "pydantic-ai", "conversation", "sqlalchemy", "database"]
13
+ classifiers = [
14
+ "Development Status :: 3 - Alpha",
15
+ "Framework :: AsyncIO",
16
+ "Intended Audience :: Developers",
17
+ "Programming Language :: Python :: 3.11",
18
+ "Programming Language :: Python :: 3.12",
19
+ "Programming Language :: Python :: 3.13",
20
+ "Programming Language :: Python :: 3.14",
21
+ "Typing :: Typed",
22
+ ]
23
+ dependencies = [
24
+ "pydantic-ai-slim>=2.66,<3",
25
+ "sqlalchemy[asyncio]>=2.0.41",
26
+ "typing-extensions>=4.12",
27
+ "xtr-ai>=6.0,<7",
28
+ "xtr-orm>=6.0,<7",
29
+ ]
30
+
31
+ [dependency-groups]
32
+ dev = [
33
+ "basedpyright>=1.21",
34
+ "ruff>=0.8",
35
+ "pytest>=8",
36
+ "pytest-cov>=5",
37
+ "anyio>=4.0",
38
+ "ty>=0.0.83",
39
+ # The store is tested against a real database: a SQLite file in the test's
40
+ # own directory, which needs an async driver for it.
41
+ "aiosqlite>=0.20",
42
+ ]
43
+
44
+ [build-system]
45
+ requires = ["uv_build>=0.9.18,<0.10.0"]
46
+ build-backend = "uv_build"
47
+
48
+ [tool.basedpyright]
49
+ typeCheckingMode = "all"
50
+ pythonVersion = "3.11"
51
+ pythonPlatform = "All"
52
+ include = ["src", "tests"]
53
+ exclude = ["**/__pycache__", "**/.venv", "**/build", "**/dist", ".tmp"]
54
+ reportUnusedCallResult = "warning"
55
+ reportUnnecessaryTypeIgnoreComment = "error"
56
+ reportUnusedVariable = "error"
57
+ reportMissingParameterType = "error"
58
+ reportPrivateUsage = "error"
59
+ reportImplicitStringConcatenation = "none"
60
+ # SQLAlchemy's rows and pydantic-ai's generic types surface as Any.
61
+ reportExplicitAny = "none"
62
+ reportAny = "none"
63
+
64
+ [[tool.basedpyright.executionEnvironments]]
65
+ root = "tests"
66
+ extraPaths = ["."]
67
+ reportPrivateUsage = "none"
68
+
69
+ [tool.ruff]
70
+ target-version = "py311"
71
+ line-length = 100
72
+ src = ["src", "tests"]
73
+
74
+ [tool.ruff.lint]
75
+ select = ["ALL"]
76
+ ignore = [
77
+ "COM812",
78
+ "ISC001",
79
+ "D203",
80
+ "D213",
81
+ "CPY001",
82
+ "FBT001",
83
+ "FBT002",
84
+ "TD002",
85
+ "TD003",
86
+ "FIX002",
87
+ "TRY003",
88
+ "EM101",
89
+ "EM102",
90
+ "ANN401",
91
+ ]
92
+ fixable = ["ALL"]
93
+ unfixable = []
94
+
95
+ [tool.ruff.lint.per-file-ignores]
96
+ "tests/**/*.py" = [
97
+ "S101", # assert is the point of pytest
98
+ "ARG", # fixtures look unused
99
+ "PLR2004", # magic numbers in test data
100
+ "SLF001", # tests reach into privates
101
+ "D", # test names are the docs
102
+ ]
103
+
104
+ [tool.ruff.lint.pydocstyle]
105
+ convention = "google"
106
+
107
+ [tool.ruff.format]
108
+ quote-style = "double"
109
+ indent-style = "space"
110
+ docstring-code-format = true
111
+ docstring-code-line-length = "dynamic"
112
+
113
+ [tool.ty.src]
114
+ include = ["src", "tests"]
115
+
116
+ [tool.pytest.ini_options]
117
+ minversion = "8.0"
118
+ testpaths = ["tests"]
119
+ addopts = ["-ra", "--strict-config", "--strict-markers"]
120
+ filterwarnings = [
121
+ "error",
122
+ # xtr-orm's database layer reaches advanced-alchemy, which still carries a decorator
123
+ # SQLAlchemy 2.1 retired; it warns once, on import, and changes nothing.
124
+ "ignore:The declarative_mixin decorator:sqlalchemy.exc.SADeprecationWarning",
125
+ ]
126
+
127
+ [tool.coverage.run]
128
+ source = ["src"]
129
+ branch = true
130
+
131
+ [tool.coverage.report]
132
+ fail_under = 95
133
+ exclude_lines = [
134
+ "pragma: no cover",
135
+ "if TYPE_CHECKING:",
136
+ "raise NotImplementedError",
137
+ ]
@@ -0,0 +1,20 @@
1
+ """Agent conversations stored in a relational database, for xtr-ai.
2
+
3
+ :class:`OrmConversationStore` is a pydantic-ai ``ConversationStore`` on two
4
+ tables, written through an xtr-orm connection's engine in short transactions of
5
+ its own. :class:`ConversationMappingDriver` adds those tables to a connection's
6
+ definitions, so ``orm:migrations:diff`` writes them.
7
+ """
8
+
9
+ from __future__ import annotations
10
+
11
+ from .conversation_mapping_driver import ConversationMappingDriver
12
+ from .conversation_tables import DEFAULT_TABLE_PREFIX, ConversationTables
13
+ from .orm_conversation_store import OrmConversationStore
14
+
15
+ __all__ = [
16
+ "DEFAULT_TABLE_PREFIX",
17
+ "ConversationMappingDriver",
18
+ "ConversationTables",
19
+ "OrmConversationStore",
20
+ ]
@@ -0,0 +1,89 @@
1
+ """Turning records and entries into rows and back."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from datetime import UTC, datetime
6
+ from typing import TYPE_CHECKING, Any, Final, Literal, cast
7
+
8
+ import pydantic
9
+ from pydantic_ai import DeferredToolRequests
10
+ from pydantic_ai.conversation_store import ConversationEntry, ConversationRecord
11
+ from pydantic_ai.messages import ModelMessage
12
+ from pydantic_ai.usage import RunUsage
13
+
14
+ if TYPE_CHECKING:
15
+ from sqlalchemy import Row
16
+
17
+ __all__ = ["entry_from_row", "entry_row", "record_from_row", "record_row"]
18
+
19
+ # Bytes travel as base64 in both directions, as pydantic-ai's own message adapter does.
20
+ _JSON: Final = pydantic.ConfigDict(ser_json_bytes="base64", val_json_bytes="base64")
21
+ _MESSAGE: Final = pydantic.TypeAdapter[ModelMessage](ModelMessage, config=_JSON)
22
+ _USAGE: Final = pydantic.TypeAdapter(RunUsage)
23
+ _DEFERRED: Final = pydantic.TypeAdapter(DeferredToolRequests)
24
+ _BRANCHES: Final = pydantic.TypeAdapter(list[str])
25
+
26
+
27
+ def record_row(record: ConversationRecord) -> dict[str, Any]:
28
+ """The ``conversation`` row holding ``record``."""
29
+ deferred = record.deferred_tool_requests
30
+ return {
31
+ "conversation_id": record.conversation_id,
32
+ "head_id": record.head_id,
33
+ "scope": record.scope,
34
+ "usage": _USAGE.dump_json(record.usage).decode(),
35
+ "deferred_tool_requests": _DEFERRED.dump_json(deferred).decode() if deferred else None,
36
+ "forked_from": record.forked_from,
37
+ "forked_at": record.forked_at,
38
+ "branches": _BRANCHES.dump_json(record.branches).decode(),
39
+ "revision": record.revision,
40
+ "updated_at": _naive_utc(record.updated_at),
41
+ "format_version": record.format_version,
42
+ }
43
+
44
+
45
+ def record_from_row(row: Row[Any]) -> ConversationRecord:
46
+ """The record a ``conversation`` row holds."""
47
+ deferred = cast("str | None", row.deferred_tool_requests)
48
+ updated_at = cast("datetime", row.updated_at)
49
+ return ConversationRecord(
50
+ conversation_id=cast("str", row.conversation_id),
51
+ head_id=cast("str | None", row.head_id),
52
+ usage=_USAGE.validate_json(cast("str", row.usage)),
53
+ deferred_tool_requests=_DEFERRED.validate_json(deferred) if deferred else None,
54
+ scope=cast("str | None", row.scope),
55
+ forked_from=cast("str | None", row.forked_from),
56
+ forked_at=cast("str | None", row.forked_at),
57
+ branches=_BRANCHES.validate_json(cast("str", row.branches)),
58
+ revision=cast("str", row.revision),
59
+ updated_at=updated_at.replace(tzinfo=UTC),
60
+ format_version=cast("Literal[1]", row.format_version),
61
+ )
62
+
63
+
64
+ def entry_row(entry: ConversationEntry) -> dict[str, Any]:
65
+ """The ``conversation_entry`` row holding ``entry``."""
66
+ return {
67
+ "id": entry.id,
68
+ "parent_id": entry.parent_id,
69
+ "kind": entry.kind,
70
+ "message": _MESSAGE.dump_json(entry.message).decode(),
71
+ }
72
+
73
+
74
+ def entry_from_row(row: Row[Any]) -> ConversationEntry:
75
+ """The entry a ``conversation_entry`` row holds."""
76
+ return ConversationEntry(
77
+ id=cast("str", row.id),
78
+ parent_id=cast("str | None", row.parent_id),
79
+ message=_MESSAGE.validate_json(cast("str", row.message)),
80
+ kind=cast("Literal['message']", row.kind),
81
+ )
82
+
83
+
84
+ def _naive_utc(instant: datetime) -> datetime:
85
+ """``instant``'s UTC wall clock without a zone, as a zone-less column keeps it everywhere."""
86
+ utc = instant.astimezone(UTC)
87
+ return datetime( # noqa: DTZ001 — a zone-less column wants exactly a naive value
88
+ utc.year, utc.month, utc.day, utc.hour, utc.minute, utc.second, utc.microsecond
89
+ )
@@ -0,0 +1,38 @@
1
+ """Adding the conversation tables to a connection's definitions."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from typing import TYPE_CHECKING, final
6
+
7
+ from typing_extensions import override
8
+ from xtr_orm import MappingDriverInterface
9
+
10
+ from .conversation_tables import DEFAULT_TABLE_PREFIX, ConversationTables
11
+
12
+ if TYPE_CHECKING:
13
+ from sqlalchemy import MetaData
14
+
15
+ __all__ = ["ConversationMappingDriver"]
16
+
17
+
18
+ @final
19
+ class ConversationMappingDriver(MappingDriverInterface):
20
+ """Adds the two conversation tables to a connection's definitions.
21
+
22
+ Register it through xtr-orm's ``OrmMappingsPass`` so the connection's
23
+ migrations know the tables; the store itself defines its own copy and needs
24
+ nothing from here.
25
+ """
26
+
27
+ __slots__ = ("_table_prefix",)
28
+
29
+ def __init__(self, table_prefix: str = DEFAULT_TABLE_PREFIX) -> None:
30
+ """Add tables whose names start with ``table_prefix``."""
31
+ self._table_prefix = table_prefix
32
+
33
+ @override
34
+ def load_metadata(self, metadata: MetaData, /) -> tuple[str, ...]:
35
+ """Add the tables ``metadata`` lacks; return the names added, ``()`` the second time."""
36
+ known = set(metadata.tables)
37
+ _ = ConversationTables.define(metadata, self._table_prefix)
38
+ return tuple(name for name in metadata.tables if name not in known)
@@ -0,0 +1,76 @@
1
+ """The two tables a conversation is stored in."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from dataclasses import dataclass
6
+ from typing import Final, final
7
+
8
+ from sqlalchemy import Column, DateTime, Index, MetaData, SmallInteger, String, Table, Text
9
+ from sqlalchemy.dialects.mysql import DATETIME, LONGTEXT
10
+
11
+ __all__ = ["DEFAULT_TABLE_PREFIX", "ConversationTables"]
12
+
13
+ DEFAULT_TABLE_PREFIX: Final = "ai_"
14
+ """What both table names start with unless told otherwise."""
15
+
16
+ _ID: Final = 64
17
+ _NAME: Final = 255
18
+
19
+ # MySQL's TEXT holds 64 KB, less than one message carrying an image, and its DATETIME keeps
20
+ # whole seconds, which leaves conversations saved in the same second unordered.
21
+ _JSON: Final = Text().with_variant(LONGTEXT(), "mysql")
22
+ _INSTANT: Final = DateTime(timezone=False).with_variant(DATETIME(fsp=6), "mysql")
23
+
24
+
25
+ @final
26
+ @dataclass(frozen=True, slots=True)
27
+ class ConversationTables:
28
+ """``{prefix}conversation``, a row per conversation; ``{prefix}conversation_entry``, a message.
29
+
30
+ An entry holds one message and points at the entry before it; entries are
31
+ shared between a conversation and its forks, so they carry no conversation
32
+ of their own. Every structured value is pydantic's JSON, as text, so binary
33
+ content keeps its base64 form on every database.
34
+ """
35
+
36
+ conversation: Table
37
+ entry: Table
38
+
39
+ @classmethod
40
+ def define(
41
+ cls, metadata: MetaData, table_prefix: str = DEFAULT_TABLE_PREFIX
42
+ ) -> ConversationTables:
43
+ """Return the tables in ``metadata``, adding them first where they are not there."""
44
+ conversation_name = f"{table_prefix}conversation"
45
+ entry_name = f"{table_prefix}conversation_entry"
46
+ conversation = metadata.tables.get(conversation_name)
47
+ if conversation is None:
48
+ conversation = Table(
49
+ conversation_name,
50
+ metadata,
51
+ Column("conversation_id", String(_NAME), primary_key=True),
52
+ Column("head_id", String(_ID), nullable=True),
53
+ Column("scope", String(_NAME), nullable=True),
54
+ Column("usage", _JSON, nullable=False),
55
+ Column("deferred_tool_requests", _JSON, nullable=True),
56
+ Column("forked_from", String(_NAME), nullable=True),
57
+ Column("forked_at", String(_ID), nullable=True),
58
+ Column("branches", _JSON, nullable=False),
59
+ Column("revision", String(_ID), nullable=False),
60
+ Column("updated_at", _INSTANT, nullable=False),
61
+ Column("format_version", SmallInteger, nullable=False),
62
+ Index(None, "scope", "updated_at"),
63
+ Index(None, "forked_from"),
64
+ )
65
+ entry = metadata.tables.get(entry_name)
66
+ if entry is None:
67
+ entry = Table(
68
+ entry_name,
69
+ metadata,
70
+ Column("id", String(_ID), primary_key=True),
71
+ Column("parent_id", String(_ID), nullable=True),
72
+ Column("kind", String(32), nullable=False),
73
+ Column("message", _JSON, nullable=False),
74
+ Index(None, "parent_id"),
75
+ )
76
+ return cls(conversation=conversation, entry=entry)
@@ -0,0 +1,243 @@
1
+ """A pydantic-ai conversation store on two database tables."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from typing import TYPE_CHECKING, final
6
+
7
+ from pydantic_ai.conversation_store import ConversationConflict, ConversationStore
8
+ from sqlalchemy import MetaData, delete, func, insert, literal, or_, select, update
9
+ from sqlalchemy.exc import IntegrityError
10
+ from typing_extensions import override
11
+
12
+ from ._rows import entry_from_row, entry_row, record_from_row, record_row
13
+ from .conversation_tables import DEFAULT_TABLE_PREFIX, ConversationTables
14
+
15
+ if TYPE_CHECKING:
16
+ from collections.abc import Iterable, Sequence
17
+
18
+ from pydantic_ai.conversation_store import ConversationEntry, ConversationRecord
19
+ from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine
20
+
21
+ __all__ = ["OrmConversationStore"]
22
+
23
+
24
+ @final
25
+ class OrmConversationStore(ConversationStore):
26
+ """Conversations in ``{prefix}conversation`` and ``{prefix}conversation_entry``.
27
+
28
+ Every call runs in a short transaction of its own on ``engine`` — give it
29
+ the primary's — and commits before it returns: a conversation is never
30
+ part of the caller's unit of work, so a long model call holds no database
31
+ connection, and a rollback of the caller's work never takes a saved turn
32
+ with it. Saving a record and the messages it adds is one transaction, so a
33
+ conflict leaves neither behind.
34
+
35
+ A conversation's history is read with one recursive query. Deleting a
36
+ conversation deletes every message only it reaches; the messages a fork of
37
+ it — or the conversation it was forked from — still reaches stay until
38
+ the last of them is deleted.
39
+
40
+ ```python
41
+ store = OrmConversationStore(engine)
42
+ conversation = await store.load(conversation_id) or Conversation(
43
+ conversation_id=conversation_id
44
+ )
45
+ result = await agent.run(prompt, conversation=conversation)
46
+ await store.save(result)
47
+ ```
48
+ """
49
+
50
+ def __init__(self, engine: AsyncEngine, *, table_prefix: str = DEFAULT_TABLE_PREFIX) -> None:
51
+ """Store through ``engine``, in tables whose names start with ``table_prefix``."""
52
+ self._engine = engine
53
+ self._tables = ConversationTables.define(MetaData(), table_prefix)
54
+
55
+ @override
56
+ async def read_record(self, conversation_id: str) -> ConversationRecord | None:
57
+ conversation = self._tables.conversation
58
+ async with self._engine.connect() as connection:
59
+ row = (
60
+ await connection.execute(
61
+ select(conversation).where(conversation.c.conversation_id == conversation_id)
62
+ )
63
+ ).first()
64
+ return record_from_row(row) if row is not None else None
65
+
66
+ @override
67
+ async def read_records(
68
+ self, *, scope: str | None = None, limit: int | None = None
69
+ ) -> Sequence[ConversationRecord]:
70
+ conversation = self._tables.conversation
71
+ query = select(conversation).order_by(
72
+ conversation.c.updated_at.desc(), conversation.c.conversation_id
73
+ )
74
+ if scope is not None:
75
+ query = query.where(conversation.c.scope == scope)
76
+ if limit is not None:
77
+ query = query.limit(limit)
78
+ async with self._engine.connect() as connection:
79
+ rows = (await connection.execute(query)).all()
80
+ return [record_from_row(row) for row in rows]
81
+
82
+ @override
83
+ async def write_record(
84
+ self,
85
+ record: ConversationRecord,
86
+ *,
87
+ new_entries: Sequence[ConversationEntry] = (),
88
+ expected_revision: str | None,
89
+ ) -> None:
90
+ conversation = self._tables.conversation
91
+ async with self._engine.begin() as connection:
92
+ if expected_revision is None:
93
+ try:
94
+ _ = await connection.execute(insert(conversation).values(record_row(record)))
95
+ except IntegrityError:
96
+ raise ConversationConflict(record.conversation_id) from None
97
+ else:
98
+ result = await connection.execute(
99
+ update(conversation)
100
+ .where(
101
+ conversation.c.conversation_id == record.conversation_id,
102
+ conversation.c.revision == expected_revision,
103
+ )
104
+ .values(record_row(record))
105
+ )
106
+ if result.rowcount != 1:
107
+ raise ConversationConflict(record.conversation_id)
108
+ if new_entries:
109
+ _ = await connection.execute(
110
+ insert(self._tables.entry), [entry_row(entry) for entry in new_entries]
111
+ )
112
+ head = record.head_id
113
+ new_ids = {entry.id for entry in new_entries}
114
+ if (
115
+ head is not None
116
+ and head not in new_ids
117
+ and not await self._has_entry(connection, head)
118
+ ):
119
+ raise ConversationConflict(record.conversation_id)
120
+
121
+ @override
122
+ async def delete_record(self, conversation_id: str, *, expected_revision: str) -> None:
123
+ conversation = self._tables.conversation
124
+ async with self._engine.begin() as connection:
125
+ family = await self._fork_family(connection, conversation_id)
126
+ deleted = await connection.execute(
127
+ delete(conversation).where(
128
+ conversation.c.conversation_id == conversation_id,
129
+ conversation.c.revision == expected_revision,
130
+ )
131
+ )
132
+ if deleted.rowcount != 1:
133
+ raise ConversationConflict(conversation_id)
134
+ own = family.pop(conversation_id)
135
+ reached = await self._reachable(connection, own)
136
+ kept = await self._reachable(
137
+ connection, [tip for tips in family.values() for tip in tips]
138
+ )
139
+ released = sorted(reached - kept)
140
+ if released:
141
+ entry = self._tables.entry
142
+ _ = await connection.execute(delete(entry).where(entry.c.id.in_(released)))
143
+
144
+ @override
145
+ async def read_entry(self, entry_id: str) -> ConversationEntry:
146
+ entry = self._tables.entry
147
+ async with self._engine.connect() as connection:
148
+ row = (await connection.execute(select(entry).where(entry.c.id == entry_id))).first()
149
+ if row is None:
150
+ raise KeyError(entry_id)
151
+ return entry_from_row(row)
152
+
153
+ @override
154
+ async def read_history(
155
+ self, head_id: str | None, *, limit: int | None = None
156
+ ) -> list[ConversationEntry]:
157
+ """Read the history up to ``head_id`` in one recursive query.
158
+
159
+ Raises:
160
+ KeyError: When an entry of the history is missing, as a concurrent
161
+ delete leaves it.
162
+ """
163
+ if head_id is None or limit == 0:
164
+ return []
165
+ entry = self._tables.entry
166
+ walk = (select(entry, literal(1).label("depth")).where(entry.c.id == head_id)).cte(
167
+ "history", recursive=True
168
+ )
169
+ step = select(entry, (walk.c.depth + 1).label("depth")).join(
170
+ walk, entry.c.id == walk.c.parent_id
171
+ )
172
+ if limit is not None:
173
+ step = step.where(walk.c.depth < limit)
174
+ walk = walk.union_all(step)
175
+ async with self._engine.connect() as connection:
176
+ rows = (await connection.execute(select(walk).order_by(walk.c.depth.desc()))).all()
177
+ if not rows:
178
+ raise KeyError(head_id)
179
+ reached_limit = limit is not None and len(rows) >= limit
180
+ if rows[0].parent_id is not None and not reached_limit:
181
+ raise KeyError(rows[0].parent_id)
182
+ return [entry_from_row(row) for row in rows]
183
+
184
+ async def _has_entry(self, connection: AsyncConnection, entry_id: str) -> bool:
185
+ entry = self._tables.entry
186
+ found = await connection.scalar(
187
+ select(func.count()).select_from(entry).where(entry.c.id == entry_id)
188
+ )
189
+ return bool(found)
190
+
191
+ async def _fork_family(
192
+ self, connection: AsyncConnection, conversation_id: str
193
+ ) -> dict[str, list[str]]:
194
+ """Every conversation sharing entries with ``conversation_id``, with its histories' tips.
195
+
196
+ Entries are shared only by forking, so the family is everything linked
197
+ to the conversation by ``forked_from``, in either direction.
198
+ """
199
+ conversation = self._tables.conversation
200
+ family: dict[str, list[str]] = {}
201
+ explored: set[str] = set()
202
+ pending: set[str] = {conversation_id}
203
+ while pending:
204
+ rows = (
205
+ await connection.execute(
206
+ select(conversation).where(
207
+ or_(
208
+ conversation.c.conversation_id.in_(pending),
209
+ conversation.c.forked_from.in_(pending),
210
+ )
211
+ )
212
+ )
213
+ ).all()
214
+ explored |= pending
215
+ pending = set()
216
+ for row in rows:
217
+ record = record_from_row(row)
218
+ family[record.conversation_id] = _tips(record)
219
+ pending.add(record.conversation_id)
220
+ if record.forked_from is not None:
221
+ pending.add(record.forked_from)
222
+ pending -= explored
223
+ return family
224
+
225
+ async def _reachable(self, connection: AsyncConnection, tips: Iterable[str]) -> set[str]:
226
+ """Every entry on the histories ending at ``tips``."""
227
+ starts = sorted(set(tips))
228
+ if not starts:
229
+ return set()
230
+ entry = self._tables.entry
231
+ walk = (select(entry.c.id, entry.c.parent_id).where(entry.c.id.in_(starts))).cte(
232
+ "reached", recursive=True
233
+ )
234
+ walk = walk.union(
235
+ select(entry.c.id, entry.c.parent_id).join(walk, entry.c.id == walk.c.parent_id)
236
+ )
237
+ rows = (await connection.execute(select(walk.c.id))).all()
238
+ return {row.id for row in rows}
239
+
240
+
241
+ def _tips(record: ConversationRecord) -> list[str]:
242
+ """Where ``record``'s histories end: its head, the branches it left, and where it was forked."""
243
+ return [tip for tip in (record.head_id, *record.branches, record.forked_at) if tip is not None]