Skip to content

Initial pass at PostgreSQL session management #1163

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.

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

Closed
wants to merge 5 commits into from

Conversation

aidandj
Copy link

@aidandj aidandj commented Jul 17, 2025

Summary

  • Add postgres_session.py
    • Implements session management using PostgreSQL
    • Accepts a pool as a default argument
    • Additional class method for creation from a connection string
  • Add optional-dependency for psycopg

Test plan

Tests were written mocking the postgres calls. This was tested locally against a docker postgres database, with a variety of input messages.

Issue number

Related to #745

Checks

  • I've added new tests (if relevant)
  • I've added/updated the relevant documentation
  • I've run make lint and make format
  • I've made sure tests pass

@aidandj
Copy link
Author

aidandj commented Jul 17, 2025

Wanted to get the code up here in case there was some discussion on how to do things. Will look into the documentation tomorrow.

@seratch seratch requested a review from rm-openai July 17, 2025 04:45
@seratch seratch added enhancement New feature or request feature:sessions labels Jul 17, 2025
@aidandj
Copy link
Author

aidandj commented Jul 17, 2025

Got a little stuck on the old_version_test needed to have the pyscopg extras installed, not sure how litellm is working. I updated the make target to include all extras when testing python 3.9. Let me know if this isn't the path you'd prefer.

I took a stab at documentation. Open to feedback on structuring. Overall I wanted the postgres session to be in extensions because psycopg can be somewhat heavy on certain platforms. But I wasn't sure if the docs should go with extensions or the other memory docs.

@aidandj aidandj force-pushed the aidan/postgres-sessions branch 4 times, most recently from ff857ed to be28fbc Compare July 20, 2025 20:15
Copy link
Collaborator

@rm-openai rm-openai left a comment

Choose a reason for hiding this comment

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

@artificial-aidan lmk if this is ready for review and i can take a pass!

@aidandj
Copy link
Author

aidandj commented Jul 21, 2025

Should be ready now. Have been running it in our testing for about a week now.

@aidandj aidandj force-pushed the aidan/postgres-sessions branch from be28fbc to 4498798 Compare July 21, 2025 19:35
@aidandj aidandj requested a review from rm-openai July 21, 2025 19:35
@aidandj aidandj force-pushed the aidan/postgres-sessions branch from 9f0d7a1 to 78aa642 Compare July 22, 2025 03:30
@aidandj
Copy link
Author

aidandj commented Jul 22, 2025

Tests fixed. Missed the test changes I had made locally.

@aidandj aidandj force-pushed the aidan/postgres-sessions branch 4 times, most recently from ef1a02d to 7d392a1 Compare July 28, 2025 17:47
def __init__(
self,
session_id: str,
pool: AsyncConnectionPool,

Choose a reason for hiding this comment

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

Why you use a hole connection pool just for one session management? One connection will be enough.

Copy link
Author

Choose a reason for hiding this comment

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

In my use case this is running in a service with many concurrent sessions. A shared connection pool is used between them. That's why you can initialize it from a connection string, or using an existing pool

Choose a reason for hiding this comment

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

Okay, that is your way, it is not bad.

aidandj added 5 commits August 5, 2025 11:06
* Add postgres_session.py
    * Implements session management using PostgreSQL
    * Accepts a pool as a default argument
    * Additional class method for creation from a connection string
* Add optional-dependency for psycopg

Signed-off-by: Aidan Jensen <aidan@artificial.com>
Signed-off-by: Aidan Jensen <aidan@artificial.com>
Signed-off-by: Aidan Jensen <aidan@artificial.com>
* Fix warning from opening a pool in the constructor: https://www.psycopg.org/psycopg3/docs/news_pool.html#psycopg-pool-3-2-2
* Add integration tests that can be run locally
* Fix get_items with limit

Signed-off-by: Aidan Jensen <aidan@artificial.com>
Signed-off-by: Aidan Jensen <aidan@artificial.com>
@aidandj aidandj force-pushed the aidan/postgres-sessions branch from 4eb7af4 to 1bfcec4 Compare August 5, 2025 18:06
@@ -38,6 +38,7 @@ voice = ["numpy>=2.2.0, <3; python_version>='3.10'", "websockets>=15.0, <16"]
viz = ["graphviz>=0.17"]
litellm = ["litellm>=1.67.4.post1, <2"]
realtime = ["websockets>=15.0, <16"]
psycopg = ["psycopg[pool]>=3.2.9,<4"]
Copy link
Contributor

Choose a reason for hiding this comment

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

I'd recommend psycopg-binary so users don't need a compiler installed to use this https://pypi.org/project/psycopg-binary/

Copy link
Author

@aidandj aidandj Aug 10, 2025

Choose a reason for hiding this comment

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

Does this cover all platform/arch combinations? By not specifying binary, any user of this library should be able to choose binary if they'd like. If a binary doesn't exist for a platform then won't over specifying the dependency cause issues? (I have not experimented with the binary library)

def __init__(
self,
session_id: str,
pool: AsyncConnectionPool,
Copy link
Contributor

Choose a reason for hiding this comment

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

It would be nice to have this as an optional parameter with the default set to None. It can then be created in the __init__ if the user doesn't pass it explicitly.

Copy link
Author

Choose a reason for hiding this comment

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

How could it be created? That was the intention of the from_connection_string classmethod

@Mukish45
Copy link

@seratch @rm-openai - May I know the expected timeline for this release?

@seratch
Copy link
Member

seratch commented Aug 21, 2025

We've accepted this SQLAlchemy based implementation, which supports not only PostgreSQL but any other RDBs: #1357

Thus, to avoid duplicated implementations in the built-in modules, let us close this pull request now. Anyone who prefers a simpler one can use this PostgreSQL implementation by copying the code in their projects.

We truly appreciate your time and efforts here. Thank you so much again!

@seratch seratch closed this Aug 21, 2025
@Mukish45
Copy link

Makes sense. Thanks @seratch

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

6 participants