import ast from datetime import datetime as dt, timezone as tz from flask import request, redirect, flash, render_template, Blueprint, session, current_app from flask_security import roles_accepted from sqlalchemy import desc from sqlalchemy.exc import SQLAlchemyError 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.type_defs.specialist_types import SPECIALIST_TYPES from common.utils.middleware import mw_before_request 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: {request.path} =====================================') @interaction_bp.after_request def log_after_request(response): 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}') 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/', methods=['GET']) @roles_accepted('Super User', 'Tenant Admin') def view_chat_session(chat_session_id): # Get chat session with user info chat_session = ChatSession.query.get_or_404(chat_session_id) # Get interactions with specialist info interactions = (Interaction.query .filter_by(chat_session_id=chat_session.id) .join(Specialist, Interaction.specialist_id == Specialist.id, isouter=True) .add_columns( Interaction.id, Interaction.question_at, Interaction.specialist_arguments, Interaction.specialist_results, Specialist.name.label('specialist_name'), Specialist.type.label('specialist_type') ).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/', 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 # Create the retriever associations selected_retrievers = form.retrievers.data for retriever in selected_retrievers: 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/specialist.html', form=form) @interaction_bp.route('/specialist/', 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': # 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 if form.validate_on_submit(): # Update the basic fields specialist.name = form.name.data specialist.description = form.description.data specialist.tuning = form.tuning.data # Update the configuration dynamic fields specialist.configuration = form.get_dynamic_data("configuration") # Get current and selected retrievers current_retrievers = {sr.retriever_id: sr for sr in specialist.retrievers} selected_retrievers = {r.id: r for r in form.retrievers.data} # Remove unselected retrievers for retriever_id in set(current_retrievers.keys()) - set(selected_retrievers.keys()): specialist_retriever = current_retrievers[retriever_id] db.session.delete(specialist_retriever) # Add new retrievers for retriever_id in set(selected_retrievers.keys()) - set(current_retrievers.keys()): 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') return redirect(prefixed_url_for('interaction_bp.specialists')) 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) 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'))