eveAI/eveai_app/views/interaction_views.py

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, cache_manager
from common.utils.model_logging_utils import set_logging_information, update_logging_information

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 common.utils.specialist_utils import initialize_specialist

from config.type_defs.specialist_types import SPECIALIST_TYPES

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/<int:chat_session_id>', 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/<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.type_version = cache_manager.specialist_config_cache.get_latest_version(new_specialist.type)
            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}!')

            # Initialize the newly create specialist
            initialize_specialist(new_specialist.id, new_specialist.type, new_specialist.type_version)

            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/<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)

    specialist_config = cache_manager.specialist_config_cache.get_config(specialist.type, specialist.type_version)
    configuration_config = specialist_config.get('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'))