hindsight-api 0.2.0__tar.gz → 0.3.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.
Files changed (89) hide show
  1. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/.gitignore +11 -2
  2. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/PKG-INFO +5 -1
  3. hindsight_api-0.3.0/hindsight_api/admin/__init__.py +1 -0
  4. hindsight_api-0.3.0/hindsight_api/admin/cli.py +252 -0
  5. hindsight_api-0.3.0/hindsight_api/alembic/versions/f1a2b3c4d5e6_add_memory_links_composite_index.py +44 -0
  6. hindsight_api-0.3.0/hindsight_api/alembic/versions/g2a3b4c5d6e7_add_tags_column.py +48 -0
  7. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/api/http.py +282 -20
  8. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/api/mcp.py +47 -52
  9. hindsight_api-0.3.0/hindsight_api/config.py +452 -0
  10. hindsight_api-0.3.0/hindsight_api/engine/cross_encoder.py +815 -0
  11. hindsight_api-0.3.0/hindsight_api/engine/db_budget.py +284 -0
  12. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/db_utils.py +11 -0
  13. hindsight_api-0.3.0/hindsight_api/engine/embeddings.py +720 -0
  14. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/entity_resolver.py +8 -5
  15. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/interface.py +8 -4
  16. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/llm_wrapper.py +241 -27
  17. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/memory_engine.py +609 -122
  18. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/query_analyzer.py +4 -3
  19. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/response_models.py +38 -0
  20. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/retain/fact_extraction.py +388 -192
  21. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/retain/fact_storage.py +34 -8
  22. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/retain/link_utils.py +24 -16
  23. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/retain/orchestrator.py +52 -17
  24. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/retain/types.py +9 -0
  25. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/search/graph_retrieval.py +42 -13
  26. hindsight_api-0.3.0/hindsight_api/engine/search/link_expansion_retrieval.py +256 -0
  27. hindsight_api-0.3.0/hindsight_api/engine/search/mpfp_retrieval.py +684 -0
  28. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/search/reranking.py +2 -2
  29. hindsight_api-0.3.0/hindsight_api/engine/search/retrieval.py +1346 -0
  30. hindsight_api-0.3.0/hindsight_api/engine/search/tags.py +172 -0
  31. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/search/think_utils.py +1 -1
  32. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/search/trace.py +12 -0
  33. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/search/tracer.py +24 -1
  34. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/search/types.py +21 -0
  35. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/task_backend.py +109 -18
  36. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/utils.py +1 -1
  37. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/extensions/context.py +10 -1
  38. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/main.py +56 -4
  39. hindsight_api-0.3.0/hindsight_api/metrics.py +626 -0
  40. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/migrations.py +141 -1
  41. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/models.py +3 -1
  42. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/pg0.py +53 -0
  43. hindsight_api-0.3.0/hindsight_api/server.py +77 -0
  44. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/pyproject.toml +14 -5
  45. hindsight_api-0.2.0/hindsight_api/config.py +0 -220
  46. hindsight_api-0.2.0/hindsight_api/engine/cross_encoder.py +0 -302
  47. hindsight_api-0.2.0/hindsight_api/engine/embeddings.py +0 -293
  48. hindsight_api-0.2.0/hindsight_api/engine/search/mpfp_retrieval.py +0 -439
  49. hindsight_api-0.2.0/hindsight_api/engine/search/retrieval.py +0 -699
  50. hindsight_api-0.2.0/hindsight_api/metrics.py +0 -241
  51. hindsight_api-0.2.0/hindsight_api/server.py +0 -40
  52. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/README.md +0 -0
  53. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/__init__.py +0 -0
  54. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/alembic/README +0 -0
  55. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/alembic/env.py +0 -0
  56. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/alembic/script.py.mako +0 -0
  57. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/alembic/versions/5a366d414dce_initial_schema.py +0 -0
  58. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/alembic/versions/b7c4d8e9f1a2_add_chunks_table.py +0 -0
  59. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/alembic/versions/c8e5f2a3b4d1_add_retain_params_to_documents.py +0 -0
  60. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/alembic/versions/d9f6a3b4c5e2_rename_bank_to_interactions.py +0 -0
  61. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/alembic/versions/e0a1b2c3d4e5_disposition_to_3_traits.py +0 -0
  62. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/alembic/versions/rename_personality_to_disposition.py +0 -0
  63. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/api/__init__.py +0 -0
  64. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/banner.py +0 -0
  65. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/daemon.py +0 -0
  66. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/__init__.py +0 -0
  67. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/retain/__init__.py +0 -0
  68. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/retain/bank_utils.py +0 -0
  69. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/retain/chunk_storage.py +0 -0
  70. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/retain/deduplication.py +0 -0
  71. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/retain/embedding_processing.py +0 -0
  72. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/retain/embedding_utils.py +0 -0
  73. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/retain/entity_processing.py +0 -0
  74. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/retain/link_creation.py +0 -0
  75. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/retain/observation_regeneration.py +0 -0
  76. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/search/__init__.py +0 -0
  77. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/search/fusion.py +0 -0
  78. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/search/observation_utils.py +0 -0
  79. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/search/scoring.py +0 -0
  80. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/engine/search/temporal_extraction.py +0 -0
  81. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/extensions/__init__.py +0 -0
  82. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/extensions/base.py +0 -0
  83. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/extensions/builtin/__init__.py +0 -0
  84. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/extensions/builtin/tenant.py +0 -0
  85. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/extensions/http.py +0 -0
  86. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/extensions/loader.py +0 -0
  87. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/extensions/operation_validator.py +0 -0
  88. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/extensions/tenant.py +0 -0
  89. {hindsight_api-0.2.0 → hindsight_api-0.3.0}/hindsight_api/mcp_local.py +0 -0
