added caching; separated everything into classes; enriched api response

Author: deadbits

server.py

@@ -1,8 +1,11 @@
 #!/usr/bin/env python3
 # github.com/deadbits/vector-embedding-api
+# server.py
 import os
 import sys
+import time
 import argparse
+import hashlib
 import logging
 import configparser
 
@@ -29,7 +32,7 @@ def __init__(self, config_file):
         logging.info(f'Loading config file: {self.config_file}')
         self.config.read(config_file)
 
-    def get(self, section, key):
+    def get_val(self, section, key):
         answer = None
 
         try:
@@ -38,44 +41,115 @@ def get(self, section, key):
             logging.error(f'Config file missing section: {section} - {err}')
 
         return answer
+    
+    def get_bool(self, section, key, default=False):
+        try:
+            return self.config.getboolean(section, key)
+        except Exception as err:
+            logging.error(f'Failed to parse boolean - returning default "False": {section} - {err}')
+            return default
+
+
+class EmbeddingCache:
+    def __init__(self):
+        logger.info('Created in-memory cache')
+        self.cache = {}
+
+    def get_cache_key(self, text, model_type):
+        return hashlib.sha256((text + model_type).encode()).hexdigest()
+
+    def get(self, text, model_type):
+        return self.cache.get(self.get_cache_key(text, model_type))
 
+    def set(self, text, model_type, embedding):
+        self.cache[self.get_cache_key(text, model_type)] = embedding
 
-def get_openai_embeddings(text: str) -> list:
-    try:
-        response = openai.Embedding.create(input=text, model='text-embedding-ada-002')
-        return response['data'][0]['embedding']
-    except Exception as err:
-        logger.error(f'Failed to get OpenAI embeddings: {err}')
-        abort(500, 'Failed to get OpenAI embeddings')
 
+class EmbeddingGenerator:
+    def __init__(self, sbert_model=None, openai_key=None):
+        self.sbert_model = sbert_model
+        if self.sbert_model is not None:
+            try:
+                self.model = SentenceTransformer(self.sbert_model)
+                logger.info(f'enabled model: {self.sbert_model}')
+            except Exception as err:
+                logger.error(f'Failed to load SentenceTransformer model "{self.sbert_model}": {err}')
+                sys.exit(1)
 
-def get_transformers_embeddings(text: str) -> list:
-    try:
-        return model.encode(text).tolist()
-    except Exception as err:
-        logger.error(f'Failed to get sentence-transformers embeddings: {err}')
-        abort(500, 'Failed to get sentence-transformers embeddings')
+        if openai_key is not None:
+            openai.api_key = openai_key
+            logger.info('enabled model: text-embedding-ada-002')
+
+    def get_openai_embeddings(self, text):
+        start_time = time.time()
+
+        try:
+            response = openai.Embedding.create(input=text, model='text-embedding-ada-002')
+            elapsed_time = (time.time() - start_time) * 1000
+            data = {
+                "embedding": response['data'][0]['embedding'],
+                "status": "success",
+                "elapsed": elapsed_time,
+                "model": "text-embedding-ada-002"
+            }
+            return data
+        except Exception as err:
+            logger.error(f'Failed to get OpenAI embeddings: {err}')
+            return {"status": "error", "message": str(err), "model": "text-embedding-ada-002"}
+
+    def get_transformers_embeddings(self, text):
+        start_time = time.time()
+
+        try:
+            embedding = self.model.encode(text).tolist()
+            elapsed_time = (time.time() - start_time) * 1000
+            data = {
+                "embedding": embedding,
+                "status": "success",
+                "elapsed": elapsed_time,
+                "model": self.sbert_model
+            }
+            return data
+        except Exception as err:
+            logger.error(f'Failed to get sentence-transformers embeddings: {err}')
+            return {"status": "error", "message": str(err), "model": self.sbert_model}
+
+    def generate(self, text, model_type):
+        if model_type == 'openai':
+            return self.get_openai_embeddings(text)
+        else:
+            return self.get_transformers_embeddings(text)
 
 
 @app.route('/submit', methods=['POST'])
 def submit_text():
     data = request.json
-    
+
     text_data = data.get('text')
     model_type = data.get('model', 'local').lower()
 
     if text_data is None:
         abort(400, 'Missing text data to embed')
-    
+
     if model_type not in ['local', 'openai']:
         abort(400, 'model field must be one of: local, openai')
 
-    if model_type == 'openai':
-        embedding_data = get_openai_embeddings(text_data)
+    if embedding_cache:
+        result = embedding_cache.get(text_data, model_type)
+        if result:
+            logger.info('found embedding in cache!')
+            result = {'embedding': result, 'cache': True, "status": 'success'}
     else:
-        embedding_data = get_transformers_embeddings(text_data)
+        result = None
+
+    if result is None:
+        result = embedding_generator.generate(text_data, model_type)
+
+        if embedding_cache and result['status'] == 'success':
+            embedding_cache.set(text_data, model_type, result['embedding'])
+            logger.info('added to cache')
 
-    return jsonify({'embedding': embedding_data, 'status': 'success'})
+    return jsonify(result)
 
 
 if __name__ == '__main__':
@@ -91,23 +165,21 @@ def submit_text():
     args = parser.parse_args()
 
     conf = Config(args.config)
-    api_key = conf.get('main', 'openai_api_key')
-    sent_model = conf.get('main', 'sent_transformers_model')
+    openai_key = conf.get_val('main', 'openai_api_key')
+    sbert_model = conf.get_val('main', 'sent_transformers_model')
+    use_cache = conf.get_bool('main', 'use_cache', default=False)
 
-    if api_key is None:
+    if openai_key is None:
         logger.warn('No OpenAI API key set in configuration file: server.conf')
-    else:
-        logger.info('Set OpenAI API key via openai.api_key')
-        openai.api_key = api_key
 
-    if sent_model is None:
+    if sbert_model is None:
         logger.warn('No transformer model set in configuration file: server.conf')
 
-    try:
-        model = SentenceTransformer(sent_model)
-    except Exception as err:
-        logger.error(f'Failed to load SentenceTransformer model "{sent_model}": {err}')
+    if openai_key is None and sbert_model is None:
+        logger.error('No sbert model set *and* no openAI key set; exiting')
         sys.exit(1)
 
-    app.run(debug=True)
+    embedding_cache = EmbeddingCache() if use_cache else None
+    embedding_generator = EmbeddingGenerator(sbert_model, openai_key)
 
+    app.run(debug=True)