post-graph 0.1.2__tar.gz → 0.1.3__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.3
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,73 @@
1
+ """Demo script verifying pgvector support and vector similarity search in post-graph."""
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")
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}_audit", realm)
21
+ data_table_ref = client._get_table_ref(f"{table_name}_data", realm)
22
+
23
+ await client._execute(f"DROP TABLE IF EXISTS {data_table_ref} CASCADE;")
24
+ await client._execute(f"DROP TABLE IF EXISTS {audit_table_ref} CASCADE;")
25
+ await client._execute(f"DROP TABLE IF EXISTS {table_ref} CASCADE;")
26
+
27
+ print("\n[+] Creating vertex table 'doc_chunks' with vector_dim=4...")
28
+ try:
29
+ await client.create_vertex_table(table_name, realm=realm, vector_dim=4)
30
+ print(" [OK] Table created successfully with vector_dim=4")
31
+ except Exception as e:
32
+ print(f" [SKIP] pgvector extension not installed in Postgres: {e}")
33
+ await client.close()
34
+ return
35
+
36
+ print("\n[+] Inserting document chunk vertices with 4D embeddings...")
37
+ v1 = await client.add_vertex(
38
+ table_name, realm=realm,
39
+ payload={"text": "PostgreSQL is a relational database system."},
40
+ embedding=[1.0, 0.0, 0.0, 0.0]
41
+ )
42
+ print(f" Inserted V1 ID={v1.id}, embedding={v1.embedding}")
43
+
44
+ v2 = await client.add_vertex(
45
+ table_name, realm=realm,
46
+ payload={"text": "Post-graph enables graph storage on PostgreSQL."},
47
+ embedding=[0.9, 0.1, 0.0, 0.0]
48
+ )
49
+ print(f" Inserted V2 ID={v2.id}, embedding={v2.embedding}")
50
+
51
+ v3 = await client.add_vertex(
52
+ table_name, realm=realm,
53
+ payload={"text": "Deep Learning models extract knowledge graphs."},
54
+ embedding=[0.0, 0.0, 1.0, 0.0]
55
+ )
56
+ print(f" Inserted V3 ID={v3.id}, embedding={v3.embedding}")
57
+
58
+ print("\n[+] Performing Cosine Vector Similarity Search for query [1.0, 0.05, 0.0, 0.0]...")
59
+ query_vec = [1.0, 0.05, 0.0, 0.0]
60
+ results = await client.vector_search(table_name, realm=realm, query_vector=query_vec, top_k=3, distance_metric="cosine")
61
+
62
+ for rank, (vertex, distance) in enumerate(results, 1):
63
+ print(f" Rank {rank}: ID={vertex.id}, Cosine Distance={distance:.4f}, Text='{vertex.payload['text']}'")
64
+
65
+ print("\n[+] Cleaning up demo table...")
66
+ await client._execute(f"DROP TABLE IF EXISTS {data_table_ref} CASCADE;")
67
+ await client._execute(f"DROP TABLE IF EXISTS {audit_table_ref} CASCADE;")
68
+ await client._execute(f"DROP TABLE IF EXISTS {table_ref} CASCADE;")
69
+ await client.close()
70
+ print("[+] Done!")
71
+
72
+ if __name__ == "__main__":
73
+ 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.3"
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,14 @@ 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
+ except Exception as e:
226
+ raise PostGraphError(f"Failed to initialize pgvector extension or embedding column for table '{table_name}': {e}")
227
+
215
228
  # 2. Create shadow audit table
