From a2f109c52ef1206114910cdc47a0c02341bb0b6f Mon Sep 17 00:00:00 2001
From: Shivam Pachchigar <45695285+ShivamHP@users.noreply.github.com>
Date: Thu, 4 Jan 2024 00:00:55 +0530
Subject: [PATCH] Added caching for repititive requests on the `/v1/complete`
 endpoint - Used Flask extension Flask-Caching - Currently uses a local Python
 dictionary for caching, but can be updated to use Redis - Created a
 `cache.py` file which stores the config for caching and it also prevents
 circular imports if written directly in `main.py`. Also, it makes the code
 better organized

---
 src/llm_vm/server/cache.py  | 16 ++++++++++++++++
 src/llm_vm/server/main.py   |  3 ++-
 src/llm_vm/server/routes.py |  2 ++
 3 files changed, 20 insertions(+), 1 deletion(-)
 create mode 100644 src/llm_vm/server/cache.py

diff --git a/src/llm_vm/server/cache.py b/src/llm_vm/server/cache.py
new file mode 100644
index 00000000..8f85e2ac
--- /dev/null
+++ b/src/llm_vm/server/cache.py
@@ -0,0 +1,16 @@
+from flask_caching import Cache
+from flask import request
+
+def make_cache_key(*args, **kwargs):
+    key = request.url + str(request.data)
+    return key
+
+config = {
+    "DEBUG": True,
+    "CACHE_TYPE": "SimpleCache",
+    "CACHE_DEFAULT_TIMEOUT": 30,
+    'CACHE_KEY_PREFIX': 'custom_prefix',
+    'CACHE_KEY_FUNC': make_cache_key
+}
+
+cache = Cache(config=config)
\ No newline at end of file
diff --git a/src/llm_vm/server/main.py b/src/llm_vm/server/main.py
index 6c4e5ba8..9b8ce338 100644
--- a/src/llm_vm/server/main.py
+++ b/src/llm_vm/server/main.py
@@ -9,6 +9,7 @@
 import hashlib
 import sys
 from flask_cors import CORS
+from cache import cache
 from contextlib import contextmanager
 import llm_vm.server.routes as routes
 from llm_vm.config import settings
@@ -16,7 +17,7 @@
 
 app = flask.Flask(__name__)
 CORS(app)
-app.config["DEBUG"] = True
+cache.init_app(app)
 
 # Register_blueprint from routes to load API
 app.register_blueprint(routes.bp)
diff --git a/src/llm_vm/server/routes.py b/src/llm_vm/server/routes.py
index ef36bd14..7688d095 100644
--- a/src/llm_vm/server/routes.py
+++ b/src/llm_vm/server/routes.py
@@ -6,6 +6,7 @@
 from llm_vm.agents.REBEL import agent
 from llm_vm.client import Client
 from llm_vm.config import settings
+from cache import cache
 # load optimizer for endpoint use
 # optimizer = LocalOptimizer(MIN_TRAIN_EXS=2,openai_key=None)
 
@@ -20,6 +21,7 @@ def home():
     return '''home'''
 
 @bp.route('/v1/complete', methods=['POST'])
+@cache.cached()
 def optimizing_complete():
     rebel_agent = agent.Agent("", [], verbose=1)
     data = json.loads(request.data)