@@ -5,7 +5,8 @@ build/
5
5
  dist/
6
6
  wheels/
7
7
  *.egg-info
8
-
8
+ .mcp.json
9
+ .osgrep
9
10
  # Virtual environments
10
11
  .venv
11
12
 
@@ -26,6 +27,10 @@ docker-compose.override.yml
26
27
  # NLTK data (will be downloaded automatically)
27
28
  nltk_data/
28
29
 
30
+ # Monitoring stack (Prometheus/Grafana binaries and data)
31
+ .monitoring/
32
+ .pgbouncer
33
+
29
34
  # Large benchmark datasets (will be downloaded automatically)
30
35
  **/longmemeval_s_cleaned.json
31
36
 
@@ -41,4 +46,8 @@ hindsight-docs/static/llms-full.txt
41
46
  hindsight-dev/benchmarks/locomo/results/
42
47
  hindsight-dev/benchmarks/longmemeval/results/
43
48
  hindsight-cli/target
44
- hindsight-clients/rust/target
49
+ hindsight-clients/rust/target
50
+ .claude
51
+ whats-next.md
52
+ TASK.md
53
+ CHANGELOG.md
@@ -1,14 +1,16 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: hindsight-api
3
- Version: 0.2.0
3
+ Version: 0.3.0
4
4
  Summary: Hindsight: Agent Memory That Works Like Human Memory
5
5
  Requires-Python: >=3.11
6
6
  Requires-Dist: alembic>=1.17.1
7
7
  Requires-Dist: anthropic>=0.40.0
8
8
  Requires-Dist: asyncpg>=0.29.0
9
+ Requires-Dist: cohere>=5.0.0
9
10
  Requires-Dist: dateparser>=1.2.2
10
11
  Requires-Dist: fastapi[standard]>=0.120.3
11
12
  Requires-Dist: fastmcp>=2.3.0
13
+ Requires-Dist: flashrank>=0.2.0
12
14
  Requires-Dist: google-genai>=1.0.0
13
15
  Requires-Dist: greenlet>=3.2.4
14
16
  Requires-Dist: httpx>=0.27.0
@@ -30,7 +32,9 @@ Requires-Dist: sqlalchemy>=2.0.44
30
32
  Requires-Dist: tiktoken>=0.12.0
