File size: 2,089 Bytes
dc58d54 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
import asyncpg
class NetworkDB:
def __init__(self, database_url):
self.pool = None
self.database_url = database_url
async def get_pool(self):
if self.pool:
return self.pool
self.pool = await asyncpg.create_pool(
self.database_url, min_size=1, max_size=10
)
return self.pool
async def post_text(self, content: str, embeddings: list[float]) -> bool:
# pool = await self.get_pool()
# async with pool.acquire() as conn:
try:
conn = await asyncpg.connect(self.database_url)
id = await conn.fetchval(
"INSERT INTO text_posts (content, embedding) VALUES ($1, $2) RETURNING id",
content,
f"{embeddings}",
)
await conn.close()
return True if id is not None else False
except Exception as e:
return False
async def get_text_post_random(self) -> str:
try:
conn = await asyncpg.connect(self.database_url)
post = await conn.fetchval(
"SELECT content from text_posts ORDER BY random() LIMIT 1"
)
await conn.close()
return post if post is not None else "[Internal Message: No post found!]"
except Exception as e:
return "[Internal Message: Server Error]"
async def get_text_post_similar(self, query_embedding: list[float]) -> str:
try:
conn = await asyncpg.connect(self.database_url)
post = await conn.fetchval(
"SELECT content FROM text_posts ORDER BY embedding <-> $1 LIMIT 1",
f"{query_embedding}",
)
await conn.close()
return (
post
if post is not None
else "[Internal Message: No similar post found!]"
)
except Exception as e:
return "[Internal Message: Server Error]"
async def disconnect(self) -> None:
if self.pool:
self.pool.close()
|