Hash postgres keys to support long documents (#62)

* [WIP] Hash the cache key in postgres

* [WIP] Format
pull/82/head
Sabri Eyuboglu 1 year ago committed by GitHub
parent e00d285e21
commit bed6773f75

@ -1,4 +1,5 @@
"""Postgres cache."""
import hashlib
import logging
from typing import Any, Dict, Union
@ -83,9 +84,10 @@ class PostgresCache(Cache):
"""Close the client."""
self.session.close()
def _normalize_table_key(self, key: str, table: str) -> str:
"""Cast key for prompt key."""
return f"{table}:{key}"
@staticmethod
def _hash_key(key: str, table: str) -> str:
"""Compute MD5 hash of the key."""
return hashlib.md5(f"{key}:{table}".encode("utf-8")).hexdigest()
def get_key(self, key: str, table: str = "default") -> Union[str, None]:
"""
@ -97,7 +99,11 @@ class PostgresCache(Cache):
key: key for cache.
table: table to get key in.
"""
request = self.session.query(Request).filter_by(key=key).first()
request = (
self.session.query(Request)
.filter_by(key=self._hash_key(key, table))
.first()
)
out = request.response if request else None
return out # type: ignore
@ -112,6 +118,7 @@ class PostgresCache(Cache):
value: new value for key.
table: table to set key in.
"""
key = self._hash_key(key, table)
request = self.session.query(Request).filter_by(key=key).first()
if request:
request.response = value # type: ignore

Loading…
Cancel
Save