better db connectin hadeling for long lasting deplyment

parent a4501a9c
{
"schemaVersion": 2,
"dockerfilePath": "./Dockerfile",
"containerHttpPort": "8000",
"env": {
"POSTGRES_HOST": "srv-captain--postgres",
"MINIO_HOST": "srv-captain--minio"
}
}
......@@ -17,7 +17,7 @@ class StudentNationality(str, Enum):
class Models(str, Enum):
chat = "gpt-5"
chat = "gpt-4o-mini"
tts = "gpt-4o-mini-tts"
embedding = "text-embedding-3-small"
transcription = "gpt-4o-transcribe"
This diff is collapsed.
import os
import psycopg2
from psycopg2.extras import RealDictCursor
from psycopg2.pool import ThreadedConnectionPool
from typing import List, Optional
# Import the pgvector adapter
from pgvector.psycopg2 import register_vector
class PGVectorService:
"""Service for managing embeddings with PostgreSQL pgvector"""
"""Service for managing embeddings with PostgreSQL pgvector using connection pooling"""
def __init__(self):
self.conn = psycopg2.connect(
self.pool = ThreadedConnectionPool(
minconn=1,
maxconn=20,
host=os.getenv("POSTGRES_HOST", "postgres"),
user=os.getenv("POSTGRES_USER"),
password=os.getenv("POSTGRES_PASSWORD"),
dbname=os.getenv("POSTGRES_DB"),
)
# Register the vector type with the connection
register_vector(self.conn)
# Test connection and register vector type to ensure the pool works
conn = self.pool.getconn()
try:
register_vector(conn)
finally:
self.pool.putconn(conn)
def _get_conn_with_vector(self):
"""Get a connection from the pool and register vector type"""
conn = self.pool.getconn()
register_vector(conn)
return conn
def insert_embedding(self, id: int, embedding: list):
"""Insert or update an embedding"""
with self.conn.cursor() as cur:
cur.execute(
"""
INSERT INTO embeddings_table (id, embedding)
VALUES (%s, %s)
ON CONFLICT (id) DO UPDATE SET embedding = EXCLUDED.embedding;
""",
(id, embedding),
)
self.conn.commit()
conn = self._get_conn_with_vector()
try:
with conn.cursor() as cur:
cur.execute(
"""
INSERT INTO embeddings_table (id, embedding)
VALUES (%s, %s)
ON CONFLICT (id) DO UPDATE SET embedding = EXCLUDED.embedding;
""",
(id, embedding),
)
conn.commit()
finally:
self.pool.putconn(conn)
def search_nearest(self, query_embedding: list, limit: int = 3):
"""Search nearest embeddings using cosine distance (<-> operator)"""
with self.conn.cursor(cursor_factory=RealDictCursor) as cur:
cur.execute(
"""
SELECT id, embedding, embedding <-> %s AS distance
FROM embeddings_table
ORDER BY embedding <-> %s
LIMIT %s;
""",
(query_embedding, query_embedding, limit),
)
return cur.fetchall()
conn = self._get_conn_with_vector()
try:
with conn.cursor(cursor_factory=RealDictCursor) as cur:
cur.execute(
"""
SELECT id, embedding, embedding <-> %s AS distance
FROM embeddings_table
ORDER BY embedding <-> %s
LIMIT %s;
""",
(query_embedding, query_embedding, limit),
)
return cur.fetchall()
finally:
self.pool.putconn(conn)
def search_filtered_nearest(
self,
......@@ -55,21 +76,25 @@ class PGVectorService:
limit: int = 3
):
"""Search nearest embeddings with filtering by grade, subject, and language"""
with self.conn.cursor(cursor_factory=RealDictCursor) as cur:
cur.execute(
"""
SELECT id, grade, subject, unit, concept, lesson, chunk_text,
is_arabic, embedding <-> %s::vector AS distance
FROM educational_chunks
WHERE grade = %s
AND subject ILIKE %s
AND is_arabic = %s
ORDER BY embedding <-> %s::vector
LIMIT %s;
""",
(query_embedding, grade, f"%{subject}%", is_arabic, query_embedding, limit),
)
return cur.fetchall()
conn = self._get_conn_with_vector()
try:
with conn.cursor(cursor_factory=RealDictCursor) as cur:
cur.execute(
"""
SELECT id, grade, subject, unit, concept, lesson, chunk_text,
is_arabic, embedding <-> %s::vector AS distance
FROM educational_chunks
WHERE grade = %s
AND subject ILIKE %s
AND is_arabic = %s
ORDER BY embedding <-> %s::vector
LIMIT %s;
""",
(query_embedding, grade, f"%{subject}%", is_arabic, query_embedding, limit),
)
return cur.fetchall()
finally:
self.pool.putconn(conn)
def search_flexible_filtered_nearest(
self,
......@@ -103,34 +128,42 @@ class PGVectorService:
params.append(query_embedding)
params.append(limit)
with self.conn.cursor(cursor_factory=RealDictCursor) as cur:
cur.execute(
f"""
SELECT id, grade, subject, unit, concept, lesson, chunk_text,
is_arabic, embedding <-> %s::vector AS distance
FROM educational_chunks
{where_clause}
ORDER BY embedding <-> %s::vector
LIMIT %s;
""",
params
)
return cur.fetchall()
conn = self._get_conn_with_vector()
try:
with conn.cursor(cursor_factory=RealDictCursor) as cur:
cur.execute(
f"""
SELECT id, grade, subject, unit, concept, lesson, chunk_text,
is_arabic, embedding <-> %s::vector AS distance
FROM educational_chunks
{where_clause}
ORDER BY embedding <-> %s::vector
LIMIT %s;
""",
params
)
return cur.fetchall()
finally:
self.pool.putconn(conn)
def get_subjects_by_grade_and_language(self, grade: int, is_arabic: bool) -> List[str]:
"""Get available subjects for a specific grade and language"""
with self.conn.cursor(cursor_factory=RealDictCursor) as cur:
cur.execute(
"""
SELECT DISTINCT subject
FROM educational_chunks
WHERE grade = %s AND is_arabic = %s
ORDER BY subject;
""",
(grade, is_arabic)
)
return [row['subject'] for row in cur.fetchall()]
conn = self._get_conn_with_vector()
try:
with conn.cursor(cursor_factory=RealDictCursor) as cur:
cur.execute(
"""
SELECT DISTINCT subject
FROM educational_chunks
WHERE grade = %s AND is_arabic = %s
ORDER BY subject;
""",
(grade, is_arabic)
)
return [row['subject'] for row in cur.fetchall()]
finally:
self.pool.putconn(conn)
def close(self):
if self.conn:
self.conn.close()
\ No newline at end of file
def close_pool(self):
if self.pool:
self.pool.closeall()
\ No newline at end of file
#!/bin/bash
set -e
host="postgres"
# Use the environment variables set in CapRover for the host and user
host="${POSTGRES_HOST}"
user="${POSTGRES_USER}"
db="${POSTGRES_DB}"
password="${POSTGRES_PASSWORD}"
echo "Waiting for PostgreSQL database to be ready..."
echo "Waiting for PostgreSQL database at $host to be ready..."
# Wait for the database server to be available
until PGPASSWORD="$password" pg_isready -h "$host" -U "$user"; do
# Wait for the database server to be available. This is sufficient.
until PGPASSWORD="${POSTGRES_PASSWORD}" pg_isready -h "$host" -U "$user"; do
echo "PostgreSQL server is unavailable - sleeping"
sleep 1
done
# Wait for the specific database to be available
until PGPASSWORD="$password" psql -h "$host" -U "$user" -d "$db" -c '\q'; do
echo "Database '$db' is not yet available - sleeping"
sleep 1
done
echo "PostgreSQL database is up and ready - executing command"
exec "$@"
\ No newline at end of file
exec "$@"
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment