Improvements to Document Interface and correcting embedding workers

common/utils/middleware.py

@@ -4,7 +4,7 @@ for handling tenant requests
 """
 
 from flask_security import current_user
-from flask import session
+from flask import session, current_app
 
 from .database import Database
 
@@ -17,12 +17,15 @@ def mw_before_request():
 
     tenant_id = session['tenant']['id']
     if not tenant_id:
-        return {"message": "You are not logged into any tenant"}, 403
+        raise Exception('Cannot switch schema for tenant: no tenant defined in session')
+
+    for role in current_user.roles:
+        current_app.logger.debug(f'In middleware: User {current_user.email} has role {role.name}')
 
     # user = User.query.get(current_user.id)
-    if current_user.has_roles(['Super User']) or current_user.tenant_id == tenant_id:
+    if current_user.has_role('Super User') or current_user.tenant_id == tenant_id:
         Database(tenant_id).switch_schema()
     else:
-        return {"message": "You are not a member of this tenant"}, 403
+        raise Exception(f'Cannot switch schema for tenant {tenant_id}: user {current_user.email} does not have access')
 
 

common/utils/model_utils.py

@@ -0,0 +1,79 @@
+from flask import current_app
+from langchain.embeddings import OpenAIEmbeddings
+from langchain.chat_models import ChatOpenAI
+from langchain.prompts import ChatPromptTemplate
+
+from common.models.document import EmbeddingSmallOpenAI
+
+
+def select_model_variables(tenant):
+    embedding_provider = tenant.embedding_model.rsplit('.', 1)[0]
+    embedding_model = tenant.embedding_model.rsplit('.', 1)[1]
+
+    llm_provider = tenant.llm_model.rsplit('.', 1)[0]
+    llm_model = tenant.llm_model.rsplit('.', 1)[1]
+
+    # Set model variables
+    model_variables = {}
+    if tenant.es_k:
+        model_variables['k'] = tenant.es_k
+    else:
+        model_variables['k'] = 5
+
+    if tenant.es_similarity_threshold:
+        model_variables['similarity_threshold'] = tenant.es_similarity_threshold
+    else:
+        model_variables['similarity_threshold'] = 0.7
+
+    if tenant.chat_RAG_temperature:
+        model_variables['RAG_temperature'] = tenant.chat_RAG_temperature
+    else:
+        model_variables['RAG_temperature'] = 0.3
+
+    if tenant.chat_no_RAG_temperature:
+        model_variables['no_RAG_temperature'] = tenant.chat_no_RAG_temperature
+    else:
+        model_variables['no_RAG_temperature'] = 0.5
+
+    # Set Embedding variables
+    match embedding_provider:
+        case 'openai':
+            match embedding_model:
+                case 'text-embedding-3-small':
+                    api_key = current_app.config.get('OPENAI_API_KEY')
+                    model_variables['embedding_model'] = OpenAIEmbeddings(api_key=api_key,
+                                                                    model='text-embedding-3-small')
+                    model_variables['embedding_db_model'] = EmbeddingSmallOpenAI
+                    model_variables['min_chunk_size'] = current_app.config.get('OAI_TE3S_MIN_CHUNK_SIZE')
+                    model_variables['max_chunk_size'] = current_app.config.get('OAI_TE3S_MAX_CHUNK_SIZE')
+                case _:
+                    raise Exception(f'Error setting model variables for tenant {tenant.id} '
+                                    f'error: Invalid embedding model')
+        case _:
+            raise Exception(f'Error setting model variables for tenant {tenant.id} '
+                            f'error: Invalid embedding provider')
+
+    # Set Chat model variables
+    match llm_provider:
+        case 'openai':
+            api_key = current_app.config.get('OPENAI_API_KEY')
+            model_variables['llm'] = ChatOpenAI(api_key=api_key,
+                                                model=llm_model,
+                                                temperature=model_variables['RAG_temperature'])
+            match llm_model:
+                case 'gpt-4-turbo' | 'gpt-4o':
+                    summary_template = current_app.config.get('GPT4_SUMMARY_TEMPLATE')
+                    rag_template = current_app.config.get('GPT4_RAG_TEMPLATE')
+                case 'gpt-3-5-turbo':
+                    summary_template = current_app.config.get('GPT3_5_SUMMARY_TEMPLATE')
+                    rag_template = current_app.config.get('GPT3_5_RAG_TEMPLATE')
+                case _:
+                    raise Exception(f'Error setting model variables for tenant {tenant.id} '
+                                    f'error: Invalid chat model')
+            model_variables['summary_prompt'] = ChatPromptTemplate.from_template(summary_template)
+            model_variables['rag_prompt'] = ChatPromptTemplate.from_template(rag_template)
+        case _:
+            raise Exception(f'Error setting model variables for tenant {tenant.id} '
+                            f'error: Invalid chat provider')
+
+    return model_variables