import psycopg
from langchain_postgres import PostgresChatMessageHistory
from psycopg_pool import ConnectionPool
from contextlib import contextmanager

# conn = psycopg.connect("postgresql://postgres:123456@10.10.10.9/ktor2")
database_name = "ai_chat_history"
pool = ConnectionPool("postgresql://postgres:123456@10.10.10.9/ktor2")

@contextmanager
def getConn():
    with pool.connection() as temp:
        temp.autocommit = True   # 如果你想所有连接默认 autocommit
        yield temp               # 把 conn 暴露给外部使用


def init():
    with getConn() as connection:
        PostgresChatMessageHistory.create_tables(connection, database_name)

init()