From 85c1bd2cd04dcdb2367e96fb49b05e40bc755f41 Mon Sep 17 00:00:00 2001
From: Benjamin <BenderV@users.noreply.github.com>
Date: Sat, 17 Dec 2022 01:47:23 +0100
Subject: [PATCH] add sqlalchemy generic cache (#361)

Created a generic SQLAlchemyCache class to plug any database supported
by SQAlchemy. (I am using Postgres).
I also based the class SQLiteCache class on this class SQLAlchemyCache.

As a side note, I'm questioning the need for two distinct class
LLMCache, FullLLMCache. Shouldn't we merge both ?
---
 docs/examples/prompts/llm_functionality.ipynb | 18 ++++++++++++++---
 langchain/cache.py                            | 20 ++++++++++++++-----
 2 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/docs/examples/prompts/llm_functionality.ipynb b/docs/examples/prompts/llm_functionality.ipynb
index 251e22a6b8e..7edfc84104e 100644
--- a/docs/examples/prompts/llm_functionality.ipynb
+++ b/docs/examples/prompts/llm_functionality.ipynb
@@ -373,12 +373,19 @@
    "id": "6053408b",
    "metadata": {},
    "outputs": [],
-   "source": []
+   "source": [
+    "# You can use SQLAlchemyCache to cache with any SQL database supported by SQLAlchemy.\n",
+    "from langchain.cache import SQLAlchemyCache\n",
+    "from sqlalchemy import create_engine\n",
+    "\n",
+    "engine = create_engine(\"postgresql://postgres:postgres@localhost:5432/postgres\")\n",
+    "langchain.llm_cache = SQLAlchemyCache(engine)"
+   ]
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "base",
    "language": "python",
    "name": "python3"
   },
@@ -392,7 +399,12 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.8"
+   "version": "3.9.12 (main, Jun  1 2022, 06:34:44) \n[Clang 12.0.0 ]"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "1235b9b19e8e9828b5c1fdb2cd89fe8d3de0fcde5ef5f3db36e4b671adb8660f"
+   }
   }
  },
  "nbformat": 4,
diff --git a/langchain/cache.py b/langchain/cache.py
index 86a9dc6b210..258c0383f55 100644
--- a/langchain/cache.py
+++ b/langchain/cache.py
@@ -3,6 +3,7 @@ from abc import ABC, abstractmethod
 from typing import Dict, List, Optional, Tuple, Union
 
 from sqlalchemy import Column, Integer, String, create_engine, select
+from sqlalchemy.engine.base import Engine
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import Session
 
@@ -61,12 +62,12 @@ class FullLLMCache(Base):  # type: ignore
     response = Column(String)
 
 
-class SQLiteCache(BaseCache):
-    """Cache that uses SQLite as a backend."""
+class SQLAlchemyCache(BaseCache):
+    """Cache that uses SQAlchemy as a backend."""
 
-    def __init__(self, database_path: str = ".langchain.db"):
-        """Initialize by creating the engine and all tables."""
-        self.engine = create_engine(f"sqlite:///{database_path}")
+    def __init__(self, engine: Engine):
+        """Initialize by creating all tables."""
+        self.engine = engine
         Base.metadata.create_all(self.engine)
 
     def lookup(self, prompt: str, llm_string: str) -> Optional[RETURN_VAL_TYPE]:
@@ -106,3 +107,12 @@ class SQLiteCache(BaseCache):
                 )
                 with Session(self.engine) as session, session.begin():
                     session.add(item)
+
+
+class SQLiteCache(SQLAlchemyCache):
+    """Cache that uses SQLite as a backend."""
+
+    def __init__(self, database_path: str = ".langchain.db"):
+        """Initialize by creating the engine and all tables."""
+        engine = create_engine(f"sqlite:///{database_path}")
+        super().__init__(engine)