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

Open
wants to merge 5 commits into
base: main
Choose a base branch
from

Conversation

artificial-aidan
Copy link

@artificial-aidan artificial-aidan 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

@artificial-aidan
Copy link
Author

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:session labels Jul 17, 2025
@artificial-aidan
Copy link
Author

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.

@artificial-aidan artificial-aidan 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!

@artificial-aidan
Copy link
Author

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

* 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 <[email protected]>
Signed-off-by: Aidan Jensen <[email protected]>
* 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 <[email protected]>
Signed-off-by: Aidan Jensen <[email protected]>
@artificial-aidan artificial-aidan force-pushed the aidan/postgres-sessions branch from 9f0d7a1 to 78aa642 Compare July 22, 2025 03:30
@artificial-aidan
Copy link
Author

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

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.

3 participants