Added caching for repititive requests on the /v1/complete endpoint

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)

src/llm_vm/server/main.py

@@ -9,14 +9,15 @@
 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
 from llm_vm.utils.ram import RAMLogger
 
 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)

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)