216
229
  audit_query = f"""
217
230
  CREATE TABLE IF NOT EXISTS {audit_table_ref} (
@@ -417,12 +430,14 @@ class AsyncPostGraph:
417
430
  realm: str,
418
431
  vertex_id: Optional[Union[str, int]] = None,
419
432
  payload: Optional[Dict[str, Any]] = None,
433
+ embedding: Optional[List[float]] = None,
420
434
  user_id: Optional[str] = None
421
435
  ) -> Vertex:
422
436
  """Add a new vertex. Raises TableExistsError if it already exists."""
423
437
  self._validate_identifier(table_name)
424
438
  payload_json = json.dumps(payload or {})
425
439
  table_ref = self._get_table_ref(table_name, realm)
440
+ vec_str = f"[{','.join(str(x) for x in embedding)}]" if embedding else None
426
441
 
427
442
  async def _op(conn):
428
443
  nonlocal vertex_id
@@ -435,17 +450,36 @@ class AsyncPostGraph:
435
450
 
436
451
  v_id_str = str(v_id_int)
437
452
 
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:
453
+ has_emb = False
454
+ if vec_str:
455
+ if self.schema_per_realm:
456
+ 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
457
+ else:
458
+ has_emb = await conn.fetchval("SELECT 1 FROM information_schema.columns WHERE table_name = $1 AND column_name = 'embedding'", table_name) is not None
459
+
460
+ if vec_str and has_emb:
461
+ query = f"""
462
+ INSERT INTO {table_ref} (realm, id, payload, embedding)
463
+ VALUES ($1, $2, $3::jsonb, $4::vector)
464
+ RETURNING realm, id, fqid, payload, created_at, updated_at, embedding::text AS embedding_text
465
+ """
466
+ row = await conn.fetchrow(query, realm, v_id_int, payload_json, vec_str)
467
+ else:
468
+ query = f"""
469
+ INSERT INTO {table_ref} (realm, id, payload)
470
+ VALUES ($1, $2, $3::jsonb)
471
+ RETURNING realm, id, fqid, payload, created_at, updated_at
472
+ """
444
473
  row = await conn.fetchrow(query, realm, v_id_int, payload_json)
474
+
475
+ try:
445
476
  if vertex_id is not None:
446
477
  await conn.execute(
447
478
  f"SELECT setval(pg_get_serial_sequence('{table_ref_pg}', 'id'), (SELECT COALESCE(MAX(id), 1) FROM {table_ref}))"
448
479
  )
480
+ emb = None
481
+ if 'embedding_text' in row and row['embedding_text']:
482
+ emb = [float(x) for x in row['embedding_text'].strip('[]').split(',') if x.strip()]
449
483
  return Vertex(
450
484
  realm=row['realm'],
451
485
  id=str(row['id']),
@@ -454,6 +488,7 @@ class AsyncPostGraph:
454
488
  created_at=row['created_at'],
455
489
  updated_at=row['updated_at'],
456
490
  table_name=table_name,
491
+ embedding=emb,
457
492
  _client=self
458
493
  )
459
494
  except asyncpg.UniqueViolationError:
@@ -471,12 +506,14 @@ class AsyncPostGraph:
471
506
  realm: str,
472
507
  vertex_id: Optional[Union[str, int]] = None,
473
508
  payload: Optional[Dict[str, Any]] = None,
509
+ embedding: Optional[List[float]] = None,
474
510
  user_id: Optional[str] = None
475
511
  ) -> Vertex:
476
512
  """Upsert a vertex (merges payload JSONB on conflict)."""
477
513
  self._validate_identifier(table_name)
478
514
  payload_json = json.dumps(payload or {})
479
515
  table_ref = self._get_table_ref(table_name, realm)
516
+ vec_str = f"[{','.join(str(x) for x in embedding)}]" if embedding else None
480
517
 
481
518
  async def _op(conn):
482
519
  nonlocal vertex_id
@@ -489,16 +526,38 @@ class AsyncPostGraph:
489
526
 
490
527
  v_id_str = str(v_id_int)
491
528
 
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:
529
+ has_emb = False
530
+ if vec_str:
531
+ if self.schema_per_realm:
532
+ 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
533
+ else:
534
+ has_emb = await conn.fetchval("SELECT 1 FROM information_schema.columns WHERE table_name = $1 AND column_name = 'embedding'", table_name) is not None
535
+
536
+ if vec_str and has_emb:
537
+ query = f"""
538
+ INSERT INTO {table_ref} (realm, id, payload, embedding)
539
+ VALUES ($1, $2, $3::jsonb, $4::vector)
540
+ ON CONFLICT (realm, id) DO UPDATE
541
+ SET payload = {table_ref}.payload || EXCLUDED.payload,
542
+ embedding = EXCLUDED.embedding,
543
+ updated_at = CURRENT_TIMESTAMP
544
+ RETURNING realm, id, fqid, payload, created_at, updated_at, embedding::text AS embedding_text
545
+ """
546
+ row = await conn.fetchrow(query, realm, v_id_int, payload_json, vec_str)
547
+ else:
548
+ query = f"""
549
+ INSERT INTO {table_ref} (realm, id, payload)
550
+ VALUES ($1, $2, $3::jsonb)
551
+ ON CONFLICT (realm, id) DO UPDATE
552
+ SET payload = {table_ref}.payload || EXCLUDED.payload,
553
+ updated_at = CURRENT_TIMESTAMP
554
+ RETURNING realm, id, fqid, payload, created_at, updated_at
555
+ """
501
556
  row = await conn.fetchrow(query, realm, v_id_int, payload_json)
557
+ try:
558
+ emb = None
559
+ if 'embedding_text' in row and row['embedding_text']:
560
+ emb = [float(x) for x in row['embedding_text'].strip('[]').split(',') if x.strip()]
502
561
  return Vertex(
503
562
  realm=row['realm'],
504
563
  id=str(row['id']),
@@ -507,6 +566,7 @@ class AsyncPostGraph:
507
566
  created_at=row['created_at'],
508
567
  updated_at=row['updated_at'],
509
568
  table_name=table_name,
569
+ embedding=emb,
510
570
  _client=self
511
571
  )
512
572
  except asyncpg.UndefinedTableError:
@@ -518,28 +578,112 @@ class AsyncPostGraph:
518
578
  """Fetch a vertex by realm and id."""
519
579
  self._validate_identifier(table_name)
520
580
  table_ref = self._get_table_ref(table_name, realm)
521
- v_str = str(vertex_id)
581
+ v_id_int = int(str(vertex_id).split('/')[-1]) if '/' in str(vertex_id) else int(vertex_id)
582
+
583
+ async def _op(conn):
584
+ query = f"""
585
+ SELECT t.realm, t.id, t.fqid, t.payload, t.created_at, t.updated_at,
586
+ to_jsonb(t)->>'embedding' AS embedding_text
587
+ FROM {table_ref} t
588
+ WHERE t.realm = $1 AND t.id = $2
589
+ """
590
+ try:
591
+ row = await conn.fetchrow(query, realm, v_id_int)
592
+ if not row:
593
+ return None
594
+
595
+ emb = None
596
+ if 'embedding_text' in row and row['embedding_text']:
597
+ emb = [float(x) for x in row['embedding_text'].strip('[]').split(',') if x.strip()]
598
+
599
+ return Vertex(
600
+ realm=row['realm'],
601
+ id=str(row['id']),
602
+ fqid=row['fqid'],
603
+ payload=row['payload'] if isinstance(row['payload'], dict) else json.loads(row['payload']),
604
+ created_at=row['created_at'],
605
+ updated_at=row['updated_at'],
606
+ table_name=table_name,
607
+ embedding=emb,
608
+ _client=self
609
+ )
610
+ except asyncpg.UndefinedTableError:
611
+ raise TableNotFoundError(f"Vertex table '{table_name}' does not exist.")
612
+
613
+ if isinstance(self.connection, asyncpg.Pool):
614
+ async with self.connection.acquire() as conn:
615
+ return await _op(conn)
616
+ else:
617
+ return await _op(self.connection)
618
+
619
+ async def vector_search(
620
+ self,
621
+ table_name: str,
622
+ realm: str,
623
+ query_vector: List[float],
624
+ top_k: int = 5,
625
+ distance_metric: str = "cosine"
626
+ ) -> List[Tuple[Vertex, float]]:
627
+ """Perform vector similarity search on vertex embeddings using pgvector.
628
+
629
+ Distance metrics:
630
+ - 'cosine': <=> (cosine distance)
631
+ - 'l2': <-> (Euclidean distance)
632
+ - 'inner_product': <#> (negative inner product)
633
+ """
634
+ self._validate_identifier(table_name)
635
+ table_ref = self._get_table_ref(table_name, realm)
636
+ vec_str = f"[{','.join(str(x) for x in query_vector)}]"
637
+
638
+ op = "<=>"
639
+ if distance_metric == "l2":
640
+ op = "<->"
641
+ elif distance_metric == "inner_product":
642
+ op = "<#>"
643
+
522
644
  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)
645
+ SELECT realm, id, fqid, payload, created_at, updated_at, embedding::text AS embedding_text,
646
+ (embedding {op} $2::vector) AS distance
647
+ FROM {table_ref}
648
+ WHERE realm = $1 AND embedding IS NOT NULL
649
+ ORDER BY embedding {op} $2::vector ASC
650
+ LIMIT $3
526
651
  """
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
540
- )
541
- except asyncpg.UndefinedTableError:
542
- raise TableNotFoundError(f"Vertex table '{table_name}' does not exist.")
652
+
653
+ async def _op(conn):
654
+ try:
655
+ rows = await conn.fetch(query, realm, vec_str, top_k)
656
+ except asyncpg.UndefinedTableError:
657
+ raise TableNotFoundError(f"Vertex table '{table_name}' does not exist.")
658
+ except asyncpg.UndefinedColumnError:
659
+ logger.warning(f"Table '{table_name}' does not have a vector embedding column.")
660
+ return []
661
+
662
+ results = []
663
+ for r in rows:
664
+ emb = None
665
+ if r['embedding_text']:
666
+ emb = [float(x) for x in r['embedding_text'].strip('[]').split(',') if x.strip()]
667
+ v = Vertex(
668
+ realm=r['realm'],
669
+ id=str(r['id']),
670
+ fqid=r['fqid'],
671
+ payload=r['payload'] if isinstance(r['payload'], dict) else json.loads(r['payload']),
672
+ created_at=r['created_at'],
673
+ updated_at=r['updated_at'],
674
+ table_name=table_name,
675
+ embedding=emb,
676
+ _client=self
677
+ )
678
+ dist = float(r['distance'])
679
+ results.append((v, dist))
680
+ return results
681
+
682
+ if isinstance(self.connection, asyncpg.Pool):
683
+ async with self.connection.acquire() as conn:
684
+ return await _op(conn)
685
+ else:
686
+ return await _op(self.connection)
543
687
 