31
33
  Requires-Dist: torch>=2.0.0
32
34
  Requires-Dist: transformers<4.46.0,>=4.30.0
35
+ Requires-Dist: typer>=0.9.0
33
36
  Requires-Dist: uvicorn>=0.38.0
37
+ Requires-Dist: uvloop>=0.22.1
34
38
  Requires-Dist: wsproto>=1.0.0
35
39
  Provides-Extra: test
36
40
  Requires-Dist: filelock>=3.0.0; extra == 'test'
@@ -0,0 +1 @@
1
+ # Admin CLI for Hindsight
@@ -0,0 +1,252 @@
1
+ """
2
+ Hindsight Admin CLI - backup and restore operations.
3
+ """
4
+
5
+ import asyncio
6
+ import io
7
+ import json
8
+ import logging
9
+ import zipfile
10
+ from datetime import datetime, timezone
11
+ from pathlib import Path
12
+ from typing import Any
13
+
14
+ import asyncpg
15
+ import typer
16
+
17
+ from ..config import HindsightConfig
18
+ from ..pg0 import parse_pg0_url, resolve_database_url
19
+
20
+
21
+ def _fq_table(table: str, schema: str) -> str:
22
+ """Get fully-qualified table name with schema prefix."""
23
+ return f"{schema}.{table}"
24
+
25
+
26
+ # Setup logging
27
+ logging.basicConfig(
28
+ level=logging.INFO,
29
+ format="%(message)s",
30
+ )
31
+ logger = logging.getLogger(__name__)
32
+
33
+ app = typer.Typer(name="hindsight-admin", help="Hindsight administrative commands")
34
+
35
+ # Tables to backup/restore in dependency order
36
+ # Import must happen in this order due to foreign key constraints
37
+ BACKUP_TABLES = [
38
+ "banks",
39
+ "documents",
40
+ "entities",
41
+ "chunks",
42
+ "memory_units",
43
+ "unit_entities",
44
+ "entity_cooccurrences",
45
+ "memory_links",
46
+ ]
47
+
48
+ MANIFEST_VERSION = "1"
49
+
50
+
51
+ async def _backup(database_url: str, output_path: Path, schema: str = "public") -> dict[str, Any]:
52
+ """Backup all tables to a zip file using binary COPY protocol."""
53
+ conn = await asyncpg.connect(database_url)
54
+ try:
55
+ tables: dict[str, Any] = {}
56
+ manifest: dict[str, Any] = {
57
+ "version": MANIFEST_VERSION,
58
+ "created_at": datetime.now(timezone.utc).isoformat(),
59
+ "schema": schema,
60
+ "tables": tables,
61
+ }
62
+
63
+ # Use a transaction with REPEATABLE READ isolation to get a consistent
64
+ # snapshot across all tables. This prevents race conditions where
65
+ # entity_cooccurrences could reference entities created after the
66
+ # entities table was backed up.
67
+ async with conn.transaction(isolation="repeatable_read"):
68
+ with zipfile.ZipFile(output_path, "w", zipfile.ZIP_DEFLATED) as zf:
69
+ for i, table in enumerate(BACKUP_TABLES, 1):
70
+ typer.echo(f" [{i}/{len(BACKUP_TABLES)}] Backing up {table}...", nl=False)
71
+
72
+ buffer = io.BytesIO()
73
+
74
+ # Use binary COPY for exact type preservation
75
+ # asyncpg requires schema_name as separate parameter
76
+ await conn.copy_from_table(table, schema_name=schema, output=buffer, format="binary")
77
+
78
+ data = buffer.getvalue()
79
+ zf.writestr(f"{table}.bin", data)
80
+
81
+ # Get row count for manifest
82
+ qualified_table = _fq_table(table, schema)
83
+ row_count = await conn.fetchval(f"SELECT COUNT(*) FROM {qualified_table}")
84
+ tables[table] = {
85
+ "rows": row_count,
86
+ "size_bytes": len(data),
87
+ }
88
+
89
+ typer.echo(f" {row_count} rows")
90
+
91
+ zf.writestr("manifest.json", json.dumps(manifest, indent=2))
92
+
93
+ return manifest
94
+ finally:
95
+ await conn.close()
96
+
97
+
98
+ async def _restore(database_url: str, input_path: Path, schema: str = "public") -> dict[str, Any]:
99
+ """Restore all tables from a zip file using binary COPY protocol."""
100
+ conn = await asyncpg.connect(database_url)
101
+ try:
102
+ with zipfile.ZipFile(input_path, "r") as zf:
103
+ # Read and validate manifest
104
+ manifest: dict[str, Any] = json.loads(zf.read("manifest.json"))
105
+ if manifest.get("version") != MANIFEST_VERSION:
106
+ raise ValueError(f"Unsupported backup version: {manifest.get('version')}")
107
+
108
+ # Use a transaction for atomic restore - either all tables are
109
+ # restored or none are, preventing partial/inconsistent state.
110
+ async with conn.transaction():
111
+ typer.echo(" Clearing existing data...")
112
+ # Truncate tables in reverse order (respects FK constraints)
113
+ for table in reversed(BACKUP_TABLES):
114
+ qualified_table = _fq_table(table, schema)
115
+ await conn.execute(f"TRUNCATE TABLE {qualified_table} CASCADE")
116
+
117
+ # Restore tables in forward order
118
+ for i, table in enumerate(BACKUP_TABLES, 1):
119
+ filename = f"{table}.bin"
120
+ if filename not in zf.namelist():
121
+ typer.echo(f" [{i}/{len(BACKUP_TABLES)}] {table}: skipped (not in backup)")
122
+ continue
123
+
124
+ expected_rows = manifest["tables"].get(table, {}).get("rows", "?")
125
+ typer.echo(f" [{i}/{len(BACKUP_TABLES)}] Restoring {table}... {expected_rows} rows")
126
+
127
+ data = zf.read(filename)
128
+ buffer = io.BytesIO(data)
129
+ # asyncpg requires schema_name as separate parameter
130
+ await conn.copy_to_table(table, schema_name=schema, source=buffer, format="binary")
131
+
132
+ # Refresh materialized view
133
+ typer.echo(" Refreshing materialized views...")
134
+ await conn.execute(f"REFRESH MATERIALIZED VIEW {_fq_table('memory_units_bm25', schema)}")
135
+
136
+ return manifest
137
+ finally:
138
+ await conn.close()
139
+
140
+
141
+ async def _run_backup(db_url: str, output: Path, schema: str = "public") -> dict[str, Any]:
142
+ """Resolve database URL and run backup."""
143
+ is_pg0, instance_name, _ = parse_pg0_url(db_url)
144
+ if is_pg0:
145
+ typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
146
+ resolved_url = await resolve_database_url(db_url)
147
+ return await _backup(resolved_url, output, schema)
148
+
149
+
150
+ async def _run_restore(db_url: str, input_file: Path, schema: str = "public") -> dict[str, Any]:
151
+ """Resolve database URL and run restore."""
152
+ is_pg0, instance_name, _ = parse_pg0_url(db_url)
153
+ if is_pg0:
154
+ typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
155
+ resolved_url = await resolve_database_url(db_url)
156
+ return await _restore(resolved_url, input_file, schema)
157
+
158
+
159
+ @app.command()
160
+ def backup(
161
+ output: Path = typer.Argument(..., help="Output file path (.zip)"),
162
+ schema: str = typer.Option("public", "--schema", "-s", help="Database schema to backup"),
163
+ ):
164
+ """Backup the Hindsight database to a zip file."""
165
+ config = HindsightConfig.from_env()
166
+
167
+ if not config.database_url:
168
+ typer.echo("Error: Database URL not configured.", err=True)
169
+ typer.echo("Set HINDSIGHT_API_DATABASE_URL environment variable.", err=True)
170
+ raise typer.Exit(1)
171
+
172
+ if output.suffix != ".zip":
173
+ output = output.with_suffix(".zip")
174
+
175
+ typer.echo(f"Backing up database (schema: {schema}) to {output}...")
176
+
177
+ manifest = asyncio.run(_run_backup(config.database_url, output, schema))
178
+
179
+ total_rows = sum(t["rows"] for t in manifest["tables"].values())
180
+ typer.echo(f"Backed up {total_rows} rows across {len(BACKUP_TABLES)} tables")
181
+ typer.echo(f"Backup saved to {output}")
182
+
183
+
184
+ @app.command()
185
+ def restore(
186
+ input_file: Path = typer.Argument(..., help="Input backup file (.zip)"),
187
+ schema: str = typer.Option("public", "--schema", "-s", help="Database schema to restore to"),
188
+ yes: bool = typer.Option(False, "--yes", "-y", help="Skip confirmation prompt"),
189
+ ):
190
+ """Restore the database from a backup file. WARNING: This deletes all existing data."""
191
+ config = HindsightConfig.from_env()
192
+
193
+ if not config.database_url:
194
+ typer.echo("Error: Database URL not configured.", err=True)
195
+ typer.echo("Set HINDSIGHT_API_DATABASE_URL environment variable.", err=True)
196
+ raise typer.Exit(1)
197
+
198
+ if not input_file.exists():
199
+ typer.echo(f"Error: File not found: {input_file}", err=True)
200
+ raise typer.Exit(1)
201
+
202
+ if not yes:
203
+ typer.confirm(
204
+ "This will DELETE all existing data and replace it with the backup. Continue?",
205
+ abort=True,
206
+ )
207
+
208
+ typer.echo(f"Restoring database (schema: {schema}) from {input_file}...")
209
+
210
+ manifest = asyncio.run(_run_restore(config.database_url, input_file, schema))
211
+
212
+ total_rows = sum(t["rows"] for t in manifest["tables"].values())
213
+ typer.echo(f"Restored {total_rows} rows across {len(BACKUP_TABLES)} tables")
214
+ typer.echo("Restore complete")
215
+
216
+
217
+ async def _run_migration(db_url: str, schema: str = "public") -> None:
218
+ """Resolve database URL and run migrations."""
219
+ from ..migrations import run_migrations
220
+
221
+ is_pg0, instance_name, _ = parse_pg0_url(db_url)
222
+ if is_pg0:
223
+ typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
224
+ resolved_url = await resolve_database_url(db_url)
225
+ run_migrations(resolved_url, schema=schema)
226
+
227
+
228
+ @app.command(name="run-db-migration")
229
+ def run_db_migration(
230
+ schema: str = typer.Option("public", "--schema", "-s", help="Database schema to run migrations on"),
231
+ ):
232
+ """Run database migrations to the latest version."""
233
+ config = HindsightConfig.from_env()
234
+
235
+ if not config.database_url:
236
+ typer.echo("Error: Database URL not configured.", err=True)
237
+ typer.echo("Set HINDSIGHT_API_DATABASE_URL environment variable.", err=True)
238
+ raise typer.Exit(1)
239
+
240
+ typer.echo(f"Running database migrations (schema: {schema})...")
241
+
242
+ asyncio.run(_run_migration(config.database_url, schema))
243
+
244
+ typer.echo("Database migrations completed successfully")
245
+
246
+
247
+ def main():
248
+ app()
249
+
250
+
251
+ if __name__ == "__main__":
252
+ main()
@@ -0,0 +1,44 @@
1
+ """add_memory_links_from_type_weight_index
2
+
3
+ Revision ID: f1a2b3c4d5e6
4
+ Revises: e0a1b2c3d4e5
5
+ Create Date: 2025-01-12
6
+
7
+ Add composite index on memory_links (from_unit_id, link_type, weight DESC)
8
+ to optimize MPFP graph traversal queries that need top-k edges per type.
9
+ """
10
+
11
+ from collections.abc import Sequence
12
+
13
+ from alembic import context, op
14
+
15
+ # revision identifiers, used by Alembic.
16
+ revision: str = "f1a2b3c4d5e6"
17
+ down_revision: str | Sequence[str] | None = "e0a1b2c3d4e5"
18
+ branch_labels: str | Sequence[str] | None = None
19
+ depends_on: str | Sequence[str] | None = None
20
+
21
+
22
+ def _get_schema_prefix() -> str:
23
+ """Get schema prefix for table names (e.g., 'tenant_x.' or '' for public)."""
24
+ schema = context.config.get_main_option("target_schema")
25
+ return f'"{schema}".' if schema else ""
26
+
27
+
28
+ def upgrade() -> None:
29
+ """Add composite index for efficient MPFP edge loading."""
30
+ schema = _get_schema_prefix()
31
+ # Create composite index for efficient top-k per (from_node, link_type) queries
32
+ # This enables LATERAL joins to use index-only scans with early termination
33
+ # Note: Not using CONCURRENTLY here as it requires running outside a transaction
34
+ # For production with large tables, consider running this manually with CONCURRENTLY
35
+ op.execute(
36
+ f"CREATE INDEX IF NOT EXISTS idx_memory_links_from_type_weight "
37
+ f"ON {schema}memory_links(from_unit_id, link_type, weight DESC)"
38
+ )
39
+
40
+
41
+ def downgrade() -> None:
42
+ """Remove the composite index."""
43
+ schema = _get_schema_prefix()
44
+ op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_links_from_type_weight")
@@ -0,0 +1,48 @@
1
+ """add_tags_column
2
+
3
+ Revision ID: g2a3b4c5d6e7
4
+ Revises: f1a2b3c4d5e6
5
+ Create Date: 2025-01-13
6
+
7
+ Add tags column to memory_units and documents tables for visibility scoping.
8
+ Tags enable filtering memories by scope (e.g., user IDs, session IDs) during recall/reflect.
9
+ """
10
+
11
+ from collections.abc import Sequence
12
+
13
+ from alembic import context, op
14
+
15
+ # revision identifiers, used by Alembic.
16
+ revision: str = "g2a3b4c5d6e7"
17
+ down_revision: str | Sequence[str] | None = "f1a2b3c4d5e6"
18
+ branch_labels: str | Sequence[str] | None = None
19
+ depends_on: str | Sequence[str] | None = None
20
+
21
+
22
+ def _get_schema_prefix() -> str:
23
+ """Get schema prefix for table names (e.g., 'tenant_x.' or '' for public)."""
24
+ schema = context.config.get_main_option("target_schema")
25
+ return f'"{schema}".' if schema else ""
26
+
27
+
28
+ def upgrade() -> None:
29
+ """Add tags column to memory_units and documents tables."""
30
+ schema = _get_schema_prefix()
31
+
32
+ # Add tags column to memory_units table
33
+ op.execute(f"ALTER TABLE {schema}memory_units ADD COLUMN IF NOT EXISTS tags VARCHAR[] NOT NULL DEFAULT '{{}}'")
34
+
35
+ # Create GIN index for efficient array containment queries (tags && ARRAY['x'])
36
+ op.execute(f"CREATE INDEX IF NOT EXISTS idx_memory_units_tags ON {schema}memory_units USING GIN (tags)")
37
+
38
+ # Add tags column to documents table for document-level tags
39
+ op.execute(f"ALTER TABLE {schema}documents ADD COLUMN IF NOT EXISTS tags VARCHAR[] NOT NULL DEFAULT '{{}}'")
40
+
41
+
42
+ def downgrade() -> None:
43
+ """Remove tags columns and index."""
44
+ schema = _get_schema_prefix()
45
+
46
+ op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_units_tags")
47
+ op.execute(f"ALTER TABLE {schema}memory_units DROP COLUMN IF EXISTS tags")
48
+ op.execute(f"ALTER TABLE {schema}documents DROP COLUMN IF EXISTS tags")