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