544
688
  async def delete_vertex(self, table_name: str, realm: str, vertex_id: str, user_id: Optional[str] = None) -> bool:
545
689
  """Delete a vertex. Cascading foreign keys will automatically delete referencing edges."""
@@ -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,14 @@ 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
+ except Exception as e:
198
+ raise PostGraphError(f"Failed to initialize pgvector extension or embedding column for table '{table_name}': {e}")
199
+
187
200
  # 2. Create shadow audit table
188
201
  audit_query = f"""
189
202
  CREATE TABLE IF NOT EXISTS {audit_table_ref} (
@@ -381,12 +394,14 @@ class SQLAlchemyPostGraph:
381
394
  realm: str,
382
395
  vertex_id: Optional[Union[str, int]] = None,
383
396
  payload: Optional[Dict[str, Any]] = None,
397
+ embedding: Optional[List[float]] = None,
384
398
  user_id: Optional[str] = None
385
399
  ) -> Vertex:
386
400
  """Add a new vertex. Raises TableExistsError if it already exists."""
387
401
  self._validate_identifier(table_name)
388
402
  payload_json = json.dumps(payload or {})
389
403
  table_ref = self._get_table_ref(table_name, realm)
404
+ vec_str = f"[{','.join(str(x) for x in embedding)}]" if embedding else None
390
405
 
391
406
  async def _op(conn):
392
407
  nonlocal vertex_id
@@ -399,17 +414,30 @@ class SQLAlchemyPostGraph:
399
414
 
400
415
  v_id_str = str(v_id_int)
401
416
 
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
- """
417
+ if vec_str:
418
+ query = f"""
419
+ INSERT INTO {table_ref} (realm, id, payload, embedding)
420
+ VALUES (:realm, :id, CAST(:payload AS JSONB), CAST(:vec AS vector))
421
+ RETURNING realm, id, fqid, payload, created_at, updated_at, CAST(embedding AS TEXT) AS embedding_text
422
+ """
423
+ kwargs = {"realm": realm, "id": v_id_int, "payload": payload_json, "vec": vec_str}
424
+ else:
425
+ query = f"""
426
+ INSERT INTO {table_ref} (realm, id, payload)
427
+ VALUES (:realm, :id, CAST(:payload AS JSONB))
428
+ RETURNING realm, id, fqid, payload, created_at, updated_at
429
+ """
430
+ kwargs = {"realm": realm, "id": v_id_int, "payload": payload_json}
431
+
407
432
  try:
