post-graph 0.1.2__tar.gz → 0.1.4__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.
- {post_graph-0.1.2 → post_graph-0.1.4}/PKG-INFO +1 -1
- post_graph-0.1.4/demo_vector.py +90 -0
- {post_graph-0.1.2 → post_graph-0.1.4}/post_graph/__init__.py +1 -1
- {post_graph-0.1.2 → post_graph-0.1.4}/post_graph/client_asyncpg.py +299 -64
- {post_graph-0.1.2 → post_graph-0.1.4}/post_graph/client_sqlalchemy.py +263 -53
- {post_graph-0.1.2 → post_graph-0.1.4}/post_graph/models.py +21 -3
- {post_graph-0.1.2 → post_graph-0.1.4}/pyproject.toml +1 -1
- {post_graph-0.1.2 → post_graph-0.1.4}/.github/workflows/publish.yml +0 -0
- {post_graph-0.1.2 → post_graph-0.1.4}/.gitignore +0 -0
- {post_graph-0.1.2 → post_graph-0.1.4}/LICENSE +0 -0
- {post_graph-0.1.2 → post_graph-0.1.4}/README.md +0 -0
- {post_graph-0.1.2 → post_graph-0.1.4}/demo.py +0 -0
- {post_graph-0.1.2 → post_graph-0.1.4}/demo_schema_per_realm.py +0 -0
- {post_graph-0.1.2 → post_graph-0.1.4}/demo_transaction.py +0 -0
- {post_graph-0.1.2 → post_graph-0.1.4}/greek_gods.py +0 -0
- {post_graph-0.1.2 → post_graph-0.1.4}/post_graph/errors.py +0 -0
- {post_graph-0.1.2 → post_graph-0.1.4}/requirements.txt +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: post-graph
|
|
3
|
-
Version: 0.1.
|
|
3
|
+
Version: 0.1.4
|
|
4
4
|
Summary: High-performance PostgreSQL-backed graph database library supporting multi-tenant realms, schema-per-realm, shadow auditing, and append-only data history.
|
|
5
5
|
Project-URL: Homepage, https://github.com/crajah/post-graph
|
|
6
6
|
Project-URL: Repository, https://github.com/crajah/post-graph
|
|
@@ -0,0 +1,90 @@
|
|
|
1
|
+
"""Demo script verifying pgvector support and vector similarity search in post-graph across main and data tables."""
|
|
2
|
+
import asyncio
|
|
3
|
+
import os
|
|
4
|
+
from post_graph import AsyncPostGraph, TableNotFoundError
|
|
5
|
+
|
|
6
|
+
DB_URI = os.getenv("POSTGRES_URI", "postgresql://crajah@localhost:5432/postgres")
|
|
7
|
+
|
|
8
|
+
async def main():
|
|
9
|
+
print("=" * 60)
|
|
10
|
+
print("POST-GRAPH PGVECTOR DEMO (MAIN & DATA TABLE SEARCH)")
|
|
11
|
+
print("=" * 60)
|
|
12
|
+
|
|
13
|
+
client = AsyncPostGraph(dsn=DB_URI)
|
|
14
|
+
await client.connect()
|
|
15
|
+
|
|
16
|
+
table_name = "doc_chunks"
|
|
17
|
+
realm = "rag_demo"
|
|
18
|
+
|
|
19
|
+
table_ref = client._get_table_ref(table_name, realm)
|
|
20
|
+
audit_table_ref = client._get_table_ref(f"{table_name}_data_audit", realm)
|
|
21
|
+
audit_table_main = client._get_table_ref(f"{table_name}_audit", realm)
|
|
22
|
+
data_table_ref = client._get_table_ref(f"{table_name}_data", realm)
|
|
23
|
+
|
|
24
|
+
await client._execute(f"DROP TABLE IF EXISTS {data_table_ref} CASCADE;")
|
|
25
|
+
await client._execute(f"DROP TABLE IF EXISTS {audit_table_main} CASCADE;")
|
|
26
|
+
await client._execute(f"DROP TABLE IF EXISTS {table_ref} CASCADE;")
|
|
27
|
+
|
|
28
|
+
print("\n[+] Creating vertex table 'doc_chunks' with vector_dim=4...")
|
|
29
|
+
try:
|
|
30
|
+
await client.create_vertex_table(table_name, realm=realm, vector_dim=4)
|
|
31
|
+
print(" [OK] Table created successfully with vector_dim=4 for main & data tables")
|
|
32
|
+
except Exception as e:
|
|
33
|
+
print(f" [SKIP] pgvector extension not installed in Postgres: {e}")
|
|
34
|
+
await client.close()
|
|
35
|
+
return
|
|
36
|
+
|
|
37
|
+
print("\n[+] Inserting document chunk vertices into MAIN table with 4D embeddings...")
|
|
38
|
+
v1 = await client.add_vertex(
|
|
39
|
+
table_name, realm=realm,
|
|
40
|
+
payload={"text": "PostgreSQL is a relational database system."},
|
|
41
|
+
embedding=[1.0, 0.0, 0.0, 0.0]
|
|
42
|
+
)
|
|
43
|
+
print(f" Inserted V1 ID={v1.id}, embedding={v1.embedding}")
|
|
44
|
+
|
|
45
|
+
v2 = await client.add_vertex(
|
|
46
|
+
table_name, realm=realm,
|
|
47
|
+
payload={"text": "Post-graph enables graph storage on PostgreSQL."},
|
|
48
|
+
embedding=[0.9, 0.1, 0.0, 0.0]
|
|
49
|
+
)
|
|
50
|
+
print(f" Inserted V2 ID={v2.id}, embedding={v2.embedding}")
|
|
51
|
+
|
|
52
|
+
print("\n[+] Adding historical data records to DATA table with 4D embeddings for V1 & V2...")
|
|
53
|
+
d1 = await v1.add_data(
|
|
54
|
+
payload={"version": "v1.1", "note": "PostgreSQL 16 release notes"},
|
|
55
|
+
embedding=[0.0, 1.0, 0.0, 0.0]
|
|
56
|
+
)
|
|
57
|
+
print(f" Added V1 DataRecord ID={d1.data_id}, embedding={d1.embedding}")
|
|
58
|
+
|
|
59
|
+
d2 = await v2.add_data(
|
|
60
|
+
payload={"version": "v1.2", "note": "Post-graph pgvector extension release"},
|
|
61
|
+
embedding=[0.0, 0.95, 0.05, 0.0]
|
|
62
|
+
)
|
|
63
|
+
print(f" Added V2 DataRecord ID={d2.data_id}, embedding={d2.embedding}")
|
|
64
|
+
|
|
65
|
+
query_vec_main = [1.0, 0.05, 0.0, 0.0]
|
|
66
|
+
print(f"\n[+] 1. Vector Search [Scope: MAIN] for query {query_vec_main}...")
|
|
67
|
+
results_main = await client.vector_search(table_name, realm=realm, query_vector=query_vec_main, top_k=2, search_scope="main")
|
|
68
|
+
for rank, (vertex, dist) in enumerate(results_main, 1):
|
|
69
|
+
print(f" Rank {rank}: Vertex ID={vertex.id}, Distance={dist:.4f}, Text='{vertex.payload['text']}'")
|
|
70
|
+
|
|
71
|
+
query_vec_data = [0.0, 1.0, 0.0, 0.0]
|
|
72
|
+
print(f"\n[+] 2. Vector Search [Scope: DATA] for query {query_vec_data}...")
|
|
73
|
+
results_data = await client.vector_search(table_name, realm=realm, query_vector=query_vec_data, top_k=2, search_scope="data")
|
|
74
|
+
for rank, (vertex, dist) in enumerate(results_data, 1):
|
|
75
|
+
print(f" Rank {rank}: Vertex ID={vertex.id}, Distance={dist:.4f}, Text='{vertex.payload['text']}'")
|
|
76
|
+
|
|
77
|
+
print(f"\n[+] 3. Vector Search [Scope: BOTH] for query {query_vec_data}...")
|
|
78
|
+
results_both = await client.vector_search(table_name, realm=realm, query_vector=query_vec_data, top_k=2, search_scope="both")
|
|
79
|
+
for rank, (vertex, dist) in enumerate(results_both, 1):
|
|
80
|
+
print(f" Rank {rank}: Vertex ID={vertex.id}, Distance={dist:.4f}, Text='{vertex.payload['text']}'")
|
|
81
|
+
|
|
82
|
+
print("\n[+] Cleaning up demo table...")
|
|
83
|
+
await client._execute(f"DROP TABLE IF EXISTS {data_table_ref} CASCADE;")
|
|
84
|
+
await client._execute(f"DROP TABLE IF EXISTS {audit_table_main} CASCADE;")
|
|
85
|
+
await client._execute(f"DROP TABLE IF EXISTS {table_ref} CASCADE;")
|
|
86
|
+
await client.close()
|
|
87
|
+
print("[+] Done!")
|
|
88
|
+
|
|
89
|
+
if __name__ == "__main__":
|
|
90
|
+
asyncio.run(main())
|
|
@@ -10,7 +10,7 @@ from post_graph.models import Vertex, Edge, TraversalStep
|
|
|
10
10
|
from post_graph.client_asyncpg import AsyncPostGraph
|
|
11
11
|
from post_graph.client_sqlalchemy import SQLAlchemyPostGraph
|
|
12
12
|
|
|
13
|
-
__version__ = "0.1.
|
|
13
|
+
__version__ = "0.1.4"
|
|
14
14
|
|
|
15
15
|
__all__ = [
|
|
16
16
|
"__version__",
|
|
@@ -129,7 +129,12 @@ class AsyncPostGraph:
|
|
|
129
129
|
row = await self._fetchrow(query, table_name)
|
|
130
130
|
return row is not None
|
|
131
131
|
|
|
132
|
-
async def create_vertex_table(
|
|
132
|
+
async def create_vertex_table(
|
|
133
|
+
self,
|
|
134
|
+
table_name: str,
|
|
135
|
+
realm: Optional[str] = None,
|
|
136
|
+
vector_dim: Optional[int] = None
|
|
137
|
+
):
|
|
133
138
|
"""Create a new vertex table and its associated shadow audit table, indexes, and triggers."""
|
|
134
139
|
self._validate_identifier(table_name)
|
|
135
140
|
if self.schema_per_realm:
|
|
@@ -212,6 +217,16 @@ class AsyncPostGraph:
|
|
|
212
217
|
await self._execute(query)
|
|
213
218
|
await self._execute(f"ALTER TABLE {table_ref} ADD COLUMN IF NOT EXISTS fqid TEXT GENERATED ALWAYS AS (realm || '/' || '{table_name}' || '/' || id::text) STORED;")
|
|
214
219
|
|
|
220
|
+
if vector_dim and vector_dim > 0:
|
|
221
|
+
try:
|
|
222
|
+
await self._execute("CREATE EXTENSION IF NOT EXISTS vector;")
|
|
223
|
+
await self._execute(f"ALTER TABLE {table_ref} ADD COLUMN IF NOT EXISTS embedding vector({vector_dim});")
|
|
224
|
+
await self._execute(f'CREATE INDEX IF NOT EXISTS "idx_{table_name}_embedding" ON {table_ref} USING hnsw (embedding vector_cosine_ops);')
|
|
225
|
+
await self._execute(f"ALTER TABLE {data_table_ref} ADD COLUMN IF NOT EXISTS embedding vector({vector_dim});")
|
|
226
|
+
await self._execute(f'CREATE INDEX IF NOT EXISTS "idx_{table_name}_data_embedding" ON {data_table_ref} USING hnsw (embedding vector_cosine_ops);')
|
|
227
|
+
except Exception as e:
|
|
228
|
+
raise PostGraphError(f"Failed to initialize pgvector extension or embedding column for table '{table_name}': {e}")
|
|
229
|
+
|
|
215
230
|
# 2. Create shadow audit table
|
|
216
231
|
audit_query = f"""
|
|
217
232
|
CREATE TABLE IF NOT EXISTS {audit_table_ref} (
|
|
@@ -417,12 +432,14 @@ class AsyncPostGraph:
|
|
|
417
432
|
realm: str,
|
|
418
433
|
vertex_id: Optional[Union[str, int]] = None,
|
|
419
434
|
payload: Optional[Dict[str, Any]] = None,
|
|
435
|
+
embedding: Optional[List[float]] = None,
|
|
420
436
|
user_id: Optional[str] = None
|
|
421
437
|
) -> Vertex:
|
|
422
438
|
"""Add a new vertex. Raises TableExistsError if it already exists."""
|
|
423
439
|
self._validate_identifier(table_name)
|
|
424
440
|
payload_json = json.dumps(payload or {})
|
|
425
441
|
table_ref = self._get_table_ref(table_name, realm)
|
|
442
|
+
vec_str = f"[{','.join(str(x) for x in embedding)}]" if embedding else None
|
|
426
443
|
|
|
427
444
|
async def _op(conn):
|
|
428
445
|
nonlocal vertex_id
|
|
@@ -435,17 +452,36 @@ class AsyncPostGraph:
|
|
|
435
452
|
|
|
436
453
|
v_id_str = str(v_id_int)
|
|
437
454
|
|
|
438
|
-
|
|
439
|
-
|
|
440
|
-
|
|
441
|
-
|
|
442
|
-
|
|
443
|
-
|
|
455
|
+
has_emb = False
|
|
456
|
+
if vec_str:
|
|
457
|
+
if self.schema_per_realm:
|
|
458
|
+
has_emb = await conn.fetchval("SELECT 1 FROM information_schema.columns WHERE table_schema = $1 AND table_name = $2 AND column_name = 'embedding'", realm, table_name) is not None
|
|
459
|
+
else:
|
|
460
|
+
has_emb = await conn.fetchval("SELECT 1 FROM information_schema.columns WHERE table_name = $1 AND column_name = 'embedding'", table_name) is not None
|
|
461
|
+
|
|
462
|
+
if vec_str and has_emb:
|
|
463
|
+
query = f"""
|
|
464
|
+
INSERT INTO {table_ref} (realm, id, payload, embedding)
|
|
465
|
+
VALUES ($1, $2, $3::jsonb, $4::vector)
|
|
466
|
+
RETURNING realm, id, fqid, payload, created_at, updated_at, embedding::text AS embedding_text
|
|
467
|
+
"""
|
|
468
|
+
row = await conn.fetchrow(query, realm, v_id_int, payload_json, vec_str)
|
|
469
|
+
else:
|
|
470
|
+
query = f"""
|
|
471
|
+
INSERT INTO {table_ref} (realm, id, payload)
|
|
472
|
+
VALUES ($1, $2, $3::jsonb)
|
|
473
|
+
RETURNING realm, id, fqid, payload, created_at, updated_at
|
|
474
|
+
"""
|
|
444
475
|
row = await conn.fetchrow(query, realm, v_id_int, payload_json)
|
|
476
|
+
|
|
477
|
+
try:
|
|
445
478
|
if vertex_id is not None:
|
|
446
479
|
await conn.execute(
|
|
447
480
|
f"SELECT setval(pg_get_serial_sequence('{table_ref_pg}', 'id'), (SELECT COALESCE(MAX(id), 1) FROM {table_ref}))"
|
|
448
481
|
)
|
|
482
|
+
emb = None
|
|
483
|
+
if 'embedding_text' in row and row['embedding_text']:
|
|
484
|
+
emb = [float(x) for x in row['embedding_text'].strip('[]').split(',') if x.strip()]
|
|
449
485
|
return Vertex(
|
|
450
486
|
realm=row['realm'],
|
|
451
487
|
id=str(row['id']),
|
|
@@ -454,6 +490,7 @@ class AsyncPostGraph:
|
|
|
454
490
|
created_at=row['created_at'],
|
|
455
491
|
updated_at=row['updated_at'],
|
|
456
492
|
table_name=table_name,
|
|
493
|
+
embedding=emb,
|
|
457
494
|
_client=self
|
|
458
495
|
)
|
|
459
496
|
except asyncpg.UniqueViolationError:
|
|
@@ -471,12 +508,14 @@ class AsyncPostGraph:
|
|
|
471
508
|
realm: str,
|
|
472
509
|
vertex_id: Optional[Union[str, int]] = None,
|
|
473
510
|
payload: Optional[Dict[str, Any]] = None,
|
|
511
|
+
embedding: Optional[List[float]] = None,
|
|
474
512
|
user_id: Optional[str] = None
|
|
475
513
|
) -> Vertex:
|
|
476
514
|
"""Upsert a vertex (merges payload JSONB on conflict)."""
|
|
477
515
|
self._validate_identifier(table_name)
|
|
478
516
|
payload_json = json.dumps(payload or {})
|
|
479
517
|
table_ref = self._get_table_ref(table_name, realm)
|
|
518
|
+
vec_str = f"[{','.join(str(x) for x in embedding)}]" if embedding else None
|
|
480
519
|
|
|
481
520
|
async def _op(conn):
|
|
482
521
|
nonlocal vertex_id
|
|
@@ -489,16 +528,38 @@ class AsyncPostGraph:
|
|
|
489
528
|
|
|
490
529
|
v_id_str = str(v_id_int)
|
|
491
530
|
|
|
492
|
-
|
|
493
|
-
|
|
494
|
-
|
|
495
|
-
|
|
496
|
-
|
|
497
|
-
|
|
498
|
-
|
|
499
|
-
|
|
500
|
-
|
|
531
|
+
has_emb = False
|
|
532
|
+
if vec_str:
|
|
533
|
+
if self.schema_per_realm:
|
|
534
|
+
has_emb = await conn.fetchval("SELECT 1 FROM information_schema.columns WHERE table_schema = $1 AND table_name = $2 AND column_name = 'embedding'", realm, table_name) is not None
|
|
535
|
+
else:
|
|
536
|
+
has_emb = await conn.fetchval("SELECT 1 FROM information_schema.columns WHERE table_name = $1 AND column_name = 'embedding'", table_name) is not None
|
|
537
|
+
|
|
538
|
+
if vec_str and has_emb:
|
|
539
|
+
query = f"""
|
|
540
|
+
INSERT INTO {table_ref} (realm, id, payload, embedding)
|
|
541
|
+
VALUES ($1, $2, $3::jsonb, $4::vector)
|
|
542
|
+
ON CONFLICT (realm, id) DO UPDATE
|
|
543
|
+
SET payload = {table_ref}.payload || EXCLUDED.payload,
|
|
544
|
+
embedding = EXCLUDED.embedding,
|
|
545
|
+
updated_at = CURRENT_TIMESTAMP
|
|
546
|
+
RETURNING realm, id, fqid, payload, created_at, updated_at, embedding::text AS embedding_text
|
|
547
|
+
"""
|
|
548
|
+
row = await conn.fetchrow(query, realm, v_id_int, payload_json, vec_str)
|
|
549
|
+
else:
|
|
550
|
+
query = f"""
|
|
551
|
+
INSERT INTO {table_ref} (realm, id, payload)
|
|
552
|
+
VALUES ($1, $2, $3::jsonb)
|
|
553
|
+
ON CONFLICT (realm, id) DO UPDATE
|
|
554
|
+
SET payload = {table_ref}.payload || EXCLUDED.payload,
|
|
555
|
+
updated_at = CURRENT_TIMESTAMP
|
|
556
|
+
RETURNING realm, id, fqid, payload, created_at, updated_at
|
|
557
|
+
"""
|
|
501
558
|
row = await conn.fetchrow(query, realm, v_id_int, payload_json)
|
|
559
|
+
try:
|
|
560
|
+
emb = None
|
|
561
|
+
if 'embedding_text' in row and row['embedding_text']:
|
|
562
|
+
emb = [float(x) for x in row['embedding_text'].strip('[]').split(',') if x.strip()]
|
|
502
563
|
return Vertex(
|
|
503
564
|
realm=row['realm'],
|
|
504
565
|
id=str(row['id']),
|
|
@@ -507,6 +568,7 @@ class AsyncPostGraph:
|
|
|
507
568
|
created_at=row['created_at'],
|
|
508
569
|
updated_at=row['updated_at'],
|
|
509
570
|
table_name=table_name,
|
|
571
|
+
embedding=emb,
|
|
510
572
|
_client=self
|
|
511
573
|
)
|
|
512
574
|
except asyncpg.UndefinedTableError:
|
|
@@ -518,28 +580,163 @@ class AsyncPostGraph:
|
|
|
518
580
|
"""Fetch a vertex by realm and id."""
|
|
519
581
|
self._validate_identifier(table_name)
|
|
520
582
|
table_ref = self._get_table_ref(table_name, realm)
|
|
521
|
-
|
|
522
|
-
|
|
523
|
-
|
|
524
|
-
|
|
525
|
-
|
|
583
|
+
v_id_int = int(str(vertex_id).split('/')[-1]) if '/' in str(vertex_id) else int(vertex_id)
|
|
584
|
+
|
|
585
|
+
async def _op(conn):
|
|
586
|
+
query = f"""
|
|
587
|
+
SELECT t.realm, t.id, t.fqid, t.payload, t.created_at, t.updated_at,
|
|
588
|
+
to_jsonb(t)->>'embedding' AS embedding_text
|
|
589
|
+
FROM {table_ref} t
|
|
590
|
+
WHERE t.realm = $1 AND t.id = $2
|
|
591
|
+
"""
|
|
592
|
+
try:
|
|
593
|
+
row = await conn.fetchrow(query, realm, v_id_int)
|
|
594
|
+
if not row:
|
|
595
|
+
return None
|
|
596
|
+
|
|
597
|
+
emb = None
|
|
598
|
+
if 'embedding_text' in row and row['embedding_text']:
|
|
599
|
+
emb = [float(x) for x in row['embedding_text'].strip('[]').split(',') if x.strip()]
|
|
600
|
+
|
|
601
|
+
return Vertex(
|
|
602
|
+
realm=row['realm'],
|
|
603
|
+
id=str(row['id']),
|
|
604
|
+
fqid=row['fqid'],
|
|
605
|
+
payload=row['payload'] if isinstance(row['payload'], dict) else json.loads(row['payload']),
|
|
606
|
+
created_at=row['created_at'],
|
|
607
|
+
updated_at=row['updated_at'],
|
|
608
|
+
table_name=table_name,
|
|
609
|
+
embedding=emb,
|
|
610
|
+
_client=self
|
|
611
|
+
)
|
|
612
|
+
except asyncpg.UndefinedTableError:
|
|
613
|
+
raise TableNotFoundError(f"Vertex table '{table_name}' does not exist.")
|
|
614
|
+
|
|
615
|
+
if isinstance(self.connection, asyncpg.Pool):
|
|
616
|
+
async with self.connection.acquire() as conn:
|
|
617
|
+
return await _op(conn)
|
|
618
|
+
else:
|
|
619
|
+
return await _op(self.connection)
|
|
620
|
+
|
|
621
|
+
async def vector_search(
|
|
622
|
+
self,
|
|
623
|
+
table_name: str,
|
|
624
|
+
realm: str,
|
|
625
|
+
query_vector: List[float],
|
|
626
|
+
top_k: int = 5,
|
|
627
|
+
distance_metric: str = "cosine",
|
|
628
|
+
search_data_table: bool = False,
|
|
629
|
+
search_scope: str = "main"
|
|
630
|
+
) -> List[Tuple[Vertex, float]]:
|
|
631
|
+
"""Perform vector similarity search on vertex embeddings using pgvector.
|
|
632
|
+
|
|
633
|
+
Parameters:
|
|
634
|
+
- search_data_table: If True and search_scope is 'main', searches the associated data table.
|
|
635
|
+
- search_scope: 'main' (search main vertex table), 'data' (search associated data table), or 'both' (search both tables).
|
|
636
|
+
|
|
637
|
+
Distance metrics:
|
|
638
|
+
- 'cosine': <=> (cosine distance)
|
|
639
|
+
- 'l2': <-> (Euclidean distance)
|
|
640
|
+
- 'inner_product': <#> (negative inner product)
|
|
526
641
|
"""
|
|
527
|
-
|
|
528
|
-
|
|
529
|
-
|
|
530
|
-
|
|
531
|
-
|
|
532
|
-
|
|
533
|
-
|
|
534
|
-
|
|
535
|
-
|
|
536
|
-
|
|
537
|
-
|
|
538
|
-
|
|
539
|
-
|
|
642
|
+
self._validate_identifier(table_name)
|
|
643
|
+
table_ref = self._get_table_ref(table_name, realm)
|
|
644
|
+
data_table_ref = self._get_table_ref(f"{table_name}_data", realm)
|
|
645
|
+
vec_str = f"[{','.join(str(x) for x in query_vector)}]"
|
|
646
|
+
|
|
647
|
+
op = "<=>"
|
|
648
|
+
if distance_metric == "l2":
|
|
649
|
+
op = "<->"
|
|
650
|
+
elif distance_metric == "inner_product":
|
|
651
|
+
op = "<#>"
|
|
652
|
+
|
|
653
|
+
scope = search_scope.lower()
|
|
654
|
+
if search_data_table and scope == "main":
|
|
655
|
+
scope = "data"
|
|
656
|
+
|
|
657
|
+
if scope == "data":
|
|
658
|
+
query = f"""
|
|
659
|
+
SELECT v.realm, v.id, v.fqid, v.payload, v.created_at, v.updated_at,
|
|
660
|
+
to_jsonb(v)->>'embedding' AS embedding_text,
|
|
661
|
+
MIN(d.embedding {op} $2::vector) AS distance
|
|
662
|
+
FROM {data_table_ref} d
|
|
663
|
+
JOIN {table_ref} v ON d.realm = v.realm AND d.id = v.id
|
|
664
|
+
WHERE d.realm = $1 AND d.embedding IS NOT NULL
|
|
665
|
+
GROUP BY v.realm, v.id, v.fqid, v.payload, v.created_at, v.updated_at, to_jsonb(v)
|
|
666
|
+
ORDER BY distance ASC
|
|
667
|
+
LIMIT $3
|
|
668
|
+
"""
|
|
669
|
+
elif scope == "both":
|
|
670
|
+
query = f"""
|
|
671
|
+
WITH combined AS (
|
|
672
|
+
SELECT realm, id, (embedding {op} $2::vector) AS distance
|
|
673
|
+
FROM {table_ref}
|
|
674
|
+
WHERE realm = $1 AND embedding IS NOT NULL
|
|
675
|
+
|
|
676
|
+
UNION ALL
|
|
677
|
+
|
|
678
|
+
SELECT realm, id, (embedding {op} $2::vector) AS distance
|
|
679
|
+
FROM {data_table_ref}
|
|
680
|
+
WHERE realm = $1 AND embedding IS NOT NULL
|
|
681
|
+
),
|
|
682
|
+
best AS (
|
|
683
|
+
SELECT realm, id, MIN(distance) AS distance
|
|
684
|
+
FROM combined
|
|
685
|
+
GROUP BY realm, id
|
|
540
686
|
)
|
|
541
|
-
|
|
542
|
-
|
|
687
|
+
SELECT v.realm, v.id, v.fqid, v.payload, v.created_at, v.updated_at,
|
|
688
|
+
to_jsonb(v)->>'embedding' AS embedding_text,
|
|
689
|
+
b.distance
|
|
690
|
+
FROM best b
|
|
691
|
+
JOIN {table_ref} v ON b.realm = v.realm AND b.id = v.id
|
|
692
|
+
ORDER BY b.distance ASC
|
|
693
|
+
LIMIT $3
|
|
694
|
+
"""
|
|
695
|
+
else:
|
|
696
|
+
query = f"""
|
|
697
|
+
SELECT t.realm, t.id, t.fqid, t.payload, t.created_at, t.updated_at,
|
|
698
|
+
to_jsonb(t)->>'embedding' AS embedding_text,
|
|
699
|
+
(t.embedding {op} $2::vector) AS distance
|
|
700
|
+
FROM {table_ref} t
|
|
701
|
+
WHERE t.realm = $1 AND t.embedding IS NOT NULL
|
|
702
|
+
ORDER BY t.embedding {op} $2::vector ASC
|
|
703
|
+
LIMIT $3
|
|
704
|
+
"""
|
|
705
|
+
|
|
706
|
+
async def _op(conn):
|
|
707
|
+
try:
|
|
708
|
+
rows = await conn.fetch(query, realm, vec_str, top_k)
|
|
709
|
+
except asyncpg.UndefinedTableError:
|
|
710
|
+
raise TableNotFoundError(f"Vertex table '{table_name}' does not exist.")
|
|
711
|
+
except asyncpg.UndefinedColumnError:
|
|
712
|
+
logger.warning(f"Table '{table_name}' or data table does not have a vector embedding column.")
|
|
713
|
+
return []
|
|
714
|
+
|
|
715
|
+
results = []
|
|
716
|
+
for r in rows:
|
|
717
|
+
emb = None
|
|
718
|
+
if r['embedding_text']:
|
|
719
|
+
emb = [float(x) for x in r['embedding_text'].strip('[]').split(',') if x.strip()]
|
|
720
|
+
v = Vertex(
|
|
721
|
+
realm=r['realm'],
|
|
722
|
+
id=str(r['id']),
|
|
723
|
+
fqid=r['fqid'],
|
|
724
|
+
payload=r['payload'] if isinstance(r['payload'], dict) else json.loads(r['payload']),
|
|
725
|
+
created_at=r['created_at'],
|
|
726
|
+
updated_at=r['updated_at'],
|
|
727
|
+
table_name=table_name,
|
|
728
|
+
embedding=emb,
|
|
729
|
+
_client=self
|
|
730
|
+
)
|
|
731
|
+
dist = float(r['distance'])
|
|
732
|
+
results.append((v, dist))
|
|
733
|
+
return results
|
|
734
|
+
|
|
735
|
+
if isinstance(self.connection, asyncpg.Pool):
|
|
736
|
+
async with self.connection.acquire() as conn:
|
|
737
|
+
return await _op(conn)
|
|
738
|
+
else:
|
|
739
|
+
return await _op(self.connection)
|
|
543
740
|
|
|
544
741
|
async def delete_vertex(self, table_name: str, realm: str, vertex_id: str, user_id: Optional[str] = None) -> bool:
|
|
545
742
|
"""Delete a vertex. Cascading foreign keys will automatically delete referencing edges."""
|
|
@@ -1175,6 +1372,7 @@ class AsyncPostGraph:
|
|
|
1175
1372
|
vertex_id: Union[str, int],
|
|
1176
1373
|
payload: Dict[str, Any],
|
|
1177
1374
|
timestamp: Optional[Any] = None,
|
|
1375
|
+
embedding: Optional[List[float]] = None,
|
|
1178
1376
|
user_id: Optional[str] = None
|
|
1179
1377
|
) -> DataRecord:
|
|
1180
1378
|
"""Append a historical record to {table_name}_data table for a vertex."""
|
|
@@ -1182,29 +1380,59 @@ class AsyncPostGraph:
|
|
|
1182
1380
|
v_id_int = int(str(vertex_id).split('/')[-1]) if '/' in str(vertex_id) else int(vertex_id)
|
|
1183
1381
|
payload_json = json.dumps(payload or {})
|
|
1184
1382
|
data_table_ref = self._get_table_ref(f"{table_name}_data", realm)
|
|
1383
|
+
vec_str = f"[{','.join(str(x) for x in embedding)}]" if embedding is not None else None
|
|
1185
1384
|
|
|
1186
1385
|
async def _op(conn):
|
|
1187
|
-
|
|
1188
|
-
|
|
1189
|
-
|
|
1190
|
-
|
|
1191
|
-
|
|
1192
|
-
|
|
1193
|
-
|
|
1386
|
+
has_emb = False
|
|
1387
|
+
if vec_str:
|
|
1388
|
+
data_t_name = f"{table_name}_data"
|
|
1389
|
+
if self.schema_per_realm:
|
|
1390
|
+
has_emb = await conn.fetchval("SELECT 1 FROM information_schema.columns WHERE table_schema = $1 AND table_name = $2 AND column_name = 'embedding'", realm, data_t_name) is not None
|
|
1391
|
+
else:
|
|
1392
|
+
has_emb = await conn.fetchval("SELECT 1 FROM information_schema.columns WHERE table_name = $1 AND column_name = 'embedding'", data_t_name) is not None
|
|
1393
|
+
|
|
1394
|
+
if vec_str and has_emb:
|
|
1395
|
+
if timestamp:
|
|
1396
|
+
query = f"""
|
|
1397
|
+
INSERT INTO {data_table_ref} (realm, id, payload, timestamp, embedding)
|
|
1398
|
+
VALUES ($1, $2, $3::jsonb, $4, $5::vector)
|
|
1399
|
+
RETURNING data_id, realm, id, payload, timestamp, to_jsonb({data_table_ref})->>'embedding' AS embedding_text
|
|
1400
|
+
"""
|
|
1401
|
+
row = await conn.fetchrow(query, realm, v_id_int, payload_json, timestamp, vec_str)
|
|
1402
|
+
else:
|
|
1403
|
+
query = f"""
|
|
1404
|
+
INSERT INTO {data_table_ref} (realm, id, payload, embedding)
|
|
1405
|
+
VALUES ($1, $2, $3::jsonb, $4::vector)
|
|
1406
|
+
RETURNING data_id, realm, id, payload, timestamp, to_jsonb({data_table_ref})->>'embedding' AS embedding_text
|
|
1407
|
+
"""
|
|
1408
|
+
row = await conn.fetchrow(query, realm, v_id_int, payload_json, vec_str)
|
|
1194
1409
|
else:
|
|
1195
|
-
|
|
1196
|
-
|
|
1197
|
-
|
|
1198
|
-
|
|
1199
|
-
|
|
1200
|
-
|
|
1410
|
+
if timestamp:
|
|
1411
|
+
query = f"""
|
|
1412
|
+
INSERT INTO {data_table_ref} (realm, id, payload, timestamp)
|
|
1413
|
+
VALUES ($1, $2, $3::jsonb, $4)
|
|
1414
|
+
RETURNING data_id, realm, id, payload, timestamp, to_jsonb({data_table_ref})->>'embedding' AS embedding_text
|
|
1415
|
+
"""
|
|
1416
|
+
row = await conn.fetchrow(query, realm, v_id_int, payload_json, timestamp)
|
|
1417
|
+
else:
|
|
1418
|
+
query = f"""
|
|
1419
|
+
INSERT INTO {data_table_ref} (realm, id, payload)
|
|
1420
|
+
VALUES ($1, $2, $3::jsonb)
|
|
1421
|
+
RETURNING data_id, realm, id, payload, timestamp, to_jsonb({data_table_ref})->>'embedding' AS embedding_text
|
|
1422
|
+
"""
|
|
1423
|
+
row = await conn.fetchrow(query, realm, v_id_int, payload_json)
|
|
1424
|
+
|
|
1425
|
+
emb = None
|
|
1426
|
+
if 'embedding_text' in row and row['embedding_text']:
|
|
1427
|
+
emb = [float(x) for x in row['embedding_text'].strip('[]').split(',') if x.strip()]
|
|
1201
1428
|
|
|
1202
1429
|
return DataRecord(
|
|
1203
1430
|
data_id=str(row['data_id']),
|
|
1204
1431
|
realm=row['realm'],
|
|
1205
1432
|
id=str(row['id']),
|
|
1206
1433
|
payload=row['payload'] if isinstance(row['payload'], dict) else json.loads(row['payload']),
|
|
1207
|
-
timestamp=row['timestamp']
|
|
1434
|
+
timestamp=row['timestamp'],
|
|
1435
|
+
embedding=emb
|
|
1208
1436
|
)
|
|
1209
1437
|
|
|
1210
1438
|
return await self._run_in_tx(_op, user_id)
|
|
@@ -1223,23 +1451,29 @@ class AsyncPostGraph:
|
|
|
1223
1451
|
|
|
1224
1452
|
limit_clause = f"LIMIT {limit}" if limit and limit > 0 else ""
|
|
1225
1453
|
query = f"""
|
|
1226
|
-
SELECT data_id, realm, id, payload, timestamp
|
|
1227
|
-
FROM {data_table_ref}
|
|
1228
|
-
WHERE realm = $1 AND id = $2
|
|
1229
|
-
ORDER BY timestamp DESC, data_id DESC
|
|
1454
|
+
SELECT d.data_id, d.realm, d.id, d.payload, d.timestamp, to_jsonb(d)->>'embedding' AS embedding_text
|
|
1455
|
+
FROM {data_table_ref} d
|
|
1456
|
+
WHERE d.realm = $1 AND d.id = $2
|
|
1457
|
+
ORDER BY d.timestamp DESC, d.data_id DESC
|
|
1230
1458
|
{limit_clause}
|
|
1231
1459
|
"""
|
|
1232
1460
|
rows = await self._fetch(query, realm, v_id_int)
|
|
1233
|
-
|
|
1234
|
-
|
|
1235
|
-
|
|
1236
|
-
|
|
1237
|
-
|
|
1238
|
-
|
|
1239
|
-
|
|
1461
|
+
results = []
|
|
1462
|
+
for r in rows:
|
|
1463
|
+
emb = None
|
|
1464
|
+
if 'embedding_text' in r and r['embedding_text']:
|
|
1465
|
+
emb = [float(x) for x in r['embedding_text'].strip('[]').split(',') if x.strip()]
|
|
1466
|
+
results.append(
|
|
1467
|
+
DataRecord(
|
|
1468
|
+
data_id=str(r['data_id']),
|
|
1469
|
+
realm=r['realm'],
|
|
1470
|
+
id=str(r['id']),
|
|
1471
|
+
payload=r['payload'] if isinstance(r['payload'], dict) else json.loads(r['payload']),
|
|
1472
|
+
timestamp=r['timestamp'],
|
|
1473
|
+
embedding=emb
|
|
1474
|
+
)
|
|
1240
1475
|
)
|
|
1241
|
-
|
|
1242
|
-
]
|
|
1476
|
+
return results
|
|
1243
1477
|
|
|
1244
1478
|
async def add_edge_data(
|
|
1245
1479
|
self,
|
|
@@ -1248,10 +1482,11 @@ class AsyncPostGraph:
|
|
|
1248
1482
|
edge_id: Union[str, int],
|
|
1249
1483
|
payload: Dict[str, Any],
|
|
1250
1484
|
timestamp: Optional[Any] = None,
|
|
1485
|
+
embedding: Optional[List[float]] = None,
|
|
1251
1486
|
user_id: Optional[str] = None
|
|
1252
1487
|
) -> DataRecord:
|
|
1253
1488
|
"""Append a historical record to {table_name}_data table for an edge."""
|
|
1254
|
-
return await self.add_vertex_data(table_name, realm, edge_id, payload, timestamp=timestamp, user_id=user_id)
|
|
1489
|
+
return await self.add_vertex_data(table_name, realm, edge_id, payload, timestamp=timestamp, embedding=embedding, user_id=user_id)
|
|
1255
1490
|
|
|
1256
1491
|
async def get_edge_data(
|
|
1257
1492
|
self,
|
|
@@ -98,7 +98,12 @@ class SQLAlchemyPostGraph:
|
|
|
98
98
|
row = await self._fetchrow(conn, query, table_name=table_name)
|
|
99
99
|
return row is not None
|
|
100
100
|
|
|
101
|
-
async def create_vertex_table(
|
|
101
|
+
async def create_vertex_table(
|
|
102
|
+
self,
|
|
103
|
+
table_name: str,
|
|
104
|
+
realm: Optional[str] = None,
|
|
105
|
+
vector_dim: Optional[int] = None
|
|
106
|
+
):
|
|
102
107
|
"""Create a new vertex table and its associated shadow audit table, indexes, and triggers."""
|
|
103
108
|
self._validate_identifier(table_name)
|
|
104
109
|
if self.schema_per_realm:
|
|
@@ -184,6 +189,16 @@ class SQLAlchemyPostGraph:
|
|
|
184
189
|
await conn.execute(text(query))
|
|
185
190
|
await conn.execute(text(f"ALTER TABLE {table_ref} ADD COLUMN IF NOT EXISTS fqid TEXT GENERATED ALWAYS AS (realm || '/' || '{table_name}' || '/' || id::text) STORED;"))
|
|
186
191
|
|
|
192
|
+
if vector_dim and vector_dim > 0:
|
|
193
|
+
try:
|
|
194
|
+
await conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector;"))
|
|
195
|
+
await conn.execute(text(f"ALTER TABLE {table_ref} ADD COLUMN IF NOT EXISTS embedding vector({vector_dim});"))
|
|
196
|
+
await conn.execute(text(f'CREATE INDEX IF NOT EXISTS "idx_{table_name}_embedding" ON {table_ref} USING hnsw (embedding vector_cosine_ops);'))
|
|
197
|
+
await conn.execute(text(f"ALTER TABLE {data_table_ref} ADD COLUMN IF NOT EXISTS embedding vector({vector_dim});"))
|
|
198
|
+
await conn.execute(text(f'CREATE INDEX IF NOT EXISTS "idx_{table_name}_data_embedding" ON {data_table_ref} USING hnsw (embedding vector_cosine_ops);'))
|
|
199
|
+
except Exception as e:
|
|
200
|
+
raise PostGraphError(f"Failed to initialize pgvector extension or embedding column for table '{table_name}': {e}")
|
|
201
|
+
|
|
187
202
|
# 2. Create shadow audit table
|
|
188
203
|
audit_query = f"""
|
|
189
204
|
CREATE TABLE IF NOT EXISTS {audit_table_ref} (
|
|
@@ -381,12 +396,14 @@ class SQLAlchemyPostGraph:
|
|
|
381
396
|
realm: str,
|
|
382
397
|
vertex_id: Optional[Union[str, int]] = None,
|
|
383
398
|
payload: Optional[Dict[str, Any]] = None,
|
|
399
|
+
embedding: Optional[List[float]] = None,
|
|
384
400
|
user_id: Optional[str] = None
|
|
385
401
|
) -> Vertex:
|
|
386
402
|
"""Add a new vertex. Raises TableExistsError if it already exists."""
|
|
387
403
|
self._validate_identifier(table_name)
|
|
388
404
|
payload_json = json.dumps(payload or {})
|
|
389
405
|
table_ref = self._get_table_ref(table_name, realm)
|
|
406
|
+
vec_str = f"[{','.join(str(x) for x in embedding)}]" if embedding else None
|
|
390
407
|
|
|
391
408
|
async def _op(conn):
|
|
392
409
|
nonlocal vertex_id
|
|
@@ -399,17 +416,30 @@ class SQLAlchemyPostGraph:
|
|
|
399
416
|
|
|
400
417
|
v_id_str = str(v_id_int)
|
|
401
418
|
|
|
402
|
-
|
|
403
|
-
|
|
404
|
-
|
|
405
|
-
|
|
406
|
-
|
|
419
|
+
if vec_str:
|
|
420
|
+
query = f"""
|
|
421
|
+
INSERT INTO {table_ref} (realm, id, payload, embedding)
|
|
422
|
+
VALUES (:realm, :id, CAST(:payload AS JSONB), CAST(:vec AS vector))
|
|
423
|
+
RETURNING realm, id, fqid, payload, created_at, updated_at, CAST(embedding AS TEXT) AS embedding_text
|
|
424
|
+
"""
|
|
425
|
+
kwargs = {"realm": realm, "id": v_id_int, "payload": payload_json, "vec": vec_str}
|
|
426
|
+
else:
|
|
427
|
+
query = f"""
|
|
428
|
+
INSERT INTO {table_ref} (realm, id, payload)
|
|
429
|
+
VALUES (:realm, :id, CAST(:payload AS JSONB))
|
|
430
|
+
RETURNING realm, id, fqid, payload, created_at, updated_at
|
|
431
|
+
"""
|
|
432
|
+
kwargs = {"realm": realm, "id": v_id_int, "payload": payload_json}
|
|
433
|
+
|
|
407
434
|
try:
|
|
408
|
-
row = await self._fetchrow(conn, query,
|
|
435
|
+
row = await self._fetchrow(conn, query, **kwargs)
|
|
409
436
|
if vertex_id is not None:
|
|
410
437
|
await conn.execute(
|
|
411
438
|
text(f"SELECT setval(pg_get_serial_sequence('{table_ref_pg}', 'id'), (SELECT COALESCE(MAX(id), 1) FROM {table_ref}))")
|
|
412
439
|
)
|
|
440
|
+
emb = None
|
|
441
|
+
if 'embedding_text' in row and row['embedding_text']:
|
|
442
|
+
emb = [float(x) for x in row['embedding_text'].strip('[]').split(',') if x.strip()]
|
|
413
443
|
return Vertex(
|
|
414
444
|
realm=row['realm'],
|
|
415
445
|
id=str(row['id']),
|
|
@@ -418,6 +448,7 @@ class SQLAlchemyPostGraph:
|
|
|
418
448
|
created_at=row['created_at'],
|
|
419
449
|
updated_at=row['updated_at'],
|
|
420
450
|
table_name=table_name,
|
|
451
|
+
embedding=emb,
|
|
421
452
|
_client=self
|
|
422
453
|
)
|
|
423
454
|
except IntegrityError as e:
|
|
@@ -439,12 +470,14 @@ class SQLAlchemyPostGraph:
|
|
|
439
470
|
realm: str,
|
|
440
471
|
vertex_id: Optional[Union[str, int]] = None,
|
|
441
472
|
payload: Optional[Dict[str, Any]] = None,
|
|
473
|
+
embedding: Optional[List[float]] = None,
|
|
442
474
|
user_id: Optional[str] = None
|
|
443
475
|
) -> Vertex:
|
|
444
476
|
"""Upsert a vertex (merges payload JSONB on conflict)."""
|
|
445
477
|
self._validate_identifier(table_name)
|
|
446
478
|
payload_json = json.dumps(payload or {})
|
|
447
479
|
table_ref = self._get_table_ref(table_name, realm)
|
|
480
|
+
vec_str = f"[{','.join(str(x) for x in embedding)}]" if embedding else None
|
|
448
481
|
|
|
449
482
|
async def _op(conn):
|
|
450
483
|
nonlocal vertex_id
|
|
@@ -457,16 +490,33 @@ class SQLAlchemyPostGraph:
|
|
|
457
490
|
|
|
458
491
|
v_id_str = str(v_id_int)
|
|
459
492
|
|
|
460
|
-
|
|
461
|
-
|
|
462
|
-
|
|
463
|
-
|
|
464
|
-
|
|
465
|
-
|
|
466
|
-
|
|
467
|
-
|
|
493
|
+
if vec_str:
|
|
494
|
+
query = f"""
|
|
495
|
+
INSERT INTO {table_ref} (realm, id, payload, embedding)
|
|
496
|
+
VALUES (:realm, :id, CAST(:payload AS JSONB), CAST(:vec AS vector))
|
|
497
|
+
ON CONFLICT (realm, id) DO UPDATE
|
|
498
|
+
SET payload = {table_ref}.payload || EXCLUDED.payload,
|
|
499
|
+
embedding = EXCLUDED.embedding,
|
|
500
|
+
updated_at = CURRENT_TIMESTAMP
|
|
501
|
+
RETURNING realm, id, fqid, payload, created_at, updated_at, CAST(embedding AS TEXT) AS embedding_text
|
|
502
|
+
"""
|
|
503
|
+
kwargs = {"realm": realm, "id": v_id_int, "payload": payload_json, "vec": vec_str}
|
|
504
|
+
else:
|
|
505
|
+
query = f"""
|
|
506
|
+
INSERT INTO {table_ref} (realm, id, payload)
|
|
507
|
+
VALUES (:realm, :id, CAST(:payload AS JSONB))
|
|
508
|
+
ON CONFLICT (realm, id) DO UPDATE
|
|
509
|
+
SET payload = {table_ref}.payload || EXCLUDED.payload,
|
|
510
|
+
updated_at = CURRENT_TIMESTAMP
|
|
511
|
+
RETURNING realm, id, fqid, payload, created_at, updated_at
|
|
512
|
+
"""
|
|
513
|
+
kwargs = {"realm": realm, "id": v_id_int, "payload": payload_json}
|
|
514
|
+
|
|
468
515
|
try:
|
|
469
|
-
row = await self._fetchrow(conn, query,
|
|
516
|
+
row = await self._fetchrow(conn, query, **kwargs)
|
|
517
|
+
emb = None
|
|
518
|
+
if 'embedding_text' in row and row['embedding_text']:
|
|
519
|
+
emb = [float(x) for x in row['embedding_text'].strip('[]').split(',') if x.strip()]
|
|
470
520
|
return Vertex(
|
|
471
521
|
realm=row['realm'],
|
|
472
522
|
id=str(row['id']),
|
|
@@ -475,6 +525,7 @@ class SQLAlchemyPostGraph:
|
|
|
475
525
|
created_at=row['created_at'],
|
|
476
526
|
updated_at=row['updated_at'],
|
|
477
527
|
table_name=table_name,
|
|
528
|
+
embedding=emb,
|
|
478
529
|
_client=self
|
|
479
530
|
)
|
|
480
531
|
except ProgrammingError as e:
|
|
@@ -488,26 +539,33 @@ class SQLAlchemyPostGraph:
|
|
|
488
539
|
"""Fetch a vertex by realm and id."""
|
|
489
540
|
self._validate_identifier(table_name)
|
|
490
541
|
table_ref = self._get_table_ref(table_name, realm)
|
|
491
|
-
|
|
492
|
-
|
|
493
|
-
SELECT realm, id, fqid, payload, created_at, updated_at
|
|
494
|
-
FROM {table_ref}
|
|
495
|
-
WHERE realm = :realm AND ((CASE WHEN :id ~ '^[0-9]+$' THEN id = CAST(:id AS BIGINT) ELSE FALSE END) OR fqid = :id)
|
|
496
|
-
"""
|
|
497
|
-
|
|
542
|
+
v_id_int = int(str(vertex_id).split('/')[-1]) if '/' in str(vertex_id) else int(vertex_id)
|
|
543
|
+
|
|
498
544
|
async def _op(conn):
|
|
545
|
+
query = f"""
|
|
546
|
+
SELECT t.realm, t.id, t.fqid, t.payload, t.created_at, t.updated_at,
|
|
547
|
+
to_jsonb(t)->>'embedding' AS embedding_text
|
|
548
|
+
FROM {table_ref} t
|
|
549
|
+
WHERE t.realm = :realm AND t.id = :id
|
|
550
|
+
"""
|
|
499
551
|
try:
|
|
500
|
-
row = await self._fetchrow(conn, query, realm=realm, id=
|
|
552
|
+
row = await self._fetchrow(conn, query, realm=realm, id=v_id_int)
|
|
501
553
|
if not row:
|
|
502
554
|
return None
|
|
555
|
+
|
|
556
|
+
emb = None
|
|
557
|
+
if 'embedding_text' in row and row['embedding_text']:
|
|
558
|
+
emb = [float(x) for x in row['embedding_text'].strip('[]').split(',') if x.strip()]
|
|
559
|
+
|
|
503
560
|
return Vertex(
|
|
504
561
|
realm=row['realm'],
|
|
505
562
|
id=str(row['id']),
|
|
506
|
-
fqid=row['fqid']
|
|
563
|
+
fqid=row['fqid'],
|
|
507
564
|
payload=row['payload'] if isinstance(row['payload'], dict) else json.loads(row['payload']),
|
|
508
565
|
created_at=row['created_at'],
|
|
509
566
|
updated_at=row['updated_at'],
|
|
510
567
|
table_name=table_name,
|
|
568
|
+
embedding=emb,
|
|
511
569
|
_client=self
|
|
512
570
|
)
|
|
513
571
|
except ProgrammingError as e:
|
|
@@ -521,6 +579,116 @@ class SQLAlchemyPostGraph:
|
|
|
521
579
|
async with self.engine_or_connection.connect() as conn:
|
|
522
580
|
return await _op(conn)
|
|
523
581
|
|
|
582
|
+
async def vector_search(
|
|
583
|
+
self,
|
|
584
|
+
table_name: str,
|
|
585
|
+
realm: str,
|
|
586
|
+
query_vector: List[float],
|
|
587
|
+
top_k: int = 5,
|
|
588
|
+
distance_metric: str = "cosine",
|
|
589
|
+
search_data_table: bool = False,
|
|
590
|
+
search_scope: str = "main"
|
|
591
|
+
) -> List[Tuple[Vertex, float]]:
|
|
592
|
+
"""Perform vector similarity search on vertex embeddings using pgvector."""
|
|
593
|
+
self._validate_identifier(table_name)
|
|
594
|
+
table_ref = self._get_table_ref(table_name, realm)
|
|
595
|
+
data_table_ref = self._get_table_ref(f"{table_name}_data", realm)
|
|
596
|
+
vec_str = f"[{','.join(str(x) for x in query_vector)}]"
|
|
597
|
+
|
|
598
|
+
op = "<=>"
|
|
599
|
+
if distance_metric == "l2":
|
|
600
|
+
op = "<->"
|
|
601
|
+
elif distance_metric == "inner_product":
|
|
602
|
+
op = "<#>"
|
|
603
|
+
|
|
604
|
+
scope = search_scope.lower()
|
|
605
|
+
if search_data_table and scope == "main":
|
|
606
|
+
scope = "data"
|
|
607
|
+
|
|
608
|
+
if scope == "data":
|
|
609
|
+
query = f"""
|
|
610
|
+
SELECT v.realm, v.id, v.fqid, v.payload, v.created_at, v.updated_at,
|
|
611
|
+
to_jsonb(v)->>'embedding' AS embedding_text,
|
|
612
|
+
MIN(d.embedding {op} CAST(:vec AS vector)) AS distance
|
|
613
|
+
FROM {data_table_ref} d
|
|
614
|
+
JOIN {table_ref} v ON d.realm = v.realm AND d.id = v.id
|
|
615
|
+
WHERE d.realm = :realm AND d.embedding IS NOT NULL
|
|
616
|
+
GROUP BY v.realm, v.id, v.fqid, v.payload, v.created_at, v.updated_at, to_jsonb(v)
|
|
617
|
+
ORDER BY distance ASC
|
|
618
|
+
LIMIT :top_k
|
|
619
|
+
"""
|
|
620
|
+
elif scope == "both":
|
|
621
|
+
query = f"""
|
|
622
|
+
WITH combined AS (
|
|
623
|
+
SELECT realm, id, (embedding {op} CAST(:vec AS vector)) AS distance
|
|
624
|
+
FROM {table_ref}
|
|
625
|
+
WHERE realm = :realm AND embedding IS NOT NULL
|
|
626
|
+
|
|
627
|
+
UNION ALL
|
|
628
|
+
|
|
629
|
+
SELECT realm, id, (embedding {op} CAST(:vec AS vector)) AS distance
|
|
630
|
+
FROM {data_table_ref}
|
|
631
|
+
WHERE realm = :realm AND embedding IS NOT NULL
|
|
632
|
+
),
|
|
633
|
+
best AS (
|
|
634
|
+
SELECT realm, id, MIN(distance) AS distance
|
|
635
|
+
FROM combined
|
|
636
|
+
GROUP BY realm, id
|
|
637
|
+
)
|
|
638
|
+
SELECT v.realm, v.id, v.fqid, v.payload, v.created_at, v.updated_at,
|
|
639
|
+
to_jsonb(v)->>'embedding' AS embedding_text,
|
|
640
|
+
b.distance
|
|
641
|
+
FROM best b
|
|
642
|
+
JOIN {table_ref} v ON b.realm = v.realm AND b.id = v.id
|
|
643
|
+
ORDER BY b.distance ASC
|
|
644
|
+
LIMIT :top_k
|
|
645
|
+
"""
|
|
646
|
+
else:
|
|
647
|
+
query = f"""
|
|
648
|
+
SELECT t.realm, t.id, t.fqid, t.payload, t.created_at, t.updated_at,
|
|
649
|
+
CAST(t.embedding AS TEXT) AS embedding_text,
|
|
650
|
+
(t.embedding {op} CAST(:vec AS vector)) AS distance
|
|
651
|
+
FROM {table_ref} t
|
|
652
|
+
WHERE t.realm = :realm AND t.embedding IS NOT NULL
|
|
653
|
+
ORDER BY t.embedding {op} CAST(:vec AS vector) ASC
|
|
654
|
+
LIMIT :top_k
|
|
655
|
+
"""
|
|
656
|
+
|
|
657
|
+
async def _op(conn):
|
|
658
|
+
try:
|
|
659
|
+
rows = await self._fetch(conn, query, realm=realm, vec=vec_str, top_k=top_k)
|
|
660
|
+
except ProgrammingError as e:
|
|
661
|
+
if "does not exist" in str(e).lower():
|
|
662
|
+
raise TableNotFoundError(f"Vertex table '{table_name}' does not exist.")
|
|
663
|
+
logger.warning(f"Vector search failed: {e}")
|
|
664
|
+
return []
|
|
665
|
+
|
|
666
|
+
results = []
|
|
667
|
+
for r in rows:
|
|
668
|
+
emb = None
|
|
669
|
+
if r['embedding_text']:
|
|
670
|
+
emb = [float(x) for x in r['embedding_text'].strip('[]').split(',') if x.strip()]
|
|
671
|
+
v = Vertex(
|
|
672
|
+
realm=r['realm'],
|
|
673
|
+
id=str(r['id']),
|
|
674
|
+
fqid=r['fqid'],
|
|
675
|
+
payload=r['payload'] if isinstance(r['payload'], dict) else json.loads(r['payload']),
|
|
676
|
+
created_at=r['created_at'],
|
|
677
|
+
updated_at=r['updated_at'],
|
|
678
|
+
table_name=table_name,
|
|
679
|
+
embedding=emb,
|
|
680
|
+
_client=self
|
|
681
|
+
)
|
|
682
|
+
dist = float(r['distance'])
|
|
683
|
+
results.append((v, dist))
|
|
684
|
+
return results
|
|
685
|
+
|
|
686
|
+
if isinstance(self.engine_or_connection, AsyncConnection):
|
|
687
|
+
return await _op(self.engine_or_connection)
|
|
688
|
+
else:
|
|
689
|
+
async with self.engine_or_connection.connect() as conn:
|
|
690
|
+
return await _op(conn)
|
|
691
|
+
|
|
524
692
|
async def delete_vertex(self, table_name: str, realm: str, vertex_id: str, user_id: Optional[str] = None) -> bool:
|
|
525
693
|
"""Delete a vertex. Cascading foreign keys will automatically delete referencing edges."""
|
|
526
694
|
self._validate_identifier(table_name)
|
|
@@ -1209,6 +1377,7 @@ class SQLAlchemyPostGraph:
|
|
|
1209
1377
|
vertex_id: Union[str, int],
|
|
1210
1378
|
payload: Dict[str, Any],
|
|
1211
1379
|
timestamp: Optional[Any] = None,
|
|
1380
|
+
embedding: Optional[List[float]] = None,
|
|
1212
1381
|
user_id: Optional[str] = None
|
|
1213
1382
|
) -> DataRecord:
|
|
1214
1383
|
"""Append a historical record to {table_name}_data table for a vertex."""
|
|
@@ -1216,29 +1385,63 @@ class SQLAlchemyPostGraph:
|
|
|
1216
1385
|
v_id_int = int(str(vertex_id).split('/')[-1]) if '/' in str(vertex_id) else int(vertex_id)
|
|
1217
1386
|
payload_json = json.dumps(payload or {})
|
|
1218
1387
|
data_table_ref = self._get_table_ref(f"{table_name}_data", realm)
|
|
1388
|
+
vec_str = f"[{','.join(str(x) for x in embedding)}]" if embedding is not None else None
|
|
1219
1389
|
|
|
1220
1390
|
async def _op(conn):
|
|
1221
|
-
|
|
1222
|
-
|
|
1223
|
-
|
|
1224
|
-
|
|
1225
|
-
|
|
1226
|
-
|
|
1227
|
-
|
|
1391
|
+
has_emb = False
|
|
1392
|
+
if vec_str:
|
|
1393
|
+
data_t_name = f"{table_name}_data"
|
|
1394
|
+
if self.schema_per_realm:
|
|
1395
|
+
check_q = "SELECT 1 FROM information_schema.columns WHERE table_schema = :schema AND table_name = :t_name AND column_name = 'embedding'"
|
|
1396
|
+
res = await conn.execute(text(check_q), {"schema": realm, "t_name": data_t_name})
|
|
1397
|
+
has_emb = res.scalar() is not None
|
|
1398
|
+
else:
|
|
1399
|
+
check_q = "SELECT 1 FROM information_schema.columns WHERE table_name = :t_name AND column_name = 'embedding'"
|
|
1400
|
+
res = await conn.execute(text(check_q), {"t_name": data_t_name})
|
|
1401
|
+
has_emb = res.scalar() is not None
|
|
1402
|
+
|
|
1403
|
+
if vec_str and has_emb:
|
|
1404
|
+
if timestamp:
|
|
1405
|
+
query = f"""
|
|
1406
|
+
INSERT INTO {data_table_ref} (realm, id, payload, timestamp, embedding)
|
|
1407
|
+
VALUES (:realm, :id, CAST(:payload AS JSONB), :timestamp, CAST(:vec AS vector))
|
|
1408
|
+
RETURNING data_id, realm, id, payload, timestamp, to_jsonb({data_table_ref})->>'embedding' AS embedding_text
|
|
1409
|
+
"""
|
|
1410
|
+
row = await self._fetchrow(conn, query, realm=realm, id=v_id_int, payload=payload_json, timestamp=timestamp, vec=vec_str)
|
|
1411
|
+
else:
|
|
1412
|
+
query = f"""
|
|
1413
|
+
INSERT INTO {data_table_ref} (realm, id, payload, embedding)
|
|
1414
|
+
VALUES (:realm, :id, CAST(:payload AS JSONB), CAST(:vec AS vector))
|
|
1415
|
+
RETURNING data_id, realm, id, payload, timestamp, to_jsonb({data_table_ref})->>'embedding' AS embedding_text
|
|
1416
|
+
"""
|
|
1417
|
+
row = await self._fetchrow(conn, query, realm=realm, id=v_id_int, payload=payload_json, vec=vec_str)
|
|
1228
1418
|
else:
|
|
1229
|
-
|
|
1230
|
-
|
|
1231
|
-
|
|
1232
|
-
|
|
1233
|
-
|
|
1234
|
-
|
|
1419
|
+
if timestamp:
|
|
1420
|
+
query = f"""
|
|
1421
|
+
INSERT INTO {data_table_ref} (realm, id, payload, timestamp)
|
|
1422
|
+
VALUES (:realm, :id, CAST(:payload AS JSONB), :timestamp)
|
|
1423
|
+
RETURNING data_id, realm, id, payload, timestamp, to_jsonb({data_table_ref})->>'embedding' AS embedding_text
|
|
1424
|
+
"""
|
|
1425
|
+
row = await self._fetchrow(conn, query, realm=realm, id=v_id_int, payload=payload_json, timestamp=timestamp)
|
|
1426
|
+
else:
|
|
1427
|
+
query = f"""
|
|
1428
|
+
INSERT INTO {data_table_ref} (realm, id, payload)
|
|
1429
|
+
VALUES (:realm, :id, CAST(:payload AS JSONB))
|
|
1430
|
+
RETURNING data_id, realm, id, payload, timestamp, to_jsonb({data_table_ref})->>'embedding' AS embedding_text
|
|
1431
|
+
"""
|
|
1432
|
+
row = await self._fetchrow(conn, query, realm=realm, id=v_id_int, payload=payload_json)
|
|
1433
|
+
|
|
1434
|
+
emb = None
|
|
1435
|
+
if 'embedding_text' in row and row['embedding_text']:
|
|
1436
|
+
emb = [float(x) for x in row['embedding_text'].strip('[]').split(',') if x.strip()]
|
|
1235
1437
|
|
|
1236
1438
|
return DataRecord(
|
|
1237
1439
|
data_id=str(row['data_id']),
|
|
1238
1440
|
realm=row['realm'],
|
|
1239
1441
|
id=str(row['id']),
|
|
1240
1442
|
payload=row['payload'] if isinstance(row['payload'], dict) else json.loads(row['payload']),
|
|
1241
|
-
timestamp=row['timestamp']
|
|
1443
|
+
timestamp=row['timestamp'],
|
|
1444
|
+
embedding=emb
|
|
1242
1445
|
)
|
|
1243
1446
|
|
|
1244
1447
|
return await self._run_in_tx(_op, user_id)
|
|
@@ -1257,25 +1460,31 @@ class SQLAlchemyPostGraph:
|
|
|
1257
1460
|
|
|
1258
1461
|
limit_clause = f"LIMIT {limit}" if limit and limit > 0 else ""
|
|
1259
1462
|
query = f"""
|
|
1260
|
-
SELECT data_id, realm, id, payload, timestamp
|
|
1261
|
-
FROM {data_table_ref}
|
|
1262
|
-
WHERE realm = :realm AND id = :id
|
|
1263
|
-
ORDER BY timestamp DESC, data_id DESC
|
|
1463
|
+
SELECT d.data_id, d.realm, d.id, d.payload, d.timestamp, to_jsonb(d)->>'embedding' AS embedding_text
|
|
1464
|
+
FROM {data_table_ref} d
|
|
1465
|
+
WHERE d.realm = :realm AND d.id = :id
|
|
1466
|
+
ORDER BY d.timestamp DESC, d.data_id DESC
|
|
1264
1467
|
{limit_clause}
|
|
1265
1468
|
"""
|
|
1266
1469
|
|
|
1267
1470
|
async def _op(conn):
|
|
1268
1471
|
rows = await self._fetch(conn, query, realm=realm, id=v_id_int)
|
|
1269
|
-
|
|
1270
|
-
|
|
1271
|
-
|
|
1272
|
-
|
|
1273
|
-
|
|
1274
|
-
|
|
1275
|
-
|
|
1472
|
+
results = []
|
|
1473
|
+
for r in rows:
|
|
1474
|
+
emb = None
|
|
1475
|
+
if 'embedding_text' in r and r['embedding_text']:
|
|
1476
|
+
emb = [float(x) for x in r['embedding_text'].strip('[]').split(',') if x.strip()]
|
|
1477
|
+
results.append(
|
|
1478
|
+
DataRecord(
|
|
1479
|
+
data_id=str(r['data_id']),
|
|
1480
|
+
realm=r['realm'],
|
|
1481
|
+
id=str(r['id']),
|
|
1482
|
+
payload=r['payload'] if isinstance(r['payload'], dict) else json.loads(r['payload']),
|
|
1483
|
+
timestamp=r['timestamp'],
|
|
1484
|
+
embedding=emb
|
|
1485
|
+
)
|
|
1276
1486
|
)
|
|
1277
|
-
|
|
1278
|
-
]
|
|
1487
|
+
return results
|
|
1279
1488
|
|
|
1280
1489
|
if isinstance(self.engine_or_connection, AsyncConnection):
|
|
1281
1490
|
return await _op(self.engine_or_connection)
|
|
@@ -1290,10 +1499,11 @@ class SQLAlchemyPostGraph:
|
|
|
1290
1499
|
edge_id: Union[str, int],
|
|
1291
1500
|
payload: Dict[str, Any],
|
|
1292
1501
|
timestamp: Optional[Any] = None,
|
|
1502
|
+
embedding: Optional[List[float]] = None,
|
|
1293
1503
|
user_id: Optional[str] = None
|
|
1294
1504
|
) -> DataRecord:
|
|
1295
1505
|
"""Append a historical record to {table_name}_data table for an edge."""
|
|
1296
|
-
return await self.add_vertex_data(table_name, realm, edge_id, payload, timestamp=timestamp, user_id=user_id)
|
|
1506
|
+
return await self.add_vertex_data(table_name, realm, edge_id, payload, timestamp=timestamp, embedding=embedding, user_id=user_id)
|
|
1297
1507
|
|
|
1298
1508
|
async def get_edge_data(
|
|
1299
1509
|
self,
|
|
@@ -12,6 +12,7 @@ class Vertex:
|
|
|
12
12
|
updated_at: Optional[datetime] = None
|
|
13
13
|
table_name: Optional[str] = None
|
|
14
14
|
fqid: Optional[str] = None
|
|
15
|
+
embedding: Optional[List[float]] = None
|
|
15
16
|
_client: Optional[Any] = field(default=None, repr=False, compare=False)
|
|
16
17
|
|
|
17
18
|
def __post_init__(self):
|
|
@@ -27,7 +28,8 @@ class Vertex:
|
|
|
27
28
|
"payload": self.payload,
|
|
28
29
|
"created_at": self.created_at.isoformat() if self.created_at else None,
|
|
29
30
|
"updated_at": self.updated_at.isoformat() if self.updated_at else None,
|
|
30
|
-
"table_name": self.table_name
|
|
31
|
+
"table_name": self.table_name,
|
|
32
|
+
"embedding": self.embedding
|
|
31
33
|
}
|
|
32
34
|
|
|
33
35
|
async def to(self, edge_table: str, direction: str = 'out') -> List['TraversalStep']:
|
|
@@ -117,7 +119,13 @@ class Vertex:
|
|
|
117
119
|
user_id=user_id
|
|
118
120
|
)
|
|
119
121
|
|
|
120
|
-
async def add_data(
|
|
122
|
+
async def add_data(
|
|
123
|
+
self,
|
|
124
|
+
payload: Dict[str, Any],
|
|
125
|
+
timestamp: Optional[datetime] = None,
|
|
126
|
+
embedding: Optional[List[float]] = None,
|
|
127
|
+
user_id: Optional[str] = None
|
|
128
|
+
) -> 'DataRecord':
|
|
121
129
|
"""Append a data record to this vertex's {table_name}_data table."""
|
|
122
130
|
if not self._client:
|
|
123
131
|
from post_graph.errors import PostGraphError
|
|
@@ -128,6 +136,7 @@ class Vertex:
|
|
|
128
136
|
vertex_id=self.id,
|
|
129
137
|
payload=payload,
|
|
130
138
|
timestamp=timestamp,
|
|
139
|
+
embedding=embedding,
|
|
131
140
|
user_id=user_id
|
|
132
141
|
)
|
|
133
142
|
|
|
@@ -189,7 +198,13 @@ class Edge:
|
|
|
189
198
|
user_id=user_id
|
|
190
199
|
)
|
|
191
200
|
|
|
192
|
-
async def add_data(
|
|
201
|
+
async def add_data(
|
|
202
|
+
self,
|
|
203
|
+
payload: Dict[str, Any],
|
|
204
|
+
timestamp: Optional[datetime] = None,
|
|
205
|
+
embedding: Optional[List[float]] = None,
|
|
206
|
+
user_id: Optional[str] = None
|
|
207
|
+
) -> 'DataRecord':
|
|
193
208
|
"""Append a data record to this edge's {table_name}_data table."""
|
|
194
209
|
if not self._client:
|
|
195
210
|
from post_graph.errors import PostGraphError
|
|
@@ -200,6 +215,7 @@ class Edge:
|
|
|
200
215
|
edge_id=self.id,
|
|
201
216
|
payload=payload,
|
|
202
217
|
timestamp=timestamp,
|
|
218
|
+
embedding=embedding,
|
|
203
219
|
user_id=user_id
|
|
204
220
|
)
|
|
205
221
|
|
|
@@ -223,6 +239,7 @@ class DataRecord:
|
|
|
223
239
|
id: str
|
|
224
240
|
payload: Dict[str, Any] = field(default_factory=dict)
|
|
225
241
|
timestamp: Optional[datetime] = None
|
|
242
|
+
embedding: Optional[List[float]] = None
|
|
226
243
|
|
|
227
244
|
def to_dict(self) -> Dict[str, Any]:
|
|
228
245
|
"""Convert DataRecord object to a dictionary."""
|
|
@@ -232,6 +249,7 @@ class DataRecord:
|
|
|
232
249
|
"id": self.id,
|
|
233
250
|
"payload": self.payload,
|
|
234
251
|
"timestamp": self.timestamp.isoformat() if self.timestamp else None,
|
|
252
|
+
"embedding": self.embedding
|
|
235
253
|
}
|
|
236
254
|
|
|
237
255
|
|
|
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
|
|
4
4
|
|
|
5
5
|
[project]
|
|
6
6
|
name = "post-graph"
|
|
7
|
-
version = "0.1.
|
|
7
|
+
version = "0.1.4"
|
|
8
8
|
description = "High-performance PostgreSQL-backed graph database library supporting multi-tenant realms, schema-per-realm, shadow auditing, and append-only data history."
|
|
9
9
|
readme = "README.md"
|
|
10
10
|
requires-python = ">=3.9"
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|