rfnry-memory 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- rfnry_memory-0.1.0/.gitignore +15 -0
- rfnry_memory-0.1.0/LICENSE +21 -0
- rfnry_memory-0.1.0/PKG-INFO +29 -0
- rfnry_memory-0.1.0/README.md +6 -0
- rfnry_memory-0.1.0/pyproject.toml +84 -0
- rfnry_memory-0.1.0/src/rfnry_memory/__init__.py +153 -0
- rfnry_memory-0.1.0/src/rfnry_memory/chunking/__init__.py +0 -0
- rfnry_memory-0.1.0/src/rfnry_memory/chunking/chunker.py +39 -0
- rfnry_memory-0.1.0/src/rfnry_memory/chunking/splitter.py +93 -0
- rfnry_memory-0.1.0/src/rfnry_memory/chunking/structure.py +110 -0
- rfnry_memory-0.1.0/src/rfnry_memory/cli.py +178 -0
- rfnry_memory-0.1.0/src/rfnry_memory/errors.py +22 -0
- rfnry_memory-0.1.0/src/rfnry_memory/filters.py +104 -0
- rfnry_memory-0.1.0/src/rfnry_memory/fusion/__init__.py +13 -0
- rfnry_memory-0.1.0/src/rfnry_memory/fusion/provenance.py +14 -0
- rfnry_memory-0.1.0/src/rfnry_memory/fusion/relative_score.py +38 -0
- rfnry_memory-0.1.0/src/rfnry_memory/fusion/rrf.py +35 -0
- rfnry_memory-0.1.0/src/rfnry_memory/graph/__init__.py +0 -0
- rfnry_memory-0.1.0/src/rfnry_memory/graph/extraction.py +81 -0
- rfnry_memory-0.1.0/src/rfnry_memory/graph/identity.py +187 -0
- rfnry_memory-0.1.0/src/rfnry_memory/graph/placement.py +65 -0
- rfnry_memory-0.1.0/src/rfnry_memory/index/__init__.py +13 -0
- rfnry_memory-0.1.0/src/rfnry_memory/index/hnsw.py +39 -0
- rfnry_memory-0.1.0/src/rfnry_memory/index/ivfflat.py +33 -0
- rfnry_memory-0.1.0/src/rfnry_memory/judge.py +61 -0
- rfnry_memory-0.1.0/src/rfnry_memory/memory.py +1954 -0
- rfnry_memory-0.1.0/src/rfnry_memory/models.py +634 -0
- rfnry_memory-0.1.0/src/rfnry_memory/postgres/__init__.py +11 -0
- rfnry_memory-0.1.0/src/rfnry_memory/postgres/base.py +28 -0
- rfnry_memory-0.1.0/src/rfnry_memory/postgres/corpus.py +116 -0
- rfnry_memory-0.1.0/src/rfnry_memory/postgres/graph.py +422 -0
- rfnry_memory-0.1.0/src/rfnry_memory/postgres/language.py +52 -0
- rfnry_memory-0.1.0/src/rfnry_memory/postgres/locks.py +14 -0
- rfnry_memory-0.1.0/src/rfnry_memory/postgres/migrations.py +354 -0
- rfnry_memory-0.1.0/src/rfnry_memory/postgres/records.py +110 -0
- rfnry_memory-0.1.0/src/rfnry_memory/postgres/schema.py +93 -0
- rfnry_memory-0.1.0/src/rfnry_memory/postgres/search.py +286 -0
- rfnry_memory-0.1.0/src/rfnry_memory/postgres/sources.py +374 -0
- rfnry_memory-0.1.0/src/rfnry_memory/postgres/sql.py +133 -0
- rfnry_memory-0.1.0/src/rfnry_memory/postgres/stories.py +279 -0
- rfnry_memory-0.1.0/src/rfnry_memory/postgres/vectors.py +324 -0
- rfnry_memory-0.1.0/src/rfnry_memory/py.typed +0 -0
- rfnry_memory-0.1.0/src/rfnry_memory/ranking/__init__.py +13 -0
- rfnry_memory-0.1.0/src/rfnry_memory/ranking/bm25.py +46 -0
- rfnry_memory-0.1.0/src/rfnry_memory/ranking/cover_density.py +36 -0
- rfnry_memory-0.1.0/src/rfnry_memory/rows.py +110 -0
- rfnry_memory-0.1.0/src/rfnry_memory/structured.py +57 -0
- rfnry_memory-0.1.0/src/rfnry_memory/tools.py +897 -0
- rfnry_memory-0.1.0/tests/conftest.py +110 -0
- rfnry_memory-0.1.0/tests/database.py +60 -0
- rfnry_memory-0.1.0/tests/e2e/__init__.py +0 -0
- rfnry_memory-0.1.0/tests/e2e/conftest.py +22 -0
- rfnry_memory-0.1.0/tests/e2e/test_live.py +104 -0
- rfnry_memory-0.1.0/tests/integration/__init__.py +0 -0
- rfnry_memory-0.1.0/tests/integration/test_cli.py +261 -0
- rfnry_memory-0.1.0/tests/integration/test_migrations.py +350 -0
- rfnry_memory-0.1.0/tests/integration/test_postgres.py +187 -0
- rfnry_memory-0.1.0/tests/integration/test_storage.py +3706 -0
- rfnry_memory-0.1.0/tests/integration/test_tenancy.py +238 -0
- rfnry_memory-0.1.0/tests/integration/test_vector_index.py +1249 -0
- rfnry_memory-0.1.0/tests/unit/__init__.py +0 -0
- rfnry_memory-0.1.0/tests/unit/conftest.py +8 -0
- rfnry_memory-0.1.0/tests/unit/fakes.py +86 -0
- rfnry_memory-0.1.0/tests/unit/lookalike.py +144 -0
- rfnry_memory-0.1.0/tests/unit/test_api.py +197 -0
- rfnry_memory-0.1.0/tests/unit/test_chunker.py +46 -0
- rfnry_memory-0.1.0/tests/unit/test_comments.py +23 -0
- rfnry_memory-0.1.0/tests/unit/test_context.py +88 -0
- rfnry_memory-0.1.0/tests/unit/test_corpora.py +285 -0
- rfnry_memory-0.1.0/tests/unit/test_docs.py +10 -0
- rfnry_memory-0.1.0/tests/unit/test_docs_examples.py +237 -0
- rfnry_memory-0.1.0/tests/unit/test_document_embedding.py +63 -0
- rfnry_memory-0.1.0/tests/unit/test_errors.py +15 -0
- rfnry_memory-0.1.0/tests/unit/test_expand.py +325 -0
- rfnry_memory-0.1.0/tests/unit/test_family.py +49 -0
- rfnry_memory-0.1.0/tests/unit/test_filters.py +44 -0
- rfnry_memory-0.1.0/tests/unit/test_fingerprint.py +249 -0
- rfnry_memory-0.1.0/tests/unit/test_foreign_contracts.py +151 -0
- rfnry_memory-0.1.0/tests/unit/test_fusion.py +176 -0
- rfnry_memory-0.1.0/tests/unit/test_graph.py +658 -0
- rfnry_memory-0.1.0/tests/unit/test_graph_search.py +347 -0
- rfnry_memory-0.1.0/tests/unit/test_history.py +204 -0
- rfnry_memory-0.1.0/tests/unit/test_index_hnsw.py +30 -0
- rfnry_memory-0.1.0/tests/unit/test_index_ivfflat.py +19 -0
- rfnry_memory-0.1.0/tests/unit/test_judge.py +215 -0
- rfnry_memory-0.1.0/tests/unit/test_judge_subsets.py +63 -0
- rfnry_memory-0.1.0/tests/unit/test_language.py +233 -0
- rfnry_memory-0.1.0/tests/unit/test_memory_read.py +211 -0
- rfnry_memory-0.1.0/tests/unit/test_memory_write.py +909 -0
- rfnry_memory-0.1.0/tests/unit/test_messages.py +113 -0
- rfnry_memory-0.1.0/tests/unit/test_model_shapes.py +129 -0
- rfnry_memory-0.1.0/tests/unit/test_models.py +223 -0
- rfnry_memory-0.1.0/tests/unit/test_postgres_layout.py +48 -0
- rfnry_memory-0.1.0/tests/unit/test_public_api.py +98 -0
- rfnry_memory-0.1.0/tests/unit/test_query.py +259 -0
- rfnry_memory-0.1.0/tests/unit/test_ranking.py +159 -0
- rfnry_memory-0.1.0/tests/unit/test_scale.py +236 -0
- rfnry_memory-0.1.0/tests/unit/test_scope.py +244 -0
- rfnry_memory-0.1.0/tests/unit/test_structured.py +63 -0
- rfnry_memory-0.1.0/tests/unit/test_surface.py +214 -0
- rfnry_memory-0.1.0/tests/unit/test_sync.py +479 -0
- rfnry_memory-0.1.0/tests/unit/test_tools.py +1075 -0
- rfnry_memory-0.1.0/uv.lock +605 -0
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 rfnry (github.com/rfnry)
|
|
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,29 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: rfnry-memory
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Hybrid-retrieval memory system for AI agents
|
|
5
|
+
Project-URL: Repository, https://github.com/rfnry/memory
|
|
6
|
+
Project-URL: Documentation, https://github.com/rfnry/memory/blob/main/docs/memory.md
|
|
7
|
+
Project-URL: Issues, https://github.com/rfnry/memory/issues
|
|
8
|
+
Author-email: rfnry <contact@rfnry.dev>
|
|
9
|
+
License-Expression: MIT
|
|
10
|
+
License-File: LICENSE
|
|
11
|
+
Requires-Python: >=3.12
|
|
12
|
+
Requires-Dist: pydantic>=2.13.4
|
|
13
|
+
Provides-Extra: dev
|
|
14
|
+
Requires-Dist: asyncpg<1.0,>=0.31.0; extra == 'dev'
|
|
15
|
+
Requires-Dist: mypy>=2.3.1; extra == 'dev'
|
|
16
|
+
Requires-Dist: poethepoet>=0.48.0; extra == 'dev'
|
|
17
|
+
Requires-Dist: pytest-asyncio>=1.4.0; extra == 'dev'
|
|
18
|
+
Requires-Dist: pytest>=9.1.1; extra == 'dev'
|
|
19
|
+
Requires-Dist: ruff>=0.16.4; extra == 'dev'
|
|
20
|
+
Provides-Extra: postgres
|
|
21
|
+
Requires-Dist: asyncpg<1.0,>=0.31.0; extra == 'postgres'
|
|
22
|
+
Description-Content-Type: text/markdown
|
|
23
|
+
|
|
24
|
+
# rfnry-memory
|
|
25
|
+
|
|
26
|
+
Hybrid-retrieval memory system for AI agents.
|
|
27
|
+
|
|
28
|
+
[`docs/memory.md`](https://github.com/rfnry/memory/blob/main/docs/memory.md).
|
|
29
|
+
[`docs/api.md`](https://github.com/rfnry/memory/blob/main/docs/api.md).
|
|
@@ -0,0 +1,84 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "rfnry-memory"
|
|
3
|
+
version = "0.1.0"
|
|
4
|
+
description = "Hybrid-retrieval memory system for AI agents"
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
license = "MIT"
|
|
7
|
+
requires-python = ">=3.12"
|
|
8
|
+
authors = [{name = "rfnry", email = "contact@rfnry.dev"}]
|
|
9
|
+
dependencies = [
|
|
10
|
+
"pydantic>=2.13.4",
|
|
11
|
+
]
|
|
12
|
+
|
|
13
|
+
[project.urls]
|
|
14
|
+
Repository = "https://github.com/rfnry/memory"
|
|
15
|
+
Documentation = "https://github.com/rfnry/memory/blob/main/docs/memory.md"
|
|
16
|
+
Issues = "https://github.com/rfnry/memory/issues"
|
|
17
|
+
|
|
18
|
+
[project.scripts]
|
|
19
|
+
rfnry-memory = "rfnry_memory.cli:main"
|
|
20
|
+
|
|
21
|
+
[project.optional-dependencies]
|
|
22
|
+
postgres = ["asyncpg>=0.31.0,<1.0"]
|
|
23
|
+
dev = [
|
|
24
|
+
"pytest>=9.1.1",
|
|
25
|
+
"pytest-asyncio>=1.4.0",
|
|
26
|
+
"ruff>=0.16.4",
|
|
27
|
+
"mypy>=2.3.1",
|
|
28
|
+
"poethepoet>=0.48.0",
|
|
29
|
+
"asyncpg>=0.31.0,<1.0",
|
|
30
|
+
]
|
|
31
|
+
|
|
32
|
+
[build-system]
|
|
33
|
+
requires = ["hatchling"]
|
|
34
|
+
build-backend = "hatchling.build"
|
|
35
|
+
|
|
36
|
+
[tool.hatch.build.targets.wheel]
|
|
37
|
+
packages = ["src/rfnry_memory"]
|
|
38
|
+
|
|
39
|
+
[tool.pytest.ini_options]
|
|
40
|
+
asyncio_mode = "auto"
|
|
41
|
+
asyncio_default_fixture_loop_scope = "session"
|
|
42
|
+
asyncio_default_test_loop_scope = "session"
|
|
43
|
+
testpaths = ["tests"]
|
|
44
|
+
pythonpath = ["src", "tests"]
|
|
45
|
+
markers = [
|
|
46
|
+
"live: calls real models (set RFNRY_LIVE_ADAPTERS to a module path)",
|
|
47
|
+
]
|
|
48
|
+
addopts = "-m 'not live'"
|
|
49
|
+
|
|
50
|
+
[tool.ruff]
|
|
51
|
+
line-length = 100
|
|
52
|
+
target-version = "py312"
|
|
53
|
+
|
|
54
|
+
[tool.ruff.lint]
|
|
55
|
+
select = ["E", "F", "I", "W", "B", "UP", "SIM"]
|
|
56
|
+
ignore = ["E501"]
|
|
57
|
+
|
|
58
|
+
[tool.ruff.lint.isort]
|
|
59
|
+
combine-as-imports = true
|
|
60
|
+
|
|
61
|
+
[tool.ruff.lint.flake8-bugbear]
|
|
62
|
+
extend-immutable-calls = [
|
|
63
|
+
"rfnry_memory.ranking.BM25",
|
|
64
|
+
"rfnry_memory.ranking.bm25.BM25",
|
|
65
|
+
"rfnry_memory.fusion.RRF",
|
|
66
|
+
"rfnry_memory.fusion.rrf.RRF",
|
|
67
|
+
]
|
|
68
|
+
|
|
69
|
+
[tool.poe.tasks]
|
|
70
|
+
format = { cmd = "ruff format src tests" }
|
|
71
|
+
check = { cmd = "ruff check src tests" }
|
|
72
|
+
"check:fix" = { cmd = "ruff check --fix src tests" }
|
|
73
|
+
typecheck = { cmd = "mypy" }
|
|
74
|
+
test = { cmd = "pytest -q" }
|
|
75
|
+
"test:integration" = { cmd = "pytest -q tests/integration" }
|
|
76
|
+
"test:e2e" = { cmd = "pytest -q -m live" }
|
|
77
|
+
|
|
78
|
+
[tool.mypy]
|
|
79
|
+
strict = true
|
|
80
|
+
files = ["src/rfnry_memory", "tests"]
|
|
81
|
+
|
|
82
|
+
[[tool.mypy.overrides]]
|
|
83
|
+
module = ["asyncpg.*"]
|
|
84
|
+
ignore_missing_imports = true
|
|
@@ -0,0 +1,153 @@
|
|
|
1
|
+
from rfnry_memory.errors import (
|
|
2
|
+
ConfigurationError,
|
|
3
|
+
InputError,
|
|
4
|
+
MemoryException,
|
|
5
|
+
ModelError,
|
|
6
|
+
NotFoundError,
|
|
7
|
+
StorageError,
|
|
8
|
+
)
|
|
9
|
+
from rfnry_memory.fusion import RRF, Fusion, RelativeScore
|
|
10
|
+
from rfnry_memory.index import HNSW, IVFFlat, VectorIndex
|
|
11
|
+
from rfnry_memory.memory import Memory
|
|
12
|
+
from rfnry_memory.models import (
|
|
13
|
+
Answer,
|
|
14
|
+
AnswerLike,
|
|
15
|
+
Audio,
|
|
16
|
+
Check,
|
|
17
|
+
CheckAnswer,
|
|
18
|
+
Choice,
|
|
19
|
+
ChoiceAnswer,
|
|
20
|
+
Chunk,
|
|
21
|
+
Chunking,
|
|
22
|
+
Content,
|
|
23
|
+
CorpusDescription,
|
|
24
|
+
CorpusFingerprint,
|
|
25
|
+
Counts,
|
|
26
|
+
Decided,
|
|
27
|
+
DecidedLike,
|
|
28
|
+
DecisionModel,
|
|
29
|
+
Document,
|
|
30
|
+
DocumentEmbeddingModel,
|
|
31
|
+
EmbeddingModel,
|
|
32
|
+
Entity,
|
|
33
|
+
EntityIn,
|
|
34
|
+
GenerationModel,
|
|
35
|
+
Hit,
|
|
36
|
+
Hits,
|
|
37
|
+
Image,
|
|
38
|
+
Ingested,
|
|
39
|
+
Judgement,
|
|
40
|
+
JudgementName,
|
|
41
|
+
Kind,
|
|
42
|
+
Language,
|
|
43
|
+
Match,
|
|
44
|
+
Message,
|
|
45
|
+
Neighborhood,
|
|
46
|
+
Part,
|
|
47
|
+
PillarReport,
|
|
48
|
+
Question,
|
|
49
|
+
QuestionLike,
|
|
50
|
+
Ref,
|
|
51
|
+
Relation,
|
|
52
|
+
RelationIn,
|
|
53
|
+
Remembered,
|
|
54
|
+
Reply,
|
|
55
|
+
ReplyLike,
|
|
56
|
+
Score,
|
|
57
|
+
ScoreAnswer,
|
|
58
|
+
Section,
|
|
59
|
+
Source,
|
|
60
|
+
SourceIn,
|
|
61
|
+
State,
|
|
62
|
+
Story,
|
|
63
|
+
Synced,
|
|
64
|
+
Text,
|
|
65
|
+
ToolCall,
|
|
66
|
+
ToolSpec,
|
|
67
|
+
Usage,
|
|
68
|
+
confidence,
|
|
69
|
+
)
|
|
70
|
+
from rfnry_memory.postgres import PostgresStorage
|
|
71
|
+
from rfnry_memory.ranking import BM25, CoverDensity, TextRanking
|
|
72
|
+
from rfnry_memory.tools import MemoryTools, Tool, Tools
|
|
73
|
+
|
|
74
|
+
__version__ = "0.1.0"
|
|
75
|
+
|
|
76
|
+
__all__ = [
|
|
77
|
+
"Answer",
|
|
78
|
+
"AnswerLike",
|
|
79
|
+
"Check",
|
|
80
|
+
"CheckAnswer",
|
|
81
|
+
"Choice",
|
|
82
|
+
"ChoiceAnswer",
|
|
83
|
+
"Chunk",
|
|
84
|
+
"Chunking",
|
|
85
|
+
"ConfigurationError",
|
|
86
|
+
"CorpusDescription",
|
|
87
|
+
"CorpusFingerprint",
|
|
88
|
+
"Counts",
|
|
89
|
+
"Decided",
|
|
90
|
+
"DecidedLike",
|
|
91
|
+
"DecisionModel",
|
|
92
|
+
"DocumentEmbeddingModel",
|
|
93
|
+
"EmbeddingModel",
|
|
94
|
+
"Entity",
|
|
95
|
+
"EntityIn",
|
|
96
|
+
"GenerationModel",
|
|
97
|
+
"HNSW",
|
|
98
|
+
"BM25",
|
|
99
|
+
"CoverDensity",
|
|
100
|
+
"TextRanking",
|
|
101
|
+
"RRF",
|
|
102
|
+
"RelativeScore",
|
|
103
|
+
"Fusion",
|
|
104
|
+
"Hit",
|
|
105
|
+
"Hits",
|
|
106
|
+
"IVFFlat",
|
|
107
|
+
"Ingested",
|
|
108
|
+
"InputError",
|
|
109
|
+
"Judgement",
|
|
110
|
+
"JudgementName",
|
|
111
|
+
"Kind",
|
|
112
|
+
"Match",
|
|
113
|
+
"Language",
|
|
114
|
+
"Memory",
|
|
115
|
+
"MemoryException",
|
|
116
|
+
"MemoryTools",
|
|
117
|
+
"Message",
|
|
118
|
+
"ModelError",
|
|
119
|
+
"Neighborhood",
|
|
120
|
+
"NotFoundError",
|
|
121
|
+
"PillarReport",
|
|
122
|
+
"PostgresStorage",
|
|
123
|
+
"Question",
|
|
124
|
+
"QuestionLike",
|
|
125
|
+
"Ref",
|
|
126
|
+
"Relation",
|
|
127
|
+
"RelationIn",
|
|
128
|
+
"Remembered",
|
|
129
|
+
"Reply",
|
|
130
|
+
"ReplyLike",
|
|
131
|
+
"Score",
|
|
132
|
+
"ScoreAnswer",
|
|
133
|
+
"Section",
|
|
134
|
+
"Source",
|
|
135
|
+
"SourceIn",
|
|
136
|
+
"State",
|
|
137
|
+
"StorageError",
|
|
138
|
+
"Story",
|
|
139
|
+
"Synced",
|
|
140
|
+
"Tool",
|
|
141
|
+
"ToolCall",
|
|
142
|
+
"ToolSpec",
|
|
143
|
+
"Tools",
|
|
144
|
+
"Usage",
|
|
145
|
+
"VectorIndex",
|
|
146
|
+
"confidence",
|
|
147
|
+
"Text",
|
|
148
|
+
"Image",
|
|
149
|
+
"Audio",
|
|
150
|
+
"Document",
|
|
151
|
+
"Part",
|
|
152
|
+
"Content",
|
|
153
|
+
]
|
|
File without changes
|
|
@@ -0,0 +1,39 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from pydantic import BaseModel
|
|
4
|
+
|
|
5
|
+
from rfnry_memory.chunking.splitter import RecursiveTextSplitter
|
|
6
|
+
from rfnry_memory.chunking.structure import (
|
|
7
|
+
build_heading_spans,
|
|
8
|
+
find_atomic_regions,
|
|
9
|
+
section_path_at,
|
|
10
|
+
)
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
class Piece(BaseModel):
|
|
14
|
+
title: str | None
|
|
15
|
+
text: str
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def chunk_text(text: str, *, chunk_chars: int, overlap_chars: int) -> list[Piece]:
|
|
19
|
+
text = text.strip()
|
|
20
|
+
if not text:
|
|
21
|
+
return []
|
|
22
|
+
regions = find_atomic_regions(text)
|
|
23
|
+
spans = build_heading_spans(text, regions)
|
|
24
|
+
splitter = RecursiveTextSplitter(chunk_size=chunk_chars, chunk_overlap=overlap_chars)
|
|
25
|
+
pieces: list[Piece] = []
|
|
26
|
+
|
|
27
|
+
def free(start: int, end: int) -> None:
|
|
28
|
+
for piece, offset in splitter.split_text_with_offsets(text[start:end]):
|
|
29
|
+
pieces.append(Piece(title=section_path_at(start + offset, spans), text=piece))
|
|
30
|
+
|
|
31
|
+
cursor = 0
|
|
32
|
+
for region in regions:
|
|
33
|
+
free(cursor, region.start)
|
|
34
|
+
pieces.append(
|
|
35
|
+
Piece(title=section_path_at(region.start, spans), text=region.content.strip())
|
|
36
|
+
)
|
|
37
|
+
cursor = region.end
|
|
38
|
+
free(cursor, len(text))
|
|
39
|
+
return [piece for piece in pieces if piece.text]
|
|
@@ -0,0 +1,93 @@
|
|
|
1
|
+
import re
|
|
2
|
+
|
|
3
|
+
SEPARATORS = ["\n\n", "\n", ". ", " ", ""]
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
def _split(text: str, separator: str) -> list[tuple[int, int]]:
|
|
7
|
+
if not separator:
|
|
8
|
+
return [(i, i + 1) for i in range(len(text))]
|
|
9
|
+
spans: list[tuple[int, int]] = []
|
|
10
|
+
prev = 0
|
|
11
|
+
for m in re.finditer(re.escape(separator), text):
|
|
12
|
+
spans.append((prev, m.start()))
|
|
13
|
+
prev = m.end()
|
|
14
|
+
spans.append((prev, len(text)))
|
|
15
|
+
return [(s, e) for s, e in spans if e > s]
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
class RecursiveTextSplitter:
|
|
19
|
+
def __init__(self, chunk_size: int, chunk_overlap: int) -> None:
|
|
20
|
+
self._chunk_size = chunk_size
|
|
21
|
+
self._chunk_overlap = chunk_overlap
|
|
22
|
+
|
|
23
|
+
def split_text_with_offsets(self, text: str) -> list[tuple[str, int]]:
|
|
24
|
+
return self._split(text, SEPARATORS, 0)
|
|
25
|
+
|
|
26
|
+
def _emit(self, text: str, start: int, end: int, base: int) -> list[tuple[str, int]]:
|
|
27
|
+
chunk = text[start:end]
|
|
28
|
+
s = start + len(chunk) - len(chunk.lstrip())
|
|
29
|
+
e = end - (len(chunk) - len(chunk.rstrip()))
|
|
30
|
+
if e <= s:
|
|
31
|
+
return []
|
|
32
|
+
return [(text[s:e], base + s)]
|
|
33
|
+
|
|
34
|
+
def _merge(self, spans: list[tuple[int, int]], separator: str) -> list[tuple[int, int]]:
|
|
35
|
+
joiner = len(separator)
|
|
36
|
+
chunks: list[tuple[int, int]] = []
|
|
37
|
+
current: list[tuple[int, int]] = []
|
|
38
|
+
total = 0
|
|
39
|
+
for span in spans:
|
|
40
|
+
size = span[1] - span[0]
|
|
41
|
+
if total + size + (joiner if current else 0) > self._chunk_size and current:
|
|
42
|
+
chunks.append((current[0][0], current[-1][1]))
|
|
43
|
+
while total > self._chunk_overlap or (
|
|
44
|
+
total + size + (joiner if current else 0) > self._chunk_size and total > 0
|
|
45
|
+
):
|
|
46
|
+
total -= current[0][1] - current[0][0] + (joiner if len(current) > 1 else 0)
|
|
47
|
+
current = current[1:]
|
|
48
|
+
current.append(span)
|
|
49
|
+
total += size + (joiner if len(current) > 1 else 0)
|
|
50
|
+
if current:
|
|
51
|
+
chunks.append((current[0][0], current[-1][1]))
|
|
52
|
+
return chunks
|
|
53
|
+
|
|
54
|
+
def _split(self, text: str, separators: list[str], base: int) -> list[tuple[str, int]]:
|
|
55
|
+
if not text:
|
|
56
|
+
return []
|
|
57
|
+
if len(text) <= self._chunk_size:
|
|
58
|
+
return self._emit(text, 0, len(text), base)
|
|
59
|
+
separator = separators[-1]
|
|
60
|
+
remaining: list[str] = []
|
|
61
|
+
for i, sep in enumerate(separators):
|
|
62
|
+
if not sep:
|
|
63
|
+
separator = sep
|
|
64
|
+
break
|
|
65
|
+
if sep in text:
|
|
66
|
+
separator = sep
|
|
67
|
+
remaining = separators[i + 1 :]
|
|
68
|
+
break
|
|
69
|
+
found: list[tuple[str, int]] = []
|
|
70
|
+
good: list[tuple[int, int]] = []
|
|
71
|
+
for span in _split(text, separator):
|
|
72
|
+
piece = text[span[0] : span[1]]
|
|
73
|
+
if len(piece) < self._chunk_size:
|
|
74
|
+
good.append(span)
|
|
75
|
+
continue
|
|
76
|
+
for s, e in self._merge(good, separator):
|
|
77
|
+
found.extend(self._emit(text, s, e, base))
|
|
78
|
+
good = []
|
|
79
|
+
if remaining:
|
|
80
|
+
found.extend(self._split(piece, remaining, base + span[0]))
|
|
81
|
+
elif len(piece) > self._chunk_size:
|
|
82
|
+
step = max(1, self._chunk_size - self._chunk_overlap)
|
|
83
|
+
for start in range(0, len(piece), step):
|
|
84
|
+
found.extend(
|
|
85
|
+
self._emit(
|
|
86
|
+
piece, start, min(start + self._chunk_size, len(piece)), base + span[0]
|
|
87
|
+
)
|
|
88
|
+
)
|
|
89
|
+
else:
|
|
90
|
+
found.extend(self._emit(text, span[0], span[1], base))
|
|
91
|
+
for s, e in self._merge(good, separator):
|
|
92
|
+
found.extend(self._emit(text, s, e, base))
|
|
93
|
+
return found
|
|
@@ -0,0 +1,110 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import re
|
|
4
|
+
from dataclasses import dataclass
|
|
5
|
+
|
|
6
|
+
_FENCE_OPEN_RE = re.compile(r"```[a-zA-Z0-9_+-]*")
|
|
7
|
+
_TABLE_RE = re.compile(
|
|
8
|
+
r"(?:^\|[^\n]+\|\n)"
|
|
9
|
+
r"(?:^\|[\s\-:|]+\|\n)"
|
|
10
|
+
r"(?:^\|[^\n]+\|\n?)+",
|
|
11
|
+
re.MULTILINE,
|
|
12
|
+
)
|
|
13
|
+
_HEADING_RE = re.compile(r"^(#{1,6})\s+(.+?)$", re.MULTILINE)
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
@dataclass
|
|
17
|
+
class AtomicRegion:
|
|
18
|
+
start: int
|
|
19
|
+
end: int
|
|
20
|
+
content: str
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
@dataclass
|
|
24
|
+
class HeadingSpan:
|
|
25
|
+
start: int
|
|
26
|
+
end: int
|
|
27
|
+
path: tuple[str, ...]
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def _iter_code_fences(text: str) -> list[AtomicRegion]:
|
|
31
|
+
lines = text.splitlines(keepends=True)
|
|
32
|
+
starts: list[int] = []
|
|
33
|
+
pos = 0
|
|
34
|
+
for line in lines:
|
|
35
|
+
starts.append(pos)
|
|
36
|
+
pos += len(line)
|
|
37
|
+
|
|
38
|
+
regions: list[AtomicRegion] = []
|
|
39
|
+
i = 0
|
|
40
|
+
n = len(lines)
|
|
41
|
+
while i < n:
|
|
42
|
+
if _FENCE_OPEN_RE.fullmatch(lines[i].rstrip("\r\n")):
|
|
43
|
+
j = i + 1
|
|
44
|
+
while j < n and not lines[j].startswith("```"):
|
|
45
|
+
j += 1
|
|
46
|
+
if j < n:
|
|
47
|
+
start = starts[i]
|
|
48
|
+
end = starts[j] + 3
|
|
49
|
+
regions.append(AtomicRegion(start, end, text[start:end]))
|
|
50
|
+
i = j + 1
|
|
51
|
+
continue
|
|
52
|
+
i += 1
|
|
53
|
+
return regions
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
def find_atomic_regions(text: str) -> list[AtomicRegion]:
|
|
57
|
+
regions: list[AtomicRegion] = []
|
|
58
|
+
regions.extend(_iter_code_fences(text))
|
|
59
|
+
for m in _TABLE_RE.finditer(text):
|
|
60
|
+
regions.append(AtomicRegion(m.start(), m.end(), m.group(0)))
|
|
61
|
+
regions.sort(key=lambda r: (r.start, -r.end))
|
|
62
|
+
merged: list[AtomicRegion] = []
|
|
63
|
+
for region in regions:
|
|
64
|
+
if merged and region.end <= merged[-1].end:
|
|
65
|
+
continue
|
|
66
|
+
if merged and region.start < merged[-1].end:
|
|
67
|
+
prev = merged[-1]
|
|
68
|
+
merged[-1] = AtomicRegion(prev.start, region.end, text[prev.start : region.end])
|
|
69
|
+
continue
|
|
70
|
+
merged.append(region)
|
|
71
|
+
return merged
|
|
72
|
+
|
|
73
|
+
|
|
74
|
+
def build_heading_spans(
|
|
75
|
+
text: str,
|
|
76
|
+
exclude_regions: list[AtomicRegion] | None = None,
|
|
77
|
+
) -> list[HeadingSpan]:
|
|
78
|
+
headings: list[tuple[int, int, str]] = []
|
|
79
|
+
exclude = exclude_regions or []
|
|
80
|
+
for m in _HEADING_RE.finditer(text):
|
|
81
|
+
offset = m.start()
|
|
82
|
+
if any(r.start <= offset < r.end for r in exclude):
|
|
83
|
+
continue
|
|
84
|
+
title = m.group(2).strip()
|
|
85
|
+
title = title.rstrip("#").rstrip() or title
|
|
86
|
+
headings.append((offset, len(m.group(1)), title))
|
|
87
|
+
if not headings:
|
|
88
|
+
return []
|
|
89
|
+
spans: list[HeadingSpan] = []
|
|
90
|
+
stack: list[tuple[int, str]] = []
|
|
91
|
+
for i, (offset, level, title) in enumerate(headings):
|
|
92
|
+
while stack and stack[-1][0] >= level:
|
|
93
|
+
stack.pop()
|
|
94
|
+
stack.append((level, title))
|
|
95
|
+
end = headings[i + 1][0] if i + 1 < len(headings) else len(text)
|
|
96
|
+
spans.append(
|
|
97
|
+
HeadingSpan(
|
|
98
|
+
start=offset,
|
|
99
|
+
end=end,
|
|
100
|
+
path=tuple(t for _, t in stack),
|
|
101
|
+
)
|
|
102
|
+
)
|
|
103
|
+
return spans
|
|
104
|
+
|
|
105
|
+
|
|
106
|
+
def section_path_at(offset: int, spans: list[HeadingSpan]) -> str | None:
|
|
107
|
+
for s in spans:
|
|
108
|
+
if s.start <= offset < s.end:
|
|
109
|
+
return " > ".join(s.path)
|
|
110
|
+
return None
|