@ -56,18 +56,19 @@ class FullLLMCache(Base): # type: ignore
class SQLAlchemyCache ( BaseCache ) :
class SQLAlchemyCache ( BaseCache ) :
""" Cache that uses SQAlchemy as a backend. """
""" Cache that uses SQAlchemy as a backend. """
def __init__ ( self , engine : Engine ):
def __init__ ( self , engine : Engine , cache_schema : Any = FullLLMCache ):
""" Initialize by creating all tables. """
""" Initialize by creating all tables. """
self . engine = engine
self . engine = engine
self . cache_schema = cache_schema
Base . metadata . create_all ( self . engine )
Base . metadata . create_all ( self . engine )
def lookup ( self , prompt : str , llm_string : str ) - > Optional [ RETURN_VAL_TYPE ] :
def lookup ( self , prompt : str , llm_string : str ) - > Optional [ RETURN_VAL_TYPE ] :
""" Look up based on prompt and llm_string. """
""" Look up based on prompt and llm_string. """
stmt = (
stmt = (
select ( FullLLMCache . response )
select ( self . cache_schema . response )
. where ( FullLLMCache . prompt == prompt )
. where ( self . cache_schema . prompt == prompt )
. where ( FullLLMCache . llm == llm_string )
. where ( self . cache_schema . llm == llm_string )
. order_by ( FullLLMCache . idx )
. order_by ( self . cache_schema . idx )
)
)
with Session ( self . engine ) as session :
with Session ( self . engine ) as session :
generations = [ ]
generations = [ ]
@ -80,7 +81,7 @@ class SQLAlchemyCache(BaseCache):
def update ( self , prompt : str , llm_string : str , return_val : RETURN_VAL_TYPE ) - > None :
def update ( self , prompt : str , llm_string : str , return_val : RETURN_VAL_TYPE ) - > None :
""" Look up based on prompt and llm_string. """
""" Look up based on prompt and llm_string. """
for i , generation in enumerate ( return_val ) :
for i , generation in enumerate ( return_val ) :
item = FullLLMCache (
item = self . cache_schema (
prompt = prompt , llm = llm_string , response = generation . text , idx = i
prompt = prompt , llm = llm_string , response = generation . text , idx = i
)
)
with Session ( self . engine ) as session , session . begin ( ) :
with Session ( self . engine ) as session , session . begin ( ) :