Skip to content

fix: Use singleton PostgresDBClient (Sqlalchemy engine) #321

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 21 commits into from
Jun 23, 2025
Merged

fix: Use singleton PostgresDBClient (Sqlalchemy engine) #321

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
21 commits
Select commit Hold shift + click to select a range
baf3708
use single db.PostgresDBClient instance
yoomlam Jun 17, 2025
0a335fe
allow promptfoo to use more than 1 thread
yoomlam Jun 17, 2025
e85f9fb
add Gemini models
yoomlam Jun 17, 2025
b2ce4aa
temp: use gemini 2.5 pro for promptfoo eval
yoomlam Jun 17, 2025
e709932
de-lint
yoomlam Jun 17, 2025
ad5c1f0
add GEMINI_API_KEY to promptfoo action
yoomlam Jun 17, 2025
e340876
add app_config fixture to failing tests
yoomlam Jun 17, 2025
5ff8620
TEMP: ignore some tests
yoomlam Jun 17, 2025
bb12097
fix more tests
yoomlam Jun 17, 2025
cf2fd49
WIP
yoomlam Jun 17, 2025
5b8411e
add GEMINI_API_KEY to terraform
yoomlam Jun 18, 2025
7906d97
revert debugging code
yoomlam Jun 18, 2025
0e3c0a8
add Gemini LLM
yoomlam Jun 18, 2025
9b001fd
remove temporary code
yoomlam Jun 18, 2025
dcd8e45
add to terraform
yoomlam Jun 18, 2025
9697231
compare by chunk id
yoomlam Jun 18, 2025
c22c07d
fix more tests
yoomlam Jun 18, 2025
8e89b00
remove print()
yoomlam Jun 18, 2025
c17be59
Merge branch 'yl/add-gemini-llm' into yl/pgclient-singleton
yoomlam Jun 18, 2025
4c44d8d
Merge branch 'main' into yl/pgclient-singleton
yoomlam Jun 23, 2025
a960bb4
remove old code
yoomlam Jun 23, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions .github/workflows/promptfoo-googlesheet-evaluation.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,7 @@ jobs:
env:
GOOGLE_APPLICATION_CREDENTIALS: /tmp/gcp-creds.json
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }}
PROMPTFOO_FAILED_TEST_EXIT_CODE: 0
run: |
if [ ! -f "$GOOGLE_APPLICATION_CREDENTIALS" ]; then
Expand All @@ -118,9 +119,9 @@ jobs:
EVAL_OUTPUT_FILE="/tmp/promptfoo-output.txt"

if [ -n "$PROMPTFOO_API_KEY" ]; then
promptfoo eval --max-concurrency 1 --config "/tmp/promptfooconfig.processed.yaml" --share --output "${OUTPUT_JSON_FILE}" --no-cache | tee "${EVAL_OUTPUT_FILE}"
promptfoo eval --config "/tmp/promptfooconfig.processed.yaml" --share --output "${OUTPUT_JSON_FILE}" --no-cache | tee "${EVAL_OUTPUT_FILE}"
else
promptfoo eval --max-concurrency 1 --config "/tmp/promptfooconfig.processed.yaml" --output "${OUTPUT_JSON_FILE}" --no-cache | tee "${EVAL_OUTPUT_FILE}"
promptfoo eval --config "/tmp/promptfooconfig.processed.yaml" --output "${OUTPUT_JSON_FILE}" --no-cache | tee "${EVAL_OUTPUT_FILE}"
Comment on lines -121 to +123
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reverts #317, so it should default back to 4 threads

fi

if [ -f "${OUTPUT_JSON_FILE}" ]; then
Expand Down
6 changes: 5 additions & 1 deletion app/src/app_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,8 +42,12 @@ class AppConfig(PydanticBaseEnvConfig):
# If set, used instead of LITERAL_API_KEY for API
literal_api_key_for_api: str | None = None

@cached_property
def db_client(self) -> db.PostgresDBClient:
return db.PostgresDBClient()
Comment on lines +45 to +47
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This basically creates a singleton PostgresDBClient instance, which holds the Sqlalchemy engine.

“the engine is thread safe yes. individual Connection objects are not. we try to describe this at Working with Engines and Connections — SQLAlchemy 2.0 Documentation”


def db_session(self) -> db.Session:
return db.PostgresDBClient().get_session()
return self.db_client.get_session()
Comment on lines -46 to +50
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A new session uses an available connection (from the connection pool), so threads (e.g., those created from API calls) will not share a connection.

Our pool size is 20 – what happens when max sqlalchemy pool size is reached?

Google’s AI states:

Requests are queued: Any new requests for a database connection are placed in a queue, waiting for a connection to become available.

Timeout begins: SQLAlchemy starts a timeout period (defaulting to 30 seconds, but configurable) to see if a connection is released back into the pool.

Connection timeout error: If a connection doesn't become available within the timeout period, an exception (e.g., TimeoutError) is thrown, indicating a connection timeout.

SQLAlchemy connection pooling, what are checked out connections? : “If all the connections are simultaneously checked out then you can expect an error (there will be a timeout period during which SQLAlchemy waits to see if a connection gets freed up; this is also configurable).”


@cached_property
def embedding_model(self) -> EmbeddingModel:
Expand Down
1 change: 1 addition & 0 deletions app/src/chat_engine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -335,6 +335,7 @@ class ImagineLaEngine(BaseEngine):
]

engine_id: str = "imagine-la"
llm: str = "gemini/gemini-2.5-pro-preview-06-05"
name: str = "SBN Chat Engine"
datasets = [
"CA EDD",
Expand Down
5 changes: 5 additions & 0 deletions app/src/generate.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,11 @@ def get_models() -> dict[str, str]:
models |= {"OpenAI GPT-4o": "gpt-4o"}
if "ANTHROPIC_API_KEY" in os.environ:
models |= {"Anthropic Claude 3.5 Sonnet": "claude-3-5-sonnet-20240620"}
if "GEMINI_API_KEY" in os.environ:
models |= {
"Google Gemini 1.5 Pro": "gemini/gemini-1.5-pro",
"Google Gemini 2.5 Pro": "gemini/gemini-2.5-pro-preview-06-05",
}
if _has_aws_access():
# If you get "You don't have access to the model with the specified model ID." error,
# remember to request access to Bedrock models ...aws.amazon.com/bedrock/home?region=us-east-1#/modelaccess
Expand Down
Loading