36 lines
1.4 KiB
Python
36 lines
1.4 KiB
Python
from __future__ import annotations
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class LLMConfigStore:
|
|
def __init__(self, pg_pool):
|
|
self._pool = pg_pool
|
|
|
|
async def get_backend(self, caller):
|
|
async with self._pool.acquire(timeout=10) as conn:
|
|
row = await conn.fetchrow(
|
|
'SELECT backend FROM ab_llm_config WHERE caller = $1', caller)
|
|
return row['backend'] if row else None
|
|
|
|
async def set_backend(self, caller, backend, set_by, note=None):
|
|
async with self._pool.acquire(timeout=10) as conn:
|
|
await conn.execute(
|
|
"""INSERT INTO ab_llm_config (caller, backend, set_by, note)
|
|
VALUES ($1, $2, $3, $4)
|
|
ON CONFLICT (caller) DO UPDATE
|
|
SET backend=$2, set_by=$3, set_at=NOW(), note=$4""",
|
|
caller, backend, set_by, note)
|
|
logger.info('LLMConfigStore.set_backend caller=%s backend=%s set_by=%s', caller, backend, set_by)
|
|
|
|
async def get_all(self):
|
|
async with self._pool.acquire(timeout=10) as conn:
|
|
rows = await conn.fetch('SELECT * FROM ab_llm_config ORDER BY caller')
|
|
return [dict(r) for r in rows]
|
|
|
|
async def reset(self, caller):
|
|
async with self._pool.acquire(timeout=10) as conn:
|
|
await conn.execute('DELETE FROM ab_llm_config WHERE caller = $1', caller)
|
|
logger.info('LLMConfigStore.reset caller=%s', caller)
|