275 lines
12 KiB
Python
275 lines
12 KiB
Python
import ast
|
|
import os
|
|
from datetime import datetime as dt, timezone as tz
|
|
|
|
import chardet
|
|
from flask import request, redirect, flash, render_template, Blueprint, session, current_app
|
|
from flask_security import roles_accepted, current_user
|
|
from sqlalchemy import desc
|
|
from sqlalchemy.orm import joinedload
|
|
from werkzeug.datastructures import FileStorage
|
|
from werkzeug.utils import secure_filename
|
|
from sqlalchemy.exc import SQLAlchemyError
|
|
import requests
|
|
from requests.exceptions import SSLError
|
|
from urllib.parse import urlparse
|
|
import io
|
|
|
|
from common.models.document import Embedding, DocumentVersion, Retriever
|
|
from common.models.interaction import ChatSession, Interaction, InteractionEmbedding, Specialist, SpecialistRetriever
|
|
from common.extensions import db
|
|
from common.utils.document_utils import set_logging_information, update_logging_information
|
|
from config.specialist_types import SPECIALIST_TYPES
|
|
from .document_forms import AddDocumentForm, AddURLForm, EditDocumentForm, EditDocumentVersionForm
|
|
from common.utils.middleware import mw_before_request
|
|
from common.utils.celery_utils import current_celery
|
|
from common.utils.nginx_utils import prefixed_url_for
|
|
from common.utils.view_assistants import form_validation_failed, prepare_table_for_macro
|
|
from .interaction_forms import SpecialistForm, EditSpecialistForm
|
|
|
|
interaction_bp = Blueprint('interaction_bp', __name__, url_prefix='/interaction')
|
|
|
|
|
|
@interaction_bp.before_request
|
|
def log_before_request():
|
|
current_app.logger.debug(f"Before request (interaction_bp): {request.method} {request.url}")
|
|
|
|
|
|
@interaction_bp.after_request
|
|
def log_after_request(response):
|
|
current_app.logger.debug(
|
|
f"After request (interaction_bp): {request.method} {request.url} - Status: {response.status}")
|
|
return response
|
|
|
|
|
|
@interaction_bp.before_request
|
|
def before_request():
|
|
try:
|
|
mw_before_request()
|
|
except Exception as e:
|
|
current_app.logger.error(f'Error switching schema in Interaction Blueprint: {e}')
|
|
for role in current_user.roles:
|
|
current_app.logger.debug(f'User {current_user.email} has role {role.name}')
|
|
raise
|
|
|
|
|
|
@interaction_bp.route('/chat_sessions', methods=['GET', 'POST'])
|
|
def chat_sessions():
|
|
page = request.args.get('page', 1, type=int)
|
|
per_page = request.args.get('per_page', 10, type=int)
|
|
|
|
query = ChatSession.query.order_by(desc(ChatSession.session_start))
|
|
|
|
pagination = query.paginate(page=page, per_page=per_page, error_out=False)
|
|
docs = pagination.items
|
|
|
|
rows = prepare_table_for_macro(docs, [('id', ''), ('session_id', ''), ('session_start', ''), ('session_end', '')])
|
|
|
|
return render_template('interaction/chat_sessions.html', rows=rows, pagination=pagination)
|
|
|
|
|
|
@interaction_bp.route('/handle_chat_session_selection', methods=['POST'])
|
|
@roles_accepted('Super User', 'Tenant Admin')
|
|
def handle_chat_session_selection():
|
|
chat_session_identification = request.form['selected_row']
|
|
cs_id = ast.literal_eval(chat_session_identification).get('value')
|
|
|
|
action = request.form['action']
|
|
|
|
match action:
|
|
case 'view_chat_session':
|
|
return redirect(prefixed_url_for('interaction_bp.view_chat_session', chat_session_id=cs_id))
|
|
|
|
# Add more conditions for other actions
|
|
return redirect(prefixed_url_for('interaction_bp.chat_sessions'))
|
|
|
|
|
|
@interaction_bp.route('/view_chat_session/<int:chat_session_id>', methods=['GET'])
|
|
@roles_accepted('Super User', 'Tenant Admin')
|
|
def view_chat_session(chat_session_id):
|
|
chat_session = ChatSession.query.get_or_404(chat_session_id)
|
|
interactions = (Interaction.query
|
|
.filter_by(chat_session_id=chat_session.id)
|
|
.order_by(Interaction.question_at)
|
|
.all())
|
|
|
|
# Fetch all related embeddings for the interactions in this session
|
|
embedding_query = (db.session.query(InteractionEmbedding.interaction_id,
|
|
DocumentVersion.url,
|
|
DocumentVersion.object_name)
|
|
.join(Embedding, InteractionEmbedding.embedding_id == Embedding.id)
|
|
.join(DocumentVersion, Embedding.doc_vers_id == DocumentVersion.id)
|
|
.filter(InteractionEmbedding.interaction_id.in_([i.id for i in interactions])))
|
|
|
|
# Create a dictionary to store embeddings for each interaction
|
|
embeddings_dict = {}
|
|
for interaction_id, url, object_name in embedding_query:
|
|
if interaction_id not in embeddings_dict:
|
|
embeddings_dict[interaction_id] = []
|
|
embeddings_dict[interaction_id].append({'url': url, 'object_name': object_name})
|
|
|
|
return render_template('interaction/view_chat_session.html',
|
|
chat_session=chat_session,
|
|
interactions=interactions,
|
|
embeddings_dict=embeddings_dict)
|
|
|
|
|
|
@interaction_bp.route('/view_chat_session_by_session_id/<session_id>', methods=['GET'])
|
|
@roles_accepted('Super User', 'Tenant Admin')
|
|
def view_chat_session_by_session_id(session_id):
|
|
chat_session = ChatSession.query.filter_by(session_id=session_id).first_or_404()
|
|
show_chat_session(chat_session)
|
|
|
|
|
|
def show_chat_session(chat_session):
|
|
interactions = Interaction.query.filter_by(chat_session_id=chat_session.id).all()
|
|
return render_template('interaction/view_chat_session.html', chat_session=chat_session, interactions=interactions)
|
|
|
|
|
|
@interaction_bp.route('/specialist', methods=['GET', 'POST'])
|
|
@roles_accepted('Super User', 'Tenant Admin')
|
|
def specialist():
|
|
form = SpecialistForm()
|
|
|
|
if form.validate_on_submit():
|
|
tenant_id = session.get('tenant').get('id')
|
|
try:
|
|
new_specialist = Specialist()
|
|
|
|
# Populate fields individually instead of using populate_obj (gives problem with QueryMultipleSelectField)
|
|
new_specialist.name = form.name.data
|
|
new_specialist.description = form.description.data
|
|
new_specialist.type = form.type.data
|
|
new_specialist.tuning = form.tuning.data
|
|
|
|
set_logging_information(new_specialist, dt.now(tz.utc))
|
|
|
|
db.session.add(new_specialist)
|
|
db.session.flush() # This assigns the ID to the specialist without committing the transaction
|
|
|
|
current_app.logger.debug(
|
|
f'New Specialist after flush - id: {new_specialist.id}, name: {new_specialist.name}')
|
|
|
|
# Create the retriever associations
|
|
selected_retrievers = form.retrievers.data
|
|
current_app.logger.debug(f'Selected Retrievers - {selected_retrievers}')
|
|
for retriever in selected_retrievers:
|
|
current_app.logger.debug(f'Creating association for Retriever - {retriever.id}')
|
|
specialist_retriever = SpecialistRetriever(
|
|
specialist_id=new_specialist.id,
|
|
retriever_id=retriever.id
|
|
)
|
|
db.session.add(specialist_retriever)
|
|
|
|
# Commit everything in one transaction
|
|
db.session.commit()
|
|
flash('Specialist successfully added!', 'success')
|
|
current_app.logger.info(f'Specialist {new_specialist.name} successfully added for tenant {tenant_id}!')
|
|
|
|
return redirect(prefixed_url_for('interaction_bp.edit_specialist', specialist_id=new_specialist.id))
|
|
|
|
except Exception as e:
|
|
db.session.rollback()
|
|
current_app.logger.error(f'Failed to add specialist. Error: {str(e)}', exc_info=True)
|
|
flash(f'Failed to add specialist. Error: {str(e)}', 'danger')
|
|
return render_template('interaction/specialist.html', form=form)
|
|
|
|
return render_template('interaction/specialists.html', form=form)
|
|
|
|
|
|
@interaction_bp.route('/specialist/<int:specialist_id>', methods=['GET', 'POST'])
|
|
@roles_accepted('Super User', 'Tenant Admin')
|
|
def edit_specialist(specialist_id):
|
|
specialist = Specialist.query.get_or_404(specialist_id)
|
|
form = EditSpecialistForm(request.form, obj=specialist)
|
|
|
|
configuration_config = SPECIALIST_TYPES[specialist.type]["configuration"]
|
|
form.add_dynamic_fields("configuration", configuration_config, specialist.configuration)
|
|
|
|
if request.method == 'GET':
|
|
# Pre-populate the retrievers field with current associations
|
|
current_app.logger.debug(f'Specialist retrievers: {specialist.retrievers}')
|
|
current_app.logger.debug(f'Form Retrievers Data Before: {form.retrievers.data}')
|
|
|
|
# Get the actual Retriever objects for the associated retriever_ids
|
|
retriever_objects = Retriever.query.filter(
|
|
Retriever.id.in_([sr.retriever_id for sr in specialist.retrievers])
|
|
).all()
|
|
form.retrievers.data = retriever_objects
|
|
|
|
current_app.logger.debug(f'Form Retrievers Data After: {form.retrievers.data}')
|
|
|
|
if form.validate_on_submit():
|
|
# Update the basic fields
|
|
form.populate_obj(specialist)
|
|
# Update the configuration dynamic fields
|
|
specialist.configuration = form.get_dynamic_data("configuration")
|
|
|
|
# Update retriever associations
|
|
current_retrievers = set(sr.retriever_id for sr in specialist.retrievers)
|
|
selected_retrievers = set(r.id for r in form.retrievers.data)
|
|
|
|
# Remove unselected retrievers
|
|
for sr in specialist.retrievers[:]:
|
|
if sr.retriever_id not in selected_retrievers:
|
|
db.session.delete(sr)
|
|
|
|
# Add new retrievers
|
|
for retriever_id in selected_retrievers - current_retrievers:
|
|
specialist_retriever = SpecialistRetriever(
|
|
specialist_id=specialist.id,
|
|
retriever_id=retriever_id
|
|
)
|
|
db.session.add(specialist_retriever)
|
|
|
|
# Update logging information
|
|
update_logging_information(specialist, dt.now(tz.utc))
|
|
|
|
try:
|
|
db.session.commit()
|
|
flash('Specialist updated successfully!', 'success')
|
|
current_app.logger.info(f'Specialist {specialist.id} updated successfully')
|
|
except SQLAlchemyError as e:
|
|
db.session.rollback()
|
|
flash(f'Failed to update specialist. Error: {str(e)}', 'danger')
|
|
current_app.logger.error(f'Failed to update specialist {specialist_id}. Error: {str(e)}')
|
|
return render_template('interaction/edit_specialist.html', form=form, specialist_id=specialist_id)
|
|
|
|
return redirect(prefixed_url_for('interaction_bp.specialists'))
|
|
else:
|
|
form_validation_failed(request, form)
|
|
|
|
return render_template('interaction/edit_specialist.html', form=form, specialist_id=specialist_id)
|
|
|
|
|
|
@interaction_bp.route('/specialists', methods=['GET', 'POST'])
|
|
@roles_accepted('Super User', 'Tenant Admin')
|
|
def specialists():
|
|
page = request.args.get('page', 1, type=int)
|
|
per_page = request.args.get('per_page', 10, type=int)
|
|
|
|
query = Specialist.query.order_by(Specialist.id)
|
|
|
|
pagination = query.paginate(page=page, per_page=per_page)
|
|
the_specialists = pagination.items
|
|
|
|
# prepare table data
|
|
rows = prepare_table_for_macro(the_specialists,
|
|
[('id', ''), ('name', ''), ('type', '')])
|
|
|
|
# Render the catalogs in a template
|
|
return render_template('interaction/specialists.html', rows=rows, pagination=pagination)
|
|
|
|
|
|
@interaction_bp.route('/handle_specialist_selection', methods=['POST'])
|
|
@roles_accepted('Super User', 'Tenant Admin')
|
|
def handle_specialist_selection():
|
|
specialist_identification = request.form.get('selected_row')
|
|
specialist_id = ast.literal_eval(specialist_identification).get('value')
|
|
action = request.form.get('action')
|
|
|
|
if action == "edit_specialist":
|
|
return redirect(prefixed_url_for('interaction_bp.edit_specialist', specialist_id=specialist_id))
|
|
|
|
return redirect(prefixed_url_for('interaction_bp.specialists'))
|