408
- row = await self._fetchrow(conn, query, realm=realm, id=v_id_int, payload=payload_json)
433
+ row = await self._fetchrow(conn, query, **kwargs)
409
434
  if vertex_id is not None:
410
435
  await conn.execute(
411
436
  text(f"SELECT setval(pg_get_serial_sequence('{table_ref_pg}', 'id'), (SELECT COALESCE(MAX(id), 1) FROM {table_ref}))")
412
437
  )
438
+ emb = None
439
+ if 'embedding_text' in row and row['embedding_text']:
440
+ emb = [float(x) for x in row['embedding_text'].strip('[]').split(',') if x.strip()]
413
441
  return Vertex(
414
442
  realm=row['realm'],
415
443
  id=str(row['id']),
@@ -418,6 +446,7 @@ class SQLAlchemyPostGraph:
418
446
  created_at=row['created_at'],
419
447
  updated_at=row['updated_at'],
420
448
  table_name=table_name,
449
+ embedding=emb,
421
450
  _client=self
422
451
  )
423
452
  except IntegrityError as e:
@@ -439,12 +468,14 @@ class SQLAlchemyPostGraph:
439
468
  realm: str,
440
469
  vertex_id: Optional[Union[str, int]] = None,
441
470
  payload: Optional[Dict[str, Any]] = None,
