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.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: post-graph
3
- Version: 0.1.2
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.2"
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(self, table_name: str, realm: Optional[str] = None):
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
- query = f"""
439
- INSERT INTO {table_ref} (realm, id, payload)
440
- VALUES ($1, $2, $3::jsonb)
441
- RETURNING realm, id, fqid, payload, created_at, updated_at
442
- """
443
- try:
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
- query = f"""
493
- INSERT INTO {table_ref} (realm, id, payload)
494
- VALUES ($1, $2, $3::jsonb)
495
- ON CONFLICT (realm, id) DO UPDATE
496
- SET payload = {table_ref}.payload || EXCLUDED.payload,
497
- updated_at = CURRENT_TIMESTAMP
498
- RETURNING realm, id, fqid, payload, created_at, updated_at
499
- """
500
- try:
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
- v_str = str(vertex_id)
522
- query = f"""
523
- SELECT realm, id, fqid, payload, created_at, updated_at
524
- FROM {table_ref}
525
- WHERE realm = $1 AND ((CASE WHEN $2 ~ '^[0-9]+$' THEN id = $2::bigint ELSE FALSE END) OR fqid = $2)
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
- try:
528
- row = await self._fetchrow(query, realm, v_str)
529
- if not row:
530
- return None
531
- return Vertex(
532
- realm=row['realm'],
533
- id=str(row['id']),
534
- fqid=row['fqid'] or f"{row['realm']}/{table_name}/{row['id']}",
535
- payload=row['payload'] if isinstance(row['payload'], dict) else json.loads(row['payload']),
536
- created_at=row['created_at'],
537
- updated_at=row['updated_at'],
538
- table_name=table_name,
539
- _client=self
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
- except asyncpg.UndefinedTableError:
542
- raise TableNotFoundError(f"Vertex table '{table_name}' does not exist.")
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
- if timestamp:
1188
- query = f"""
1189
- INSERT INTO {data_table_ref} (realm, id, payload, timestamp)
1190
- VALUES ($1, $2, $3::jsonb, $4)
1191
- RETURNING data_id, realm, id, payload, timestamp
1192
- """
1193
- row = await conn.fetchrow(query, realm, v_id_int, payload_json, timestamp)
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
- query = f"""
1196
- INSERT INTO {data_table_ref} (realm, id, payload)
1197
- VALUES ($1, $2, $3::jsonb)
1198
- RETURNING data_id, realm, id, payload, timestamp
1199
- """
1200
- row = await conn.fetchrow(query, realm, v_id_int, payload_json)
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
- return [
1234
- DataRecord(
1235
- data_id=str(r['data_id']),
1236
- realm=r['realm'],
1237
- id=str(r['id']),
1238
- payload=r['payload'] if isinstance(r['payload'], dict) else json.loads(r['payload']),
1239
- timestamp=r['timestamp']
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
- for r in rows
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(self, table_name: str, realm: Optional[str] = None):
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
- query = f"""
403
- INSERT INTO {table_ref} (realm, id, payload)
404
- VALUES (:realm, :id, CAST(:payload AS JSONB))
405
- RETURNING realm, id, fqid, payload, created_at, updated_at
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, realm=realm, id=v_id_int, payload=payload_json)
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
- query = f"""
461
- INSERT INTO {table_ref} (realm, id, payload)
462
- VALUES (:realm, :id, CAST(:payload AS JSONB))
463
- ON CONFLICT (realm, id) DO UPDATE
464
- SET payload = {table_ref}.payload || EXCLUDED.payload,
465
- updated_at = CURRENT_TIMESTAMP
466
- RETURNING realm, id, fqid, payload, created_at, updated_at
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, realm=realm, id=v_id_int, payload=payload_json)
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
- v_str = str(vertex_id)
492
- query = f"""
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=v_str)
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'] or f"{row['realm']}/{table_name}/{row['id']}",
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
- if timestamp:
1222
- query = f"""
1223
- INSERT INTO {data_table_ref} (realm, id, payload, timestamp)
1224
- VALUES (:realm, :id, CAST(:payload AS JSONB), :timestamp)
1225
- RETURNING data_id, realm, id, payload, timestamp
1226
- """
1227
- row = await self._fetchrow(conn, query, realm=realm, id=v_id_int, payload=payload_json, timestamp=timestamp)
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
- query = f"""
1230
- INSERT INTO {data_table_ref} (realm, id, payload)
1231
- VALUES (:realm, :id, CAST(:payload AS JSONB))
1232
- RETURNING data_id, realm, id, payload, timestamp
1233
- """
1234
- row = await self._fetchrow(conn, query, realm=realm, id=v_id_int, payload=payload_json)
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
- return [
1270
- DataRecord(
1271
- data_id=str(r['data_id']),
1272
- realm=r['realm'],
1273
- id=str(r['id']),
1274
- payload=r['payload'] if isinstance(r['payload'], dict) else json.loads(r['payload']),
1275
- timestamp=r['timestamp']
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
- for r in rows
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(self, payload: Dict[str, Any], timestamp: Optional[datetime] = None, user_id: Optional[str] = None) -> 'DataRecord':
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(self, payload: Dict[str, Any], timestamp: Optional[datetime] = None, user_id: Optional[str] = None) -> 'DataRecord':
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.2"
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