32 lines
1.1 KiB
Python
32 lines
1.1 KiB
Python
from __future__ import annotations
|
|
import json, logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class KnowledgeStore:
|
|
def __init__(self, pool):
|
|
self._pool = pool
|
|
|
|
async def upsert(self, entity_type, entity_key, facts):
|
|
async with self._pool.acquire(timeout=10) as conn:
|
|
await conn.execute(
|
|
"""INSERT INTO ab_knowledge_store (entity_type, entity_key, facts)
|
|
VALUES ($1, $2, $3)
|
|
ON CONFLICT (entity_type, entity_key)
|
|
DO UPDATE SET facts = $3, updated_at = NOW()""",
|
|
entity_type, entity_key, json.dumps(facts))
|
|
|
|
async def get(self, entity_type, entity_key):
|
|
async with self._pool.acquire(timeout=10) as conn:
|
|
row = await conn.fetchrow(
|
|
'SELECT facts FROM ab_knowledge_store WHERE entity_type=$1 AND entity_key=$2',
|
|
entity_type, entity_key)
|
|
if not row:
|
|
return {}
|
|
f = row['facts']
|
|
return json.loads(f) if isinstance(f, str) else f
|
|
|
|
async def get_client_profile(self, partner_id):
|
|
return await self.get('client', f'partner_{partner_id}')
|