471
+ embedding: Optional[List[float]] = None,
442
472
  user_id: Optional[str] = None
443
473
  ) -> Vertex:
444
474
  """Upsert a vertex (merges payload JSONB on conflict)."""
445
475
  self._validate_identifier(table_name)
446
476
  payload_json = json.dumps(payload or {})
447
477
  table_ref = self._get_table_ref(table_name, realm)
478
+ vec_str = f"[{','.join(str(x) for x in embedding)}]" if embedding else None
448
479
 
449
480
  async def _op(conn):
450
481
  nonlocal vertex_id
@@ -457,16 +488,33 @@ class SQLAlchemyPostGraph:
457
488
 
458
489
  v_id_str = str(v_id_int)
459
490
 
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
- """
491
+ if vec_str:
492
+ query = f"""
493
+ INSERT INTO {table_ref} (realm, id, payload, embedding)
494
+ VALUES (:realm, :id, CAST(:payload AS JSONB), CAST(:vec AS vector))
495
+ ON CONFLICT (realm, id) DO UPDATE
496
+ SET payload = {table_ref}.payload || EXCLUDED.payload,
497
+ embedding = EXCLUDED.embedding,
498
+ updated_at = CURRENT_TIMESTAMP
499
+ RETURNING realm, id, fqid, payload, created_at, updated_at, CAST(embedding AS TEXT) AS embedding_text
500
+ """
501
+ kwargs = {"realm": realm, "id": v_id_int, "payload": payload_json, "vec": vec_str}
502
+ else:
503
+ query = f"""
504
+ INSERT INTO {table_ref} (realm, id, payload)
505
+ VALUES (:realm, :id, CAST(:payload AS JSONB))
506
+ ON CONFLICT (realm, id) DO UPDATE
507
+ SET payload = {table_ref}.payload || EXCLUDED.payload,
508
+ updated_at = CURRENT_TIMESTAMP
509
+ RETURNING realm, id, fqid, payload, created_at, updated_at
510
+ """
511
+ kwargs = {"realm": realm, "id": v_id_int, "payload": payload_json}
512
+
468
513
  try:
469
- row = await self._fetchrow(conn, query, realm=realm, id=v_id_int, payload=payload_json)
514
+ row = await self._fetchrow(conn, query, **kwargs)
515
+ emb = None
516
+ if 'embedding_text' in row and row['embedding_text']:
517
+ emb = [float(x) for x in row['embedding_text'].strip('[]').split(',') if x.strip()]
470
518
  return Vertex(
471
519
  realm=row['realm'],
472
520
  id=str(row['id']),
@@ -475,6 +523,7 @@ class SQLAlchemyPostGraph:
475
523
  created_at=row['created_at'],
476
524
  updated_at=row['updated_at'],
477
525
  table_name=table_name,
526
+ embedding=emb,
478
527
  _client=self
479
528
  )
480
529
  except ProgrammingError as e:
@@ -488,26 +537,33 @@ class SQLAlchemyPostGraph:
488
537
  """Fetch a vertex by realm and id."""
489
538
  self._validate_identifier(table_name)
490
539
  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
-
540
+ v_id_int = int(str(vertex_id).split('/')[-1]) if '/' in str(vertex_id) else int(vertex_id)
541
+
498
542
  async def _op(conn):
543
+ query = f"""
544
+ SELECT t.realm, t.id, t.fqid, t.payload, t.created_at, t.updated_at,
545
+ to_jsonb(t)->>'embedding' AS embedding_text
546
+ FROM {table_ref} t
547
+ WHERE t.realm = :realm AND t.id = :id
548
+ """
499
549
  try:
500
- row = await self._fetchrow(conn, query, realm=realm, id=v_str)
550
+ row = await self._fetchrow(conn, query, realm=realm, id=v_id_int)
501
551
  if not row:
502
552
  return None
553
+
554
+ emb = None
555
+ if 'embedding_text' in row and row['embedding_text']:
556
+ emb = [float(x) for x in row['embedding_text'].strip('[]').split(',') if x.strip()]
557
+
503
558
  return Vertex(
504
559
  realm=row['realm'],
505
560
  id=str(row['id']),
506
- fqid=row['fqid'] or f"{row['realm']}/{table_name}/{row['id']}",
561
+ fqid=row['fqid'],
507
562
  payload=row['payload'] if isinstance(row['payload'], dict) else json.loads(row['payload']),
508
563
  created_at=row['created_at'],
509
564
  updated_at=row['updated_at'],
510
565
  table_name=table_name,
566
+ embedding=emb,
511
567
  _client=self
512
568
  )
513
569
  except ProgrammingError as e:
@@ -521,6 +577,69 @@ class SQLAlchemyPostGraph:
521
577
  async with self.engine_or_connection.connect() as conn:
522
578
  return await _op(conn)
523
579
 
580
+ async def vector_search(
581
+ self,
582
+ table_name: str,
583
+ realm: str,
584
+ query_vector: List[float],
585
+ top_k: int = 5,
586
+ distance_metric: str = "cosine"
587
+ ) -> List[Tuple[Vertex, float]]:
588
+ """Perform vector similarity search on vertex embeddings using pgvector."""
589
+ self._validate_identifier(table_name)
590
+ table_ref = self._get_table_ref(table_name, realm)
591
+ vec_str = f"[{','.join(str(x) for x in query_vector)}]"
592
+
593
+ op = "<=>"
594
+ if distance_metric == "l2":
595
+ op = "<->"
596
+ elif distance_metric == "inner_product":
597
+ op = "<#>"
598
+
599
+ query = f"""
600
+ SELECT realm, id, fqid, payload, created_at, updated_at, CAST(embedding AS TEXT) AS embedding_text,
601
+ (embedding {op} CAST(:vec AS vector)) AS distance
602
+ FROM {table_ref}
603
+ WHERE realm = :realm AND embedding IS NOT NULL
604
+ ORDER BY embedding {op} CAST(:vec AS vector) ASC
605
+ LIMIT :top_k
606
+ """
607
+
608
+ async def _op(conn):
609
+ try:
610
+ rows = await self._fetch(conn, query, realm=realm, vec=vec_str, top_k=top_k)
611
+ except ProgrammingError as e:
612
+ if "does not exist" in str(e).lower():
613
+ raise TableNotFoundError(f"Vertex table '{table_name}' does not exist.")
614
+ logger.warning(f"Vector search failed: {e}")
615
+ return []
616
+
617
+ results = []
618
+ for r in rows:
619
+ emb = None
620
+ if r['embedding_text']:
621
+ emb = [float(x) for x in r['embedding_text'].strip('[]').split(',') if x.strip()]
622
+ v = Vertex(
623
+ realm=r['realm'],
624
+ id=str(r['id']),
625
+ fqid=r['fqid'],
626
+ payload=r['payload'] if isinstance(r['payload'], dict) else json.loads(r['payload']),
627
+ created_at=r['created_at'],
628
+ updated_at=r['updated_at'],
629
+ table_name=table_name,
630
+ embedding=emb,
631
+ _client=self
632
+ )
633
+ dist = float(r['distance'])
634
+ results.append((v, dist))
635
+ return results
636
+
637
+ if isinstance(self.engine_or_connection, AsyncConnection):
638
+ return await _op(self.engine_or_connection)
639
+ else:
640
+ async with self.engine_or_connection.connect() as conn:
641
+ return await _op(conn)
642
+
524
643
  async def delete_vertex(self, table_name: str, realm: str, vertex_id: str, user_id: Optional[str] = None) -> bool:
525
644
  """Delete a vertex. Cascading foreign keys will automatically delete referencing edges."""
526
645
  self._validate_identifier(table_name)
@@ -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']:
@@ -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